From e5be70796089d1b0f65445e9aec6f2396c4d8e9a Mon Sep 17 00:00:00 2001
From: nbari <nbari@tequila.io>
Date: Sat, 8 Oct 2016 15:40:06 +0200
Subject: [PATCH] Implemented option to select keys -k N

---
 cache.go              | 19 +++++++++++++------
 cmd/ssh-vault/main.go |  6 ++++--
 getkey.go             | 14 +++++++++-----
 vault.go              | 23 ++++++++++++++++++-----
 4 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/cache.go b/cache.go
index 1ef4b13..08ad5ca 100644
--- a/cache.go
+++ b/cache.go
@@ -24,16 +24,23 @@ func Cache() *cache {
 }
 
 // Get return ssh-key
-func (c *cache) Get(u string) (string, error) {
-	uKey := fmt.Sprintf("%s/%s.key", c.dir, u)
+func (c *cache) Get(u string, k int) (string, error) {
+	uKey := fmt.Sprintf("%s/%s.key-%d", c.dir, u, k)
 	if !c.IsFile(uKey) {
-		key, err := GetKey(u)
+		keys, err := GetKey(u)
 		if err != nil {
 			return "", err
 		}
-		err = ioutil.WriteFile(uKey, []byte(key), 0644)
-		if err != nil {
-			log.Println(err)
+		for k, v := range keys {
+			err = ioutil.WriteFile(fmt.Sprintf("%s/%s.key-%d", c.dir, u, k+1),
+				[]byte(v),
+				0644)
+			if err != nil {
+				log.Println(err)
+			}
+		}
+		if !c.IsFile(uKey) {
+			return "", fmt.Errorf("key index not found, try -k with a value between 1 and %d", len(keys))
 		}
 		return uKey, nil
 	}
diff --git a/cmd/ssh-vault/main.go b/cmd/ssh-vault/main.go
index 4a6fe69..abbafb5 100644
--- a/cmd/ssh-vault/main.go
+++ b/cmd/ssh-vault/main.go
@@ -44,8 +44,10 @@ func main() {
 	}
 
 	usr, _ := user.Current()
-	if (*k)[:2] == "~/" {
-		*k = filepath.Join(usr.HomeDir, (*k)[2:])
+	if len(*k) > 2 {
+		if (*k)[:2] == "~/" {
+			*k = filepath.Join(usr.HomeDir, (*k)[2:])
+		}
 	}
 
 	vault, err := sv.New(*k, *u, flag.Arg(0), flag.Arg(1))
diff --git a/getkey.go b/getkey.go
index 7393f98..55a644e 100644
--- a/getkey.go
+++ b/getkey.go
@@ -13,7 +13,7 @@ import (
 const GITHUB = "https://github.com"
 
 // GetKey fetches ssh-key from url
-func GetKey(u string) (string, error) {
+func GetKey(u string) ([]string, error) {
 	client := &http.Client{}
 	// create a new request
 	req, _ := http.NewRequest("GET", fmt.Sprintf("%s/%s.keys",
@@ -23,19 +23,23 @@ func GetKey(u string) (string, error) {
 	req.Header.Set("User-Agent", "ssh-vault")
 	res, err := client.Do(req)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	defer res.Body.Close()
 	reader := bufio.NewReader(res.Body)
 	tp := textproto.NewReader(reader)
+	keys := []string{}
 	for {
 		if line, err := tp.ReadLine(); err != nil {
 			if err == io.EOF {
-				return "", fmt.Errorf("key %q not found", u)
+				if len(keys) == 0 {
+					return nil, fmt.Errorf("key %q not found", u)
+				}
+				return keys, nil
 			}
-			return "", err
+			return nil, err
 		} else if strings.HasPrefix(line, "ssh-rsa") {
-			return line, nil
+			keys = append(keys, line)
 		}
 	}
 }
diff --git a/vault.go b/vault.go
index 1328a16..cbf0aa6 100644
--- a/vault.go
+++ b/vault.go
@@ -9,13 +9,13 @@ import (
 	"encoding/pem"
 	"fmt"
 	"os/exec"
+	"strconv"
 	"strings"
 )
 
 // Vault structure
 type vault struct {
 	key         string
-	option      string
 	vault       string
 	PublicKey   *rsa.PublicKey
 	Fingerprint string
@@ -30,17 +30,30 @@ func New(k, u, o, v string) (*vault, error) {
 	)
 	cache := Cache()
 	if u != "" {
-		keyPath, err = cache.Get(u)
+		// use -k N where N is the index to use when multiple keys
+		// are available
+		var ki int
+		if ki, err = strconv.Atoi(k); err != nil {
+			ki = 1
+		}
+		if ki <= 1 {
+			ki = 1
+		}
+		keyPath, err = cache.Get(u, ki)
 		if err != nil {
 			return nil, err
 		}
 	} else if !cache.IsFile(keyPath) {
 		return nil, fmt.Errorf("key not found or unable to read")
 	}
+	if o == "create" {
+		if cache.IsFile(v) {
+			return nil, fmt.Errorf("File already exists: %q", v)
+		}
+	}
 	return &vault{
-		key:    keyPath,
-		option: o,
-		vault:  v,
+		key:   keyPath,
+		vault: v,
 	}, nil
 }