mirror of
https://github.com/Luzifer/vault-totp.git
synced 2024-12-23 06:11: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"
|
"io/ioutil"
|
||||||
"log"
|
"log"
|
||||||
"os"
|
"os"
|
||||||
"strings"
|
"sort"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/Luzifer/rconfig"
|
"github.com/Luzifer/rconfig"
|
||||||
"github.com/hashicorp/vault/api"
|
"github.com/hashicorp/vault/api"
|
||||||
homedir "github.com/mitchellh/go-homedir"
|
homedir "github.com/mitchellh/go-homedir"
|
||||||
"github.com/pquerna/otp/totp"
|
"github.com/sethgrid/curse"
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
|
@ -67,72 +67,70 @@ func main() {
|
||||||
}
|
}
|
||||||
|
|
||||||
var (
|
var (
|
||||||
secret string
|
secrets []token
|
||||||
err error
|
err error
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.TestSecret == "" {
|
if cfg.TestSecret == "" {
|
||||||
secret, err = getSecretFromVault()
|
secrets, err = getSecretsFromVault()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Fatalf("Error to retrieve secret: %s", err)
|
log.Fatalf("Error to retrieve secret: %s", err)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
secret = cfg.TestSecret
|
secrets = []token{{Name: "Test", Secret: cfg.TestSecret}}
|
||||||
}
|
}
|
||||||
|
|
||||||
for range time.Tick(250 * time.Millisecond) {
|
sort.Sort(tokenList(secrets))
|
||||||
output, err := buildOutput(secret)
|
|
||||||
|
output, err := curse.New()
|
||||||
if err != nil {
|
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 {
|
if cfg.OneShot {
|
||||||
fmt.Printf("%s", output)
|
|
||||||
break
|
break
|
||||||
} else {
|
}
|
||||||
fmt.Printf("\r%s", output)
|
|
||||||
|
firstRun = false
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fmt.Println("")
|
func getSecretsFromVault() ([]token, error) {
|
||||||
}
|
|
||||||
|
|
||||||
func getSecretFromVault() (string, error) {
|
|
||||||
client, err := api.NewClient(&api.Config{
|
client, err := api.NewClient(&api.Config{
|
||||||
Address: cfg.VaultAddress,
|
Address: cfg.VaultAddress,
|
||||||
})
|
})
|
||||||
|
|
||||||
if err != nil {
|
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)
|
client.SetToken(cfg.VaultToken)
|
||||||
|
|
||||||
data, err := client.Logical().Read(rconfig.Args()[1])
|
data, err := client.Logical().Read(rconfig.Args()[1])
|
||||||
if err != nil {
|
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 {
|
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
|
return []token{{Name: rconfig.Args()[1], Secret: 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
|
|
||||||
}
|
}
|
||||||
|
|
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