// 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 = trustmanager.GenerateRSAKey(rand.Reader, 2048)
		case data.ECDSAKey:
			addedPrivKey, err = trustmanager.GenerateECDSAKey(rand.Reader)
		case data.ED25519Key:
			addedPrivKey, err = trustmanager.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, true)
}

// CreateKey creates a new key inside the cryptoservice for the given role and gun,
// returning the public key.  If the role is a root role, create an x509 key.
func CreateKey(cs signed.CryptoService, gun, role, keyAlgorithm string) (data.PublicKey, error) {
	key, err := cs.Create(role, gun, keyAlgorithm)
	if err != nil {
		return nil, err
	}
	if role == data.CanonicalRootRole {
		start := time.Now().AddDate(0, 0, -1)
		privKey, _, err := cs.GetPrivateKey(key.ID())
		if err != nil {
			return nil, err
		}
		cert, err := cryptoservice.GenerateCertificate(
			privKey, gun, start, start.AddDate(1, 0, 0),
		)
		if err != nil {
			return nil, err
		}
		// Keep the x509 key type consistent with the key's algorithm
		switch keyAlgorithm {
		case data.RSAKey:
			key = data.NewRSAx509PublicKey(trustmanager.CertToPEM(cert))
		case data.ECDSAKey:
			key = data.NewECDSAx509PublicKey(trustmanager.CertToPEM(cert))
		default:
			// This should be impossible because of the Create() call above, but just in case
			return nil, fmt.Errorf("invalid key algorithm type")
		}

	}
	return key, nil
}

func getPubKeys(cs signed.CryptoService, s *data.Signed, role string) ([]data.PublicKey, error) {
	var pubKeys []data.PublicKey
	if role == data.CanonicalRootRole {
		// if this is root metadata, we have to get the keys from the root because they
		// are certs
		root := &data.Root{}
		if err := json.Unmarshal(*s.Signed, root); err != nil {
			return nil, err
		}
		rootRole, ok := root.Roles[data.CanonicalRootRole]
		if !ok || rootRole == nil {
			return nil, tuf.ErrNotLoaded{}
		}
		for _, pubKeyID := range rootRole.KeyIDs {
			pubKeys = append(pubKeys, root.Keys[pubKeyID])
		}
	} else {
		pubKeyIDs := cs.ListKeys(role)
		for _, pubKeyID := range pubKeyIDs {
			pubKey := cs.GetKey(pubKeyID)
			if pubKey != nil {
				pubKeys = append(pubKeys, pubKey)
			}
		}
	}
	return pubKeys, nil
}

// GetOrCreateSnapshotKey either creates a new snapshot key, or returns
// the existing one. Only the PublicKey is returned. The private part
// is held by the CryptoService.
func GetOrCreateSnapshotKey(gun string, store storage.KeyStore, crypto signed.CryptoService, createAlgorithm string) (data.PublicKey, error) {
	keyAlgorithm, public, err := store.GetKey(gun, data.CanonicalSnapshotRole)
	if err == nil {
		return data.NewPublicKey(keyAlgorithm, public), nil
	}

	if _, ok := err.(*storage.ErrNoKey); ok {
		key, err := crypto.Create("snapshot", createAlgorithm)
		if err != nil {
			return nil, err
		}
		logrus.Debug("Creating new snapshot key for ", gun, ". With algo: ", key.Algorithm())
		err = store.SetKey(gun, data.CanonicalSnapshotRole, key.Algorithm(), key.Public())
		if err == nil {
			return key, nil
		}

		if _, ok := err.(*storage.ErrKeyExists); ok {
			keyAlgorithm, public, err = store.GetKey(gun, data.CanonicalSnapshotRole)
			if err != nil {
				return nil, err
			}
			return data.NewPublicKey(keyAlgorithm, public), nil
		}
		return nil, err
	}
	return nil, err
}

// GetOrCreateTimestampKey returns the timestamp key for the gun. It uses the store to
// lookup an existing timestamp key and the crypto to generate a new one if none is
// found. It attempts to handle the race condition that may occur if 2 servers try to
// create the key at the same time by simply querying the store a second time if it
// receives a conflict when writing.
func GetOrCreateTimestampKey(gun string, store storage.MetaStore, crypto signed.CryptoService, fallBackAlgorithm string) (data.PublicKey, error) {
	keyAlgorithm, public, err := store.GetTimestampKey(gun)
	if err == nil {
		return data.NewPublicKey(keyAlgorithm, public), nil
	}

	if _, ok := err.(*storage.ErrNoKey); ok {
		key, err := crypto.Create("timestamp", fallBackAlgorithm)
		if err != nil {
			return nil, err
		}
		logrus.Debug("Creating new timestamp key for ", gun, ". With algo: ", key.Algorithm())
		err = store.SetTimestampKey(gun, key.Algorithm(), key.Public())
		if err == nil {
			return key, nil
		}

		if _, ok := err.(*storage.ErrTimestampKeyExists); ok {
			keyAlgorithm, public, err = store.GetTimestampKey(gun)
			if err != nil {
				return nil, err
			}
			return data.NewPublicKey(keyAlgorithm, public), nil
		}
		return nil, err
	}
	return nil, err
}

// RotateSnapshotKey attempts to rotate a snapshot key in the signer, but might be rate-limited by the signer
func RotateSnapshotKey(gun string, store storage.MetaStore, crypto signed.CryptoService, createAlgorithm string) (data.PublicKey, error) {
	// Always attempt to create a new key, but this might be rate-limited
	key, err := crypto.Create(data.CanonicalSnapshotRole, gun, createAlgorithm)
	if err != nil {
		return nil, err
	}
	logrus.Debug("Created new pending snapshot key ", key.ID(), "to rotate to for ", gun, ". With algo: ", key.Algorithm())
	return key, nil
}

// CopyKeys copies keys of a particular role to a new cryptoservice, and returns that cryptoservice
func CopyKeys(t *testing.T, from signed.CryptoService, roles ...string) signed.CryptoService {
	memKeyStore := trustmanager.NewKeyMemoryStore(passphrase.ConstantRetriever("pass"))
	for _, role := range roles {
		for _, keyID := range from.ListKeys(role) {
			key, _, err := from.GetPrivateKey(keyID)
			require.NoError(t, err)
			memKeyStore.AddKey(trustmanager.KeyInfo{Role: role}, key)
		}
	}
	return cryptoservice.NewCryptoService(memKeyStore)
}

// CreateListKeyCryptoServiceInterfaceBehaviorTests tests expected behavior for
// creating keys in a signed.CryptoService and other read operations on the
// crypto service after keys are present
// 1.  Creating a key succeeds and returns a non-nil public key
// 2.  Listing returns the correct number of keys and right roles
// We allow skipping some tests because for now, signer does not support role checking or listing keys.
func CreateListKeyCryptoServiceInterfaceBehaviorTests(t *testing.T, cs signed.CryptoService, algo string) {
	expectedRolesToKeys := make(map[string]string)
	for i := 0; i < 2; i++ {
		role := data.BaseRoles[i+1]
		createdPubKey, err := cs.Create(role, "docker.io/notary", algo)
		require.NoError(t, err)
		require.NotNil(t, createdPubKey)
		expectedRolesToKeys[role] = createdPubKey.ID()
	}

	testListKeys(t, cs, expectedRolesToKeys)
}

// EmptyCryptoServiceInterfaceBehaviorTests tests expected behavior for
// an empty signed.CryptoService:
// 1.  Getting the public key of a key that doesn't exist should fail
// 2.  Listing an empty cryptoservice returns no keys
// 3.  Removing a non-existent key succeeds (no-op)
func EmptyCryptoServiceInterfaceBehaviorTests(t *testing.T, empty signed.CryptoService) {
	for _, role := range append(data.BaseRoles, "targets/delegation", "invalid") {
		keys := empty.ListKeys(role)
		require.Len(t, keys, 0)
	}
	keys := empty.ListAllKeys()
	require.Len(t, keys, 0)

	require.NoError(t, empty.RemoveKey("nonexistent"))

	require.Nil(t, empty.GetKey("nonexistent"))

	k, role, err := empty.GetPrivateKey("nonexistent")
	require.Error(t, err)
	require.Nil(t, k)
	require.Equal(t, "", role)
}

// The signer does not yet support listing keys or tracking roles, so skip those parts of this test if we're testing
// the signer
func testListKeys(t *testing.T, cs signed.CryptoService, expectedRolesToKeys map[string]string) {
	for _, role := range append(data.BaseRoles, "targets/delegation", "invalid") {
		keys := cs.ListKeys(role)

		if keyID, ok := expectedRolesToKeys[role]; ok {
			require.Len(t, keys, 1)
			require.Equal(t, keyID, keys[0])
		} else {
			require.Len(t, keys, 0)
		}
	}

	keys := cs.ListAllKeys()
	require.Len(t, keys, len(expectedRolesToKeys))
	for role, keyID := range expectedRolesToKeys {
		require.Equal(t, role, keys[keyID])
	}
}

func initRepo(t *testing.T, cryptoService signed.CryptoService, keyDB *keys.KeyDB) *Repo {
	rootKey, err := cryptoService.Create("root", data.ED25519Key)
	assert.NoError(t, err)
	targetsKey, err := cryptoService.Create("targets", data.ED25519Key)
	assert.NoError(t, err)
	snapshotKey, err := cryptoService.Create("snapshot", data.ED25519Key)
	assert.NoError(t, err)
	timestampKey, err := cryptoService.Create("timestamp", data.ED25519Key)
	assert.NoError(t, err)

	keyDB.AddKey(rootKey)
	keyDB.AddKey(targetsKey)
	keyDB.AddKey(snapshotKey)
	keyDB.AddKey(timestampKey)

	rootRole := &data.Role{
		Name: "root",
		RootRole: data.RootRole{
			KeyIDs:    []string{rootKey.ID()},
			Threshold: 1,
		},
	}
	targetsRole := &data.Role{
		Name: "targets",
		RootRole: data.RootRole{
			KeyIDs:    []string{targetsKey.ID()},
			Threshold: 1,
		},
	}
	snapshotRole := &data.Role{
		Name: "snapshot",
		RootRole: data.RootRole{
			KeyIDs:    []string{snapshotKey.ID()},
			Threshold: 1,
		},
	}
	timestampRole := &data.Role{
		Name: "timestamp",
		RootRole: data.RootRole{
			KeyIDs:    []string{timestampKey.ID()},
			Threshold: 1,
		},
	}

	keyDB.AddRole(rootRole)
	keyDB.AddRole(targetsRole)
	keyDB.AddRole(snapshotRole)
	keyDB.AddRole(timestampRole)

	repo := NewRepo(keyDB, cryptoService)
	err = repo.InitRepo(false)
	assert.NoError(t, err)
	return repo
}

func createKey(cs signed.CryptoService, gun, role string) (data.PublicKey, error) {
	key, err := cs.Create(role, data.ECDSAKey)
	if err != nil {
		return nil, err
	}
	if role == data.CanonicalRootRole {
		start := time.Now().AddDate(0, 0, -1)
		privKey, _, err := cs.GetPrivateKey(key.ID())
		if err != nil {
			return nil, err
		}
		cert, err := cryptoservice.GenerateCertificate(
			privKey, gun, start, start.AddDate(1, 0, 0),
		)
		if err != nil {
			return nil, err
		}
		key = data.NewECDSAx509PublicKey(trustmanager.CertToPEM(cert))
	}
	return key, nil
}

// GetOrCreateSnapshotKey either creates a new snapshot key, or returns
// the existing one. Only the PublicKey is returned. The private part
// is held by the CryptoService.
func GetOrCreateSnapshotKey(gun string, store storage.MetaStore, crypto signed.CryptoService, createAlgorithm string) (data.PublicKey, error) {
	_, rootJSON, err := store.GetCurrent(gun, data.CanonicalRootRole)
	if err != nil {
		// If the error indicates we couldn't find the root, create a new key
		if _, ok := err.(storage.ErrNotFound); !ok {
			logrus.Errorf("Error when retrieving root role for GUN %s: %v", gun, err)
			return nil, err
		}
		return crypto.Create(data.CanonicalSnapshotRole, gun, createAlgorithm)
	}

	// If we have a current root, parse out the public key for the snapshot role, and return it
	repoSignedRoot := new(data.SignedRoot)
	if err := json.Unmarshal(rootJSON, repoSignedRoot); err != nil {
		logrus.Errorf("Failed to unmarshal existing root for GUN %s to retrieve snapshot key ID", gun)
		return nil, err
	}

	snapshotRole, err := repoSignedRoot.BuildBaseRole(data.CanonicalSnapshotRole)
	if err != nil {
		logrus.Errorf("Failed to extract snapshot role from root for GUN %s", gun)
		return nil, err
	}

	// We currently only support single keys for snapshot and timestamp, so we can return the first and only key in the map if the signer has it
	for keyID := range snapshotRole.Keys {
		if pubKey := crypto.GetKey(keyID); pubKey != nil {
			return pubKey, nil
		}
	}
	logrus.Debugf("Failed to find any snapshot keys in cryptosigner from root for GUN %s, generating new key", gun)
	return crypto.Create(data.CanonicalSnapshotRole, gun, createAlgorithm)
}

func testGetKey(t *testing.T, cs signed.CryptoService, expectedRolesToKeys map[string]string, algo string) {
	for role, keyID := range expectedRolesToKeys {
		pubKey := cs.GetKey(keyID)
		require.NotNil(t, pubKey)
		require.Equal(t, keyID, pubKey.ID())
		require.Equal(t, algo, pubKey.Algorithm())

		privKey, gotRole, err := cs.GetPrivateKey(keyID)
		require.NoError(t, err)
		require.NotNil(t, privKey)
		require.Equal(t, keyID, privKey.ID())
		require.Equal(t, algo, privKey.Algorithm())
		require.Equal(t, role, gotRole)

		require.NoError(t, cs.RemoveKey(keyID))
		require.Nil(t, cs.GetKey(keyID))
	}
}

func initRepo(t *testing.T, cryptoService signed.CryptoService) *Repo {
	rootKey, err := cryptoService.Create("root", data.ED25519Key)
	assert.NoError(t, err)
	targetsKey, err := cryptoService.Create("targets", data.ED25519Key)
	assert.NoError(t, err)
	snapshotKey, err := cryptoService.Create("snapshot", data.ED25519Key)
	assert.NoError(t, err)
	timestampKey, err := cryptoService.Create("timestamp", data.ED25519Key)
	assert.NoError(t, err)

	rootRole := data.NewBaseRole(
		data.CanonicalRootRole,
		1,
		rootKey,
	)
	targetsRole := data.NewBaseRole(
		data.CanonicalTargetsRole,
		1,
		targetsKey,
	)
	snapshotRole := data.NewBaseRole(
		data.CanonicalSnapshotRole,
		1,
		snapshotKey,
	)
	timestampRole := data.NewBaseRole(
		data.CanonicalTimestampRole,
		1,
		timestampKey,
	)

	repo := NewRepo(cryptoService)
	err = repo.InitRoot(rootRole, timestampRole, snapshotRole, targetsRole, false)
	assert.NoError(t, err)
	_, err = repo.InitTargets(data.CanonicalTargetsRole)
	assert.NoError(t, err)
	err = repo.InitSnapshot()
	assert.NoError(t, err)
	err = repo.InitTimestamp()
	assert.NoError(t, err)
	return repo
}

func initRepo(t *testing.T, cryptoService signed.CryptoService) *Repo {
	rootKey, err := cryptoService.Create("root", testGUN, data.ED25519Key)
	require.NoError(t, err)
	return initRepoWithRoot(t, cryptoService, rootKey)
}