Allow trailing wildcard to display multiple TOTP

main.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"os"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/Luzifer/rconfig"
@@ -95,7 +96,7 @@ func main() {
 
 		for _, s := range secrets {
 			output.EraseCurrentLine()
-			o, err := s.BuildOutput(len(secrets) > 1)
+			o, err := s.BuildOutput(len(secrets) > 1, tokenList(secrets).LongestName())
 			if err != nil {
 				fmt.Printf("%s: ERROR (%s)\n", s.Name, err)
 				continue
@@ -123,14 +124,54 @@ func getSecretsFromVault() ([]token, error) {
 
 	client.SetToken(cfg.VaultToken)
 
-	data, err := client.Logical().Read(rconfig.Args()[1])
+	resp := []token{}
+	keyPool := []string{}
+
+	key := rconfig.Args()[1]
+
+	if strings.HasSuffix(key, "*") {
+		scanPool := []string{strings.TrimRight(key, "*")}
+
+		for len(scanPool) > 0 {
+			s, err := client.Logical().List(scanPool[0])
 			if err != nil {
-		return nil, fmt.Errorf("Unable to read from key %q: %s", rconfig.Args()[1], err)
+				return nil, fmt.Errorf("Unable to list keys %q: %s", key, err)
+			}
+
+			if s == nil {
+				return nil, fmt.Errorf("There is no key %q", scanPool[0])
+			}
+
+			if s.Data["keys"] != nil {
+				for _, sk := range s.Data["keys"].([]interface{}) {
+					sks := sk.(string)
+					if strings.HasSuffix(sks, "/") {
+						scanPool = append(scanPool, scanPool[0]+sks)
+					} else {
+						keyPool = append(keyPool, scanPool[0]+sks)
+					}
+				}
+			}
+
+			scanPool = scanPool[1:len(scanPool)]
+		}
+
+	} else {
+		keyPool = append(keyPool, key)
+	}
+
+	for _, k := range keyPool {
+		data, err := client.Logical().Read(k)
+		if err != nil {
+			return nil, fmt.Errorf("Unable to read from key %q: %s", k, err)
 		}
 
 		if data.Data[cfg.Field] == nil {
-		return nil, fmt.Errorf("The key %q does not have a field named %q.", rconfig.Args()[1], cfg.Field)
+			return nil, fmt.Errorf("The key %q does not have a field named %q.", k, cfg.Field)
 		}
 
-	return []token{{Name: rconfig.Args()[1], Secret: data.Data[cfg.Field].(string)}}, nil
+		resp = append(resp, token{Name: k, Secret: data.Data[cfg.Field].(string)})
+	}
+
+	return resp, nil
 }

token.go

@@ -14,10 +14,16 @@ type token struct {
 }
 
 func (t token) GetCode(in time.Time) (string, error) {
-	return totp.GenerateCode(strings.ToUpper(t.Secret), in)
+	secret := t.Secret
+
+	if n := len(secret) % 8; n != 0 {
+		secret = secret + strings.Repeat("=", 8-n)
+	}
+
+	return totp.GenerateCode(strings.ToUpper(secret), in)
 }
 
-func (t token) BuildOutput(showName bool) (string, error) {
+func (t token) BuildOutput(showName bool, namePad int) (string, error) {
 	// Output: "123456 (Valid 12s)", "123456 (Valid 1s)"
 
 	n := time.Now()
@@ -28,7 +34,7 @@ func (t token) BuildOutput(showName bool) (string, error) {
 
 	var output string
 	if showName {
-		output = fmt.Sprintf("%s: ", t.Name)
+		output = fmt.Sprintf(fmt.Sprintf("%%-%ds", namePad+2), fmt.Sprintf("%s:", t.Name))
 	}
 
 	output = fmt.Sprintf("%s%s", output, code)
@@ -48,3 +54,13 @@ type tokenList []token
 func (t tokenList) Len() int           { return len(t) }
 func (t tokenList) Less(i, j int) bool { return t[i].Name < t[j].Name }
 func (t tokenList) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
+
+func (t tokenList) LongestName() (l int) {
+	for _, s := range t {
+		if ll := len(s.Name); ll > l {
+			l = ll
+		}
+	}
+
+	return
+}