1
0
Fork 0
mirror of https://github.com/Luzifer/vault-totp.git synced 2024-10-18 08:04:20 +00:00

Prepare multi-token display while retaining same functionality

This commit is contained in:
Knut Ahlers 2017-01-04 21:20:20 +01:00
parent 7c9d67490a
commit d445ab048d
Signed by: luzifer
GPG key ID: DC2729FDD34BE99E
2 changed files with 84 additions and 36 deletions

68
main.go
View file

@ -5,13 +5,13 @@ import (
"io/ioutil"
"log"
"os"
"strings"
"sort"
"time"
"github.com/Luzifer/rconfig"
"github.com/hashicorp/vault/api"
homedir "github.com/mitchellh/go-homedir"
"github.com/pquerna/otp/totp"
"github.com/sethgrid/curse"
)
var (
@ -67,72 +67,70 @@ func main() {
}
var (
secret string
secrets []token
err error
)
if cfg.TestSecret == "" {
secret, err = getSecretFromVault()
secrets, err = getSecretsFromVault()
if err != nil {
log.Fatalf("Error to retrieve secret: %s", err)
}
} else {
secret = cfg.TestSecret
secrets = []token{{Name: "Test", Secret: cfg.TestSecret}}
}
for range time.Tick(250 * time.Millisecond) {
output, err := buildOutput(secret)
sort.Sort(tokenList(secrets))
output, err := curse.New()
if err != nil {
log.Fatalf("An error ocurred while generating the code: %s", err)
log.Fatal(err)
}
firstRun := true
for range time.Tick(250 * time.Millisecond) {
if !firstRun {
output.MoveUp(len(secrets))
}
for _, s := range secrets {
output.EraseCurrentLine()
o, err := s.BuildOutput(len(secrets) > 1)
if err != nil {
fmt.Printf("%s: ERROR (%s)\n", s.Name, err)
continue
}
fmt.Println(o)
}
if cfg.OneShot {
fmt.Printf("%s", output)
break
} else {
fmt.Printf("\r%s", output)
}
}
fmt.Println("")
firstRun = false
}
}
func getSecretFromVault() (string, error) {
func getSecretsFromVault() ([]token, error) {
client, err := api.NewClient(&api.Config{
Address: cfg.VaultAddress,
})
if err != nil {
return "", fmt.Errorf("Unable to create client: %s", err)
return nil, fmt.Errorf("Unable to create client: %s", err)
}
client.SetToken(cfg.VaultToken)
data, err := client.Logical().Read(rconfig.Args()[1])
if err != nil {
return "", fmt.Errorf("Unable to read from key %q: %s", rconfig.Args()[1], err)
return nil, fmt.Errorf("Unable to read from key %q: %s", rconfig.Args()[1], err)
}
if data.Data[cfg.Field] == nil {
return "", 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.", rconfig.Args()[1], cfg.Field)
}
return data.Data[cfg.Field].(string), nil
}
func buildOutput(secret string) (string, error) {
// Output: "123456 (Valid 12s)", "123456 (Valid 1s)"
n := time.Now()
code, err := totp.GenerateCode(strings.ToUpper(secret), n)
if err != nil {
return "", err
}
if cfg.NoTime {
return fmt.Sprintf("%s", code), nil
}
remain := 30 - (n.Second() % 30)
return fmt.Sprintf("%s (Valid %ds) ", code, remain), nil
return []token{{Name: rconfig.Args()[1], Secret: data.Data[cfg.Field].(string)}}, nil
}

50
token.go Normal file
View file

@ -0,0 +1,50 @@
package main
import (
"fmt"
"strings"
"time"
"github.com/pquerna/otp/totp"
)
type token struct {
Name string
Secret string
}
func (t token) GetCode(in time.Time) (string, error) {
return totp.GenerateCode(strings.ToUpper(t.Secret), in)
}
func (t token) BuildOutput(showName bool) (string, error) {
// Output: "123456 (Valid 12s)", "123456 (Valid 1s)"
n := time.Now()
code, err := t.GetCode(n)
if err != nil {
return "", err
}
var output string
if showName {
output = fmt.Sprintf("%s: ", t.Name)
}
output = fmt.Sprintf("%s%s", output, code)
if !cfg.NoTime {
remain := 30 - (n.Second() % 30)
output = fmt.Sprintf("%s (Valid %ds) ", output, remain)
}
return output, nil
}
// Sorter interface
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] }