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

gauth.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"fmt"
 	"log"
 	"os"
@@ -17,24 +18,45 @@ 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 isBareCode && accountName != "" {
-		printBareCode(accountName, urls)
-	} else {
-		printAllCodes(urls)
+	if accountName != "" {
+		switch argument {
+		case "bare":
+			printBareCode(accountName, getUrls())
+			return
+		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) {
@@ -43,7 +65,7 @@ func getPassword() ([]byte, error) {
 	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