Features: Adding keys, removing keys and printing secrets (#69)

README.md

@@ -43,6 +43,11 @@ Usage
         $ gauth Google -b
         477615
 
+- Run `gauth KEYNAME -s` to retrieve an accounts secret from the config.
+
+        $ gauth Google -s
+        your_secret_for_google
+
 - `gauth` is convenient to use in `watch`.
 
         $ watch -n1 gauth
@@ -53,6 +58,22 @@ Usage
   from an existing Google Authenticator setup, on a phone to which you do not
   have root access), then [gauthQR](https://github.com/jbert/gauthQR) may be useful.
 
+
+Adding and removing keys
+------------------------
+
+- Run `gauth KEYNAME -a` to add a new key.
+
+        $ gauth Google -a
+        Key for Google: examplekey
+        Current OTP for Google: 306726
+
+- Run `gauth KEYNAME -r` to remove an existing key.
+
+        $ gauth Google -r
+        Are you sure you want to remove Google [y/N]: y
+        Google has been removed.
+
 Encryption
 ----------
 

gauth.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"fmt"
 	"log"
 	"os"
@@ -17,33 +18,54 @@ import (
 
 func main() {
 	accountName := ""
-	isBareCode := false
+	argument := ""
 
 	if len(os.Args) > 1 {
 		accountName = os.Args[1]
 	}
+
 	if len(os.Args) > 2 {
 		if os.Args[2] == "-b" || os.Args[2] == "-bare" {
-			isBareCode = true
+			argument = "bare"
+		} else if os.Args[2] == "-a" || os.Args[2] == "-add" {
+			argument = "add"
+		} else if os.Args[2] == "-r" || os.Args[2] == "-remove" {
+			argument = "remove"
+		} else if os.Args[2] == "-s" || os.Args[2] == "-secret" {
+			argument = "secret"
 		}
 	}
 
-	urls := getUrls()
+	if accountName != "" {
-
+		switch argument {
-	if isBareCode && accountName != "" {
+		case "bare":
-		printBareCode(accountName, urls)
+			printBareCode(accountName, getUrls())
-	} else {
+			return
-		printAllCodes(urls)
+		case "add":
+			addCode(accountName)
+			return
+		case "remove":
+			removeCode(accountName)
+			return
+		case "secret":
+			printSecret(accountName, getUrls())
+			return
+		default:
+			printAllCodes(getUrls())
+			return
 		}
 	}
 
+	printAllCodes(getUrls())
+}
+
 func getPassword() ([]byte, error) {
 	fmt.Printf("Encryption password: ")
 	defer fmt.Println()
 	return term.ReadPassword(int(syscall.Stdin))
 }
 
-func getUrls() []*otpauth.URL {
+func getConfigPath() string {
 	cfgPath := os.Getenv("GAUTH_CONFIG")
 	if cfgPath == "" {
 		user, err := user.Current()
@@ -53,6 +75,12 @@ func getUrls() []*otpauth.URL {
 		cfgPath = filepath.Join(user.HomeDir, ".config", "gauth.csv")
 	}
 
+	return cfgPath
+}
+
+func getUrls() []*otpauth.URL {
+	cfgPath := getConfigPath()
+
 	cfgContent, err := gauth.LoadConfigFile(cfgPath, getPassword)
 	if err != nil {
 		log.Fatalf("Loading config: %v", err)
@@ -79,6 +107,137 @@ func printBareCode(accountName string, urls []*otpauth.URL) {
 	}
 }
 
+func addCode(accountName string) {
+	cfgPath := getConfigPath()
+
+	// Check for encryption and ask for password if necessary
+	_, isEncrypted, err := gauth.ReadConfigFile(cfgPath)
+
+	password, err := []byte(nil), nil
+
+	if isEncrypted {
+		password, err = getPassword()
+
+		if err != nil {
+			log.Fatalf("reading passphrase: %v", err)
+		}
+	}
+
+	// Get decoded config
+	rawConfig, err := gauth.LoadConfigFile(cfgPath, func() ([]byte, error) { return password, err })
+	if err != nil {
+		log.Fatalf("Loading config: %v", err)
+	}
+
+	newConfig := strings.TrimSuffix(string(rawConfig), "\n")
+
+	// Check if account already exists
+	for _, line := range strings.Split(newConfig, "\n") {
+		if strings.HasPrefix(strings.ToLower(line), strings.ToLower(accountName)) {
+			fmt.Printf("Account \"%s\" already exists. Nothing has been added.", accountName)
+			return
+		}
+	}
+
+	// Read new key
+	fmt.Printf("Key for %s: ", accountName)
+	reader := bufio.NewReader(os.Stdin)
+	key, _ := reader.ReadString('\n')
+
+	// Append new key
+	newConfig += "\n" + accountName + ":" + key + "\n"
+
+	// Try parsing the new config and print the current OTP
+	parsedConfig, err := gauth.ParseConfig([]byte(newConfig))
+	if err != nil {
+		log.Fatalf("Parsing new config: %v", err)
+	}
+
+	fmt.Printf("Current OTP for %s: ", accountName)
+	printBareCode(accountName, parsedConfig)
+
+	// write new config
+	err = gauth.WriteConfigFile(cfgPath, password, []byte(newConfig))
+	if err != nil {
+		log.Fatalf("Error writing new config: %v", err)
+	}
+}
+
+func removeCode(accountName string) {
+	cfgPath := getConfigPath()
+
+	// Check for encryption and ask for password if necessary
+	_, isEncrypted, err := gauth.ReadConfigFile(cfgPath)
+
+	password, err := []byte(nil), nil
+
+	if isEncrypted {
+		password, err = getPassword()
+
+		if err != nil {
+			log.Fatalf("reading passphrase: %v", err)
+		}
+	}
+
+	// Get decoded config
+	rawConfig, err := gauth.LoadConfigFile(cfgPath, func() ([]byte, error) { return password, err })
+	if err != nil {
+		log.Fatalf("Loading config: %v", err)
+	}
+
+	newConfig := ""
+	anythingRemoved := false
+
+	// Iterate over config lines and search for the one to be removed
+	for _, line := range strings.Split(string(rawConfig), "\n") {
+		trim := strings.TrimSpace(line)
+		if trim == "" {
+			continue
+		}
+
+		if strings.HasPrefix(strings.ToLower(trim), strings.ToLower(accountName)) {
+			anythingRemoved = true
+			continue
+		}
+
+		newConfig += trim + "\n"
+
+	}
+
+	if !anythingRemoved {
+		fmt.Printf("Account \"%s\" was not found. Nothing has been removed.", accountName)
+		return
+	}
+
+	// Prompt for confirmation
+	fmt.Printf("Are you sure you want to remove %s [y/N]: ", accountName)
+	reader := bufio.NewReader(os.Stdin)
+	confirmation, _ := reader.ReadString('\n')
+
+	confirmation = strings.TrimSpace(confirmation)
+
+	if strings.ToLower(confirmation) != "y" {
+		return
+	}
+
+	// Write the new config
+	err = gauth.WriteConfigFile(cfgPath, password, []byte(newConfig))
+	if err != nil {
+		log.Fatalf("Error writing new config: %v", err)
+	}
+
+	fmt.Printf("%s has been removed.", accountName)
+}
+
+func printSecret(accountName string, urls []*otpauth.URL) {
+	for _, url := range urls {
+		if strings.EqualFold(strings.ToLower(accountName), strings.ToLower(url.Account)) {
+			fmt.Print(url.RawSecret)
+			break
+		}
+	}
+}
+
 func printAllCodes(urls []*otpauth.URL) {
 	_, progress := gauth.IndexNow() // TODO: do this per-code
 

gauth/gauth.go

@@ -75,15 +75,31 @@ func CodesAtTimeStep(u *otpauth.URL, timeStep uint64) (prev, curr, next string,
 	return
 }
 
+// ReadConfigFile reads the config file at path and returns its contents and
+// whether it is encrypted or not
+func ReadConfigFile(path string) ([]byte, bool, error) {
+	data, err := ioutil.ReadFile(path)
+	if err != nil {
+		return nil, false, err
+	}
+
+	if bytes.HasPrefix(data, []byte("Salted__")) {
+		return data, true, nil // encrypted
+	}
+
+	return data, false, nil
+}
+
 // LoadConfigFile reads and decrypts, if necessary, the CSV config at path.
 // The getPass function is called to obtain a password if needed.
 func LoadConfigFile(path string, getPass func() ([]byte, error)) ([]byte, error) {
-	data, err := ioutil.ReadFile(path)
+	data, isEncrypted, err := ReadConfigFile(path)
+
 	if err != nil {
 		return nil, err
 	}
 
-	if !bytes.HasPrefix(data, []byte("Salted__")) {
+	if !isEncrypted {
 		return data, nil // not encrypted
 	}
 
@@ -122,6 +138,61 @@ func LoadConfigFile(path string, getPass func() ([]byte, error)) ([]byte, error)
 	return rest[:len(rest)-pad], nil
 }
 
+// WriteConfigFile encrypts the provided newConfig using passwd, if necessary,
+// and writes it to path
+func WriteConfigFile(path string, passwd []byte, newConfig []byte) error {
+	data, isEncrypted, err := ReadConfigFile(path)
+
+	if err != nil {
+		return err
+	}
+
+	if isEncrypted {
+		// Encrypt newConfig using the same salt as in the old config
+		salt := data[8:16]
+		salting := sha256.New()
+		salting.Write(passwd)
+		salting.Write(salt)
+		sum := salting.Sum(nil)
+		key := sum[:16]
+		iv := sum[16:]
+
+		block, err := aes.NewCipher(key)
+
+		if err != nil {
+			return fmt.Errorf("creating cipher: %v", err)
+		}
+
+		mode := cipher.NewCBCEncrypter(block, iv)
+
+		// Add needed CBC block padding
+		padLength := 16 - (len(newConfig) % 16)
+		pad := make([]byte, padLength)
+
+		for i := range pad {
+			pad[i] = byte(padLength)
+		}
+
+		newConfig = append(newConfig, pad...)
+
+		// Encrypt and construct the new data to be written
+		mode.CryptBlocks(newConfig, newConfig)
+
+		saltedPrefix := []byte("Salted__")
+		saltedPrefix = append(saltedPrefix, salt...)
+
+		newConfig = append(saltedPrefix, newConfig...)
+	}
+
+	err = ioutil.WriteFile(path, newConfig, 0)
+
+	if err != nil {
+		return fmt.Errorf("writing config: %v", err)
+	}
+
+	return err
+}
+
 // ParseConfig parses the contents of data as a gauth configuration file.  Each
 // line of the file specifies a single configuration.
 //