// path and encrypted key should succeed, tests gun inference from path as well
func TestEncryptedKeyImportSuccess(t *testing.T) {
	s := NewTestImportStore()

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	originalKey := privKey.Private()
	require.NoError(t, err)

	pemBytes, err := utils.EncryptPrivateKey(privKey, data.CanonicalSnapshotRole, "somegun", cannedPassphrase)
	require.NoError(t, err)

	b, _ := pem.Decode(pemBytes)
	b.Headers["path"] = privKey.ID()
	pemBytes = pem.EncodeToMemory(b)

	in := bytes.NewBuffer(pemBytes)

	_ = ImportKeys(in, []Importer{s}, "", "", passphraseRetriever)
	require.Len(t, s.data, 1)

	keyBytes := s.data[privKey.ID()]

	bFinal, bRest := pem.Decode(keyBytes)
	require.Equal(t, "somegun", bFinal.Headers["gun"])
	require.Len(t, bRest, 0)

	// we should fail to parse it without the passphrase
	privKey, err = utils.ParsePEMPrivateKey(keyBytes, "")
	require.Equal(t, err, errors.New("could not decrypt private key"))
	require.Nil(t, privKey)

	// we should succeed to parse it with the passphrase
	privKey, err = utils.ParsePEMPrivateKey(keyBytes, cannedPassphrase)
	require.NoError(t, err)
	require.Equal(t, originalKey, privKey.Private())
}

func testRemoveKeyWithRole(t *testing.T, role string) {
	gun := "docker.com/notary"
	testExt := "key"

	// Temporary directory where test files will be created
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	require.NoError(t, err, "failed to create a temporary directory")
	defer os.RemoveAll(tempBaseDir)

	// Create our store
	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Since we're generating this manually we need to add the extension '.'
	expectedFilePath := filepath.Join(tempBaseDir, notary.PrivDir, privKey.ID()+"."+testExt)

	err = store.AddKey(KeyInfo{Role: role, Gun: gun}, privKey)
	require.NoError(t, err, "failed to add key to store")

	// Check to see if file exists
	_, err = ioutil.ReadFile(expectedFilePath)
	require.NoError(t, err, "expected file not found")

	// Call remove key
	err = store.RemoveKey(privKey.ID())
	require.NoError(t, err, "unable to remove key")

	// Check to see if file still exists
	_, err = ioutil.ReadFile(expectedFilePath)
	require.Error(t, err, "file should not exist")
}

func TestEncryption(t *testing.T) {
	s := NewTestImportStore()

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	originalKey := privKey.Private()
	require.NoError(t, err)

	pemBytes, err := utils.EncryptPrivateKey(privKey, "", "", "")
	require.NoError(t, err)

	in := bytes.NewBuffer(pemBytes)

	_ = ImportKeys(in, []Importer{s}, "", "", passphraseRetriever)
	require.Len(t, s.data, 1)

	shouldBeEnc, ok := s.data[privKey.ID()]
	// we should have got a key imported to this location
	require.True(t, ok)

	// we should fail to parse it without the passphrase
	privKey, err = utils.ParsePEMPrivateKey(shouldBeEnc, "")
	require.Equal(t, err, errors.New("could not decrypt private key"))
	require.Nil(t, privKey)

	// we should succeed to parse it with the passphrase
	privKey, err = utils.ParsePEMPrivateKey(shouldBeEnc, cannedPassphrase)
	require.NoError(t, err)
	require.Equal(t, originalKey, privKey.Private())
}

// AddGetKeyCryptoServiceInterfaceBehaviorTests tests expected behavior for
// adding keys in a signed.CryptoService and other read operations on the
// crypto service after keys are present
// 1.  Adding a key succeeds
// 2.  Getting the key should return the same key, without error
// 3.  Removing the key succeeds
func AddGetKeyCryptoServiceInterfaceBehaviorTests(t *testing.T, cs signed.CryptoService, algo string) {
	expectedRolesToKeys := make(map[string]string)
	for i := 0; i < 2; i++ {
		var (
			addedPrivKey data.PrivateKey
			err          error
		)
		role := data.BaseRoles[i+1]
		switch algo {
		case data.RSAKey:
			addedPrivKey, err = utils.GenerateRSAKey(rand.Reader, 2048)
		case data.ECDSAKey:
			addedPrivKey, err = utils.GenerateECDSAKey(rand.Reader)
		case data.ED25519Key:
			addedPrivKey, err = utils.GenerateED25519Key(rand.Reader)
		default:
			require.FailNow(t, "invalid algorithm %s", algo)
		}
		require.NoError(t, err)
		require.NotNil(t, addedPrivKey)
		require.NoError(t, cs.AddKey(role, "docker.io/notary", addedPrivKey))
		expectedRolesToKeys[role] = addedPrivKey.ID()
	}

	testGetKey(t, cs, expectedRolesToKeys, algo)
}

func testAddKey(t *testing.T, store trustmanager.KeyStore) (data.PrivateKey, error) {
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err)

	err = store.AddKey(trustmanager.KeyInfo{Role: data.CanonicalRootRole, Gun: ""}, privKey)
	return privKey, err
}

func TestAddGetKeyInfoMemStore(t *testing.T) {
	gun := "docker.com/notary"

	// Create our store
	store := NewKeyMemoryStore(passphraseRetriever)

	rootKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: data.CanonicalRootRole, Gun: ""}, rootKey)
	require.NoError(t, err, "failed to add key to store")

	// Get and validate key info
	rootInfo, err := store.GetKeyInfo(rootKey.ID())
	require.NoError(t, err)
	require.Equal(t, data.CanonicalRootRole, rootInfo.Role)
	require.Equal(t, "", rootInfo.Gun)

	targetsKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: data.CanonicalTargetsRole, Gun: gun}, targetsKey)
	require.NoError(t, err, "failed to add key to store")

	// Get and validate key info
	targetsInfo, err := store.GetKeyInfo(targetsKey.ID())
	require.NoError(t, err)
	require.Equal(t, data.CanonicalTargetsRole, targetsInfo.Role)
	require.Equal(t, gun, targetsInfo.Gun)

	delgKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: "targets/delegation", Gun: gun}, delgKey)
	require.NoError(t, err, "failed to add key to store")

	// Get and validate key info
	delgInfo, err := store.GetKeyInfo(delgKey.ID())
	require.NoError(t, err)
	require.Equal(t, "targets/delegation", delgInfo.Role)
	require.Equal(t, "", delgInfo.Gun)
}

func TestKeysAreCached(t *testing.T) {
	gun := "docker.com/notary"
	testAlias := "alias"

	// Temporary directory where test files will be created
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	require.NoError(t, err, "failed to create a temporary directory")
	defer os.RemoveAll(tempBaseDir)

	var countingPassphraseRetriever notary.PassRetriever

	numTimesCalled := 0
	countingPassphraseRetriever = func(keyId, alias string, createNew bool, attempts int) (passphrase string, giveup bool, err error) {
		numTimesCalled++
		return "password", false, nil
	}

	// Create our store
	store, err := NewKeyFileStore(tempBaseDir, countingPassphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: testAlias, Gun: gun}, privKey)
	require.NoError(t, err, "failed to add key to store")

	require.Equal(t, 1, numTimesCalled, "numTimesCalled should have been 1")

	// Call the AddKey function
	privKey2, _, err := store.GetKey(privKey.ID())
	require.NoError(t, err, "failed to add key to store")

	require.Equal(t, privKey.Public(), privKey2.Public(), "cachedPrivKey should be the same as the added privKey")
	require.Equal(t, privKey.Private(), privKey2.Private(), "cachedPrivKey should be the same as the added privKey")
	require.Equal(t, 1, numTimesCalled, "numTimesCalled should be 1 -- no additional call to passphraseRetriever")

	// Create a new store
	store2, err := NewKeyFileStore(tempBaseDir, countingPassphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	// Call the GetKey function
	privKey3, _, err := store2.GetKey(privKey.ID())
	require.NoError(t, err, "failed to get key from store")

	require.Equal(t, privKey2.Private(), privKey3.Private(), "privkey from store1 should be the same as privkey from store2")
	require.Equal(t, privKey2.Public(), privKey3.Public(), "privkey from store1 should be the same as privkey from store2")
	require.Equal(t, 2, numTimesCalled, "numTimesCalled should be 2 -- one additional call to passphraseRetriever")

	// Call the GetKey function a bunch of times
	for i := 0; i < 10; i++ {
		_, _, err := store2.GetKey(privKey.ID())
		require.NoError(t, err, "failed to get key from store")
	}
	require.Equal(t, 2, numTimesCalled, "numTimesCalled should be 2 -- no additional call to passphraseRetriever")
}

func TestListKeys(t *testing.T) {
	testName := "docker.com/notary/root"
	perms := os.FileMode(0755)

	// Temporary directory where test files will be created
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	require.NoError(t, err, "failed to create a temporary directory")
	defer os.RemoveAll(tempBaseDir)

	// Create our store
	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	roles := append(data.BaseRoles, "targets/a", "invalidRoleName")

	for i, role := range roles {
		// Make a new key for each role
		privKey, err := utils.GenerateECDSAKey(rand.Reader)
		require.NoError(t, err, "could not generate private key")

		// Call the AddKey function
		gun := filepath.Dir(testName)
		err = store.AddKey(KeyInfo{Role: role, Gun: gun}, privKey)
		require.NoError(t, err, "failed to add key to store")

		// Check to see if the keystore lists this key
		keyMap := store.ListKeys()

		// Expect to see exactly one key in the map
		require.Len(t, keyMap, i+1)
		// Expect to see privKeyID inside of the map
		listedInfo, ok := keyMap[privKey.ID()]
		require.True(t, ok)
		require.Equal(t, role, listedInfo.Role)
	}

	// Write an invalid filename to the directory
	filePath := filepath.Join(tempBaseDir, notary.PrivDir, "fakekeyname.key")
	err = ioutil.WriteFile(filePath, []byte("data"), perms)
	require.NoError(t, err, "failed to write test file")

	// Check to see if the keystore still lists two keys
	keyMap := store.ListKeys()
	require.Len(t, keyMap, len(roles))

	// Check that ListKeys() returns a copy of the state
	// so modifying its returned information does not change the underlying store's keyInfo
	for keyID := range keyMap {
		delete(keyMap, keyID)
		_, err = store.GetKeyInfo(keyID)
		require.NoError(t, err)
	}
}

func generateValidTestCert() (*x509.Certificate, string, error) {
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	if err != nil {
		return nil, "", err
	}
	keyID := privKey.ID()
	startTime := time.Now()
	endTime := startTime.AddDate(10, 0, 0)
	cert, err := cryptoservice.GenerateCertificate(privKey, "gun", startTime, endTime)
	if err != nil {
		return nil, "", err
	}
	return cert, keyID, nil
}

// no path and encrypted key import should fail
func TestEncryptedKeyImportFail(t *testing.T) {
	s := NewTestImportStore()

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err)

	pemBytes, err := utils.EncryptPrivateKey(privKey, data.CanonicalRootRole, "", cannedPassphrase)
	require.NoError(t, err)

	in := bytes.NewBuffer(pemBytes)

	_ = ImportKeys(in, []Importer{s}, "", "", passphraseRetriever)
	require.Len(t, s.data, 0)
}

func generateExpiredTestCert() (*x509.Certificate, string, error) {
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	if err != nil {
		return nil, "", err
	}
	keyID := privKey.ID()
	// Set to Unix time 0 start time, valid for one more day
	startTime := time.Unix(0, 0)
	endTime := startTime.AddDate(0, 0, 1)
	cert, err := cryptoservice.GenerateCertificate(privKey, "gun", startTime, endTime)
	if err != nil {
		return nil, "", err
	}
	return cert, keyID, nil
}

// testGetDecryptedWithInvalidPassphrase takes two keystores so it can add to
// one and get from the other (to work around caching)
func testGetDecryptedWithInvalidPassphrase(t *testing.T, store KeyStore, newStore KeyStore, expectedFailureType interface{}) {
	testAlias := data.CanonicalRootRole

	// Generate a new random RSA Key
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: testAlias, Gun: ""}, privKey)
	require.NoError(t, err, "failed to add key to store")

	// Try to decrypt the file with an invalid passphrase
	_, _, err = newStore.GetKey(privKey.ID())
	require.Error(t, err, "expected error while decrypting the content due to invalid passphrase")
	require.IsType(t, err, expectedFailureType)
}

func TestCannotSignWithKeyThatDoesntExist(t *testing.T) {
	memStore := trustmanager.NewKeyMemoryStore(constPass)

	_, conn, cleanup := setUpSignerClient(t, setUpSignerServer(t, memStore))
	defer cleanup()

	key, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate key")

	remotePrivKey := client.NewRemotePrivateKey(data.PublicKeyFromPrivate(key), pb.NewSignerClient(conn))

	msg := []byte("message!")
	_, err = remotePrivKey.Sign(rand.Reader, msg, nil)
	require.Error(t, err)
	// error translated into grpc error, so compare the text
	require.Equal(t, trustmanager.ErrKeyNotFound{KeyID: key.ID()}.Error(), grpc.ErrorDesc(err))
}

// helper function to generate private keys for the signer databases - does not implement RSA since that is not
// supported by the signer
func generatePrivateKey(algorithm string) (data.PrivateKey, error) {
	var privKey data.PrivateKey
	var err error
	switch algorithm {
	case data.ECDSAKey:
		privKey, err = utils.GenerateECDSAKey(rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("failed to generate EC key: %v", err)
		}
	case data.ED25519Key:
		privKey, err = utils.GenerateED25519Key(rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("failed to generate ED25519 key: %v", err)
		}
	default:
		return nil, fmt.Errorf("private key type not supported for key generation: %s", algorithm)
	}
	return privKey, nil
}

func TestAddGetKeyMemStore(t *testing.T) {
	testAlias := data.CanonicalRootRole

	// Create our store
	store := NewKeyMemoryStore(passphraseRetriever)

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: testAlias, Gun: ""}, privKey)
	require.NoError(t, err, "failed to add key to store")

	// Check to see if file exists
	retrievedKey, retrievedAlias, err := store.GetKey(privKey.ID())
	require.NoError(t, err, "failed to get key from store")

	require.Equal(t, retrievedAlias, testAlias)
	require.Equal(t, retrievedKey.Public(), privKey.Public())
	require.Equal(t, retrievedKey.Private(), privKey.Private())
}

func TestGetDecryptedWithTamperedCipherText(t *testing.T) {
	testExt := "key"
	testAlias := data.CanonicalRootRole

	// Temporary directory where test files will be created
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	require.NoError(t, err, "failed to create a temporary directory")
	defer os.RemoveAll(tempBaseDir)

	// Create our FileStore
	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	// Generate a new Private Key
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Call the AddEncryptedKey function
	err = store.AddKey(KeyInfo{Role: testAlias, Gun: ""}, privKey)
	require.NoError(t, err, "failed to add key to store")

	// Since we're generating this manually we need to add the extension '.'
	expectedFilePath := filepath.Join(tempBaseDir, notary.PrivDir, privKey.ID()+"."+testExt)

	// Get file description, open file
	fp, err := os.OpenFile(expectedFilePath, os.O_WRONLY, 0600)
	require.NoError(t, err, "expected file not found")

	// Tamper the file
	fp.WriteAt([]byte("a"), int64(1))

	// Recreate the KeyFileStore to avoid caching
	store, err = NewKeyFileStore(tempBaseDir, passphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	// Try to decrypt the file
	_, _, err = store.GetKey(privKey.ID())
	require.Error(t, err, "expected error while decrypting the content due to invalid cipher text")
}

// Create is used to generate keys for targets, snapshots and timestamps
func (cs *CryptoService) Create(role, gun, algorithm string) (data.PublicKey, error) {
	var privKey data.PrivateKey
	var err error

	switch algorithm {
	case data.RSAKey:
		privKey, err = utils.GenerateRSAKey(rand.Reader, notary.MinRSABitSize)
		if err != nil {
			return nil, fmt.Errorf("failed to generate RSA key: %v", err)
		}
	case data.ECDSAKey:
		privKey, err = utils.GenerateECDSAKey(rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("failed to generate EC key: %v", err)
		}
	case data.ED25519Key:
		privKey, err = utils.GenerateED25519Key(rand.Reader)
		if err != nil {
			return nil, fmt.Errorf("failed to generate ED25519 key: %v", err)
		}
	default:
		return nil, fmt.Errorf("private key type not supported for key generation: %s", algorithm)
	}
	logrus.Debugf("generated new %s key for role: %s and keyID: %s", algorithm, role, privKey.ID())

	// Store the private key into our keystore
	for _, ks := range cs.keyStores {
		err = ks.AddKey(trustmanager.KeyInfo{Role: role, Gun: gun}, privKey)
		if err == nil {
			return data.PublicKeyFromPrivate(privKey), nil
		}
	}
	if err != nil {
		return nil, fmt.Errorf("failed to add key to filestore: %v", err)
	}

	return nil, fmt.Errorf("keystores would not accept new private keys for unknown reasons")
}

func testAddKeyWithRole(t *testing.T, role string) {
	gun := "docker.com/notary"
	testExt := "key"

	// Temporary directory where test files will be created
	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
	require.NoError(t, err, "failed to create a temporary directory")
	defer os.RemoveAll(tempBaseDir)
	// Create our store
	store, err := NewKeyFileStore(tempBaseDir, passphraseRetriever)
	require.NoError(t, err, "failed to create new key filestore")

	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate private key")

	// Since we're generating this manually we need to add the extension '.'
	expectedFilePath := filepath.Join(tempBaseDir, notary.PrivDir, privKey.ID()+"."+testExt)

	// Call the AddKey function
	err = store.AddKey(KeyInfo{Role: role, Gun: gun}, privKey)
	require.NoError(t, err, "failed to add key to store")

	// Check to see if file exists
	b, err := ioutil.ReadFile(expectedFilePath)
	require.NoError(t, err, "expected file not found")
	require.Contains(t, string(b), "-----BEGIN EC PRIVATE KEY-----")

	// Check that we have the role and gun info for this key's ID
	keyInfo, ok := store.keyInfoMap[privKey.ID()]
	require.True(t, ok)
	require.Equal(t, role, keyInfo.Role)
	if role == data.CanonicalRootRole || data.IsDelegation(role) || !data.ValidRole(role) {
		require.Empty(t, keyInfo.Gun)
	} else {
		require.Equal(t, gun, keyInfo.Gun)
	}
}

func TestGetPrivateKeyAndSignWithExistingKey(t *testing.T) {
	key, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate key")

	memStore := trustmanager.NewKeyMemoryStore(constPass)
	err = memStore.AddKey(trustmanager.KeyInfo{Role: data.CanonicalTimestampRole, Gun: "gun"}, key)
	require.NoError(t, err, "could not add key to store")

	signerClient, _, cleanup := setUpSignerClient(t, setUpSignerServer(t, memStore))
	defer cleanup()

	privKey, role, err := signerClient.GetPrivateKey(key.ID())
	require.NoError(t, err)
	require.Equal(t, data.CanonicalTimestampRole, role)
	require.NotNil(t, privKey)

	msg := []byte("message!")
	sig, err := privKey.Sign(rand.Reader, msg, nil)
	require.NoError(t, err)

	err = signed.Verifiers[data.ECDSASignature].Verify(
		data.PublicKeyFromPrivate(key), sig, msg)
	require.NoError(t, err)
}

func TestGenerateCertificate(t *testing.T) {
	privKey, err := utils.GenerateECDSAKey(rand.Reader)
	require.NoError(t, err, "could not generate key")

	keyStore := trustmanager.NewKeyMemoryStore(passphraseRetriever)

	err = keyStore.AddKey(trustmanager.KeyInfo{Role: data.CanonicalRootRole, Gun: ""}, privKey)
	require.NoError(t, err, "could not add key to store")

	// Check GenerateCertificate method
	gun := "docker.com/notary"
	startTime := time.Now()
	cert, err := GenerateCertificate(privKey, gun, startTime, startTime.AddDate(10, 0, 0))
	require.NoError(t, err, "could not generate certificate")

	// Check public key
	ecdsaPrivateKey, err := x509.ParseECPrivateKey(privKey.Private())
	require.NoError(t, err)
	ecdsaPublicKey := ecdsaPrivateKey.Public()
	require.Equal(t, ecdsaPublicKey, cert.PublicKey)

	// Check CommonName
	require.Equal(t, cert.Subject.CommonName, gun)
}