add checks to CLI command for role and gun

Signed-off-by: Riyaz Faizullabhoy <riyaz.faizullabhoy@docker.com>

cmd/notary/keys.go

@@ -20,6 +20,7 @@ import (
 	"github.com/docker/notary/tuf/data"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"io/ioutil"
 )
 
 var cmdKeyTemplate = usageTemplate{
@@ -364,11 +365,49 @@ func (k *keyCommander) keysImport(cmd *cobra.Command, args []string) error {
 	}
 	defer importFile.Close()
 
+	pemBytes, err := ioutil.ReadAll(importFile)
+	if err != nil {
+		return fmt.Errorf("Error reading input file: %v", err)
+	}
+
+	pemRole := trustmanager.ReadRoleFromPEM(pemBytes)
+
+	// Rewind after reading the first time
+	_, err = importFile.Seek(0, 0)
+	if err != nil {
+		return fmt.Errorf("Error reading input file: %v", err)
+	}
+
+	if pemRole != "" && !data.ValidRole(pemRole) {
+		return fmt.Errorf("Invalid role specified for key: %s", pemRole)
+	}
+
+	if k.keysImportRole != "" && !data.ValidRole(k.keysImportRole) {
+		return fmt.Errorf("Invalid role specified for key: %s", k.keysImportRole)
+	}
+
+	// If the PEM key doesn't have a role in it, we must have --role set
+	if pemRole == "" && k.keysImportRole == "" {
+		return fmt.Errorf("Could not infer role, and no role was specified for key")
+	}
+
+	// If both  PEM role and a --role are provided and they don't match, error
+	if pemRole != "" && k.keysImportRole != "" && pemRole != k.keysImportRole {
+		return fmt.Errorf("Specified role %s does not match role %s in PEM headers", k.keysImportRole, pemRole)
+	}
+
+	// If we're importing to targets or snapshot, we need a GUN
+	if (k.keysImportRole == data.CanonicalTargetsRole || k.keysImportRole == data.CanonicalSnapshotRole) ||
+		(pemRole == data.CanonicalTargetsRole || pemRole == data.CanonicalSnapshotRole) &&
+			k.keysImportGUN == "" {
+		return fmt.Errorf("Must specify GUN for %s key", k.keysImportRole)
+	}
+
 	cs := cryptoservice.NewCryptoService(k.keysImportGUN, ks...)
 	if k.keysImportRole == data.CanonicalRootRole {
 		err = cs.ImportRootKey(importFile)
 	} else {
-		err = cs.ImportRoleKey(importFile, k.keysImportRole, k.retriever)
+		err = cs.ImportRoleKey(importFile, k.keysImportRole, k.getRetriever())
 	}
 
 	if err != nil {

cmd/notary/keys_test.go

@@ -433,3 +433,130 @@ func TestChangeKeyPassphraseNonexistentID(t *testing.T) {
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "could not retrieve local key for key ID provided")
 }
+
+func TestKeyImportInvalidFlagRole(t *testing.T) {
+	k := &keyCommander{
+		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+		keysImportRole: "invalid",
+	}
+	tempFileName := generateTempTestKeyFile(t, "invalid")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Invalid role specified for key:")
+}
+
+func TestKeyImportInvalidPEMRole(t *testing.T) {
+	k := &keyCommander{
+		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+		keysImportRole: "targets",
+	}
+	tempFileName := generateTempTestKeyFile(t, "invalid")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Invalid role specified for key:")
+}
+
+func TestKeyImportMismatchingRoles(t *testing.T) {
+	k := &keyCommander{
+		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+		keysImportRole: "targets",
+	}
+	tempFileName := generateTempTestKeyFile(t, "snapshot")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "does not match role")
+}
+
+func TestKeyImportNoGUNForTargetsPEM(t *testing.T) {
+	k := &keyCommander{
+		configGetter: func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever: func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+	}
+	tempFileName := generateTempTestKeyFile(t, "targets")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Must specify GUN")
+}
+
+func TestKeyImportNoGUNForSnapshotPEM(t *testing.T) {
+	k := &keyCommander{
+		configGetter: func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever: func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+	}
+	tempFileName := generateTempTestKeyFile(t, "snapshot")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Must specify GUN")
+}
+
+func TestKeyImportNoGUNForTargetsFlag(t *testing.T) {
+	k := &keyCommander{
+		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+		keysImportRole: "targets",
+	}
+	tempFileName := generateTempTestKeyFile(t, "")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Must specify GUN")
+}
+
+func TestKeyImportNoGUNForSnapshotFlag(t *testing.T) {
+	k := &keyCommander{
+		configGetter:   func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever:   func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+		keysImportRole: "snapshot",
+	}
+	tempFileName := generateTempTestKeyFile(t, "")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Must specify GUN")
+}
+
+func TestKeyImportNoRole(t *testing.T) {
+	k := &keyCommander{
+		configGetter: func() (*viper.Viper, error) { return viper.New(), nil },
+		getRetriever: func() passphrase.Retriever { return passphrase.ConstantRetriever("pass") },
+	}
+	tempFileName := generateTempTestKeyFile(t, "")
+	defer os.Remove(tempFileName)
+
+	err := k.keysImport(&cobra.Command{}, []string{tempFileName})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Could not infer role, and no role was specified for key")
+}
+
+func generateTempTestKeyFile(t *testing.T, role string) string {
+	privKey, err := trustmanager.GenerateECDSAKey(rand.Reader)
+	if err != nil {
+		return ""
+	}
+	keyBytes, err := trustmanager.KeyToPEM(privKey, role)
+	assert.NoError(t, err)
+
+	tempPrivFile, err := ioutil.TempFile("/tmp", "privfile")
+	assert.NoError(t, err)
+
+	// Write the private key to a file so we can import it
+	_, err = tempPrivFile.Write(keyBytes)
+	assert.NoError(t, err)
+	tempPrivFile.Close()
+	return tempPrivFile.Name()
+}

cryptoservice/import_export.go

@@ -112,6 +112,7 @@ func (cs *CryptoService) ImportRootKey(source io.Reader) error {
 // the key ID.
 func (cs *CryptoService) ImportRoleKey(source io.Reader, role string, newPassphraseRetriever passphrase.Retriever) error {
 	pemBytes, err := ioutil.ReadAll(source)
+
 	if err != nil {
 		return err
 	}

trustmanager/x509utils.go

@@ -514,6 +514,20 @@ func EncryptPrivateKey(key data.PrivateKey, role, passphrase string) ([]byte, er
 	return pem.EncodeToMemory(encryptedPEMBlock), nil
 }
 
+// ReadRoleFromPEM returns the value from the role PEM header, if it exists
+// If it doesn't exist, returns nil
+func ReadRoleFromPEM(pemBytes []byte) string {
+	pemBlock, _ := pem.Decode(pemBytes)
+	if pemBlock.Headers == nil {
+		return ""
+	}
+	role, ok := pemBlock.Headers["role"]
+	if !ok {
+		return ""
+	}
+	return role
+}
+
 // CertToKey transforms a single input certificate into its corresponding
 // PublicKey
 func CertToKey(cert *x509.Certificate) data.PublicKey {