mirror of
https://github.com/Luzifer/vault-totp.git
synced 2024-12-22 22:01:19 +00:00
Prepare multi-token display while retaining same functionality
This commit is contained in:
parent
7c9d67490a
commit
d445ab048d
2 changed files with 84 additions and 36 deletions
70
main.go
70
main.go
|
@ -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)
|
||||
}
|
||||
|
||||
firstRun = false
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Println("")
|
||||
}
|
||||
|
||||
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
50
token.go
Normal 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] }
|
Loading…
Reference in a new issue