mev-beta/pkg/security/keymanager_test.go

package security

import (
	"math/big"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/fraktal/mev-beta/internal/logger"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestNewKeyManager tests the creation of a new KeyManager
func TestNewKeyManager(t *testing.T) {
	// Test with valid configuration
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)

	require.NoError(t, err)
	assert.NotNil(t, km)
	assert.NotNil(t, km.keystore)
	assert.NotNil(t, km.keys)
	assert.NotNil(t, km.encryptionKey)
	assert.Equal(t, config, km.config)

	// Test with nil configuration (should use defaults)
	km2, err := NewKeyManager(nil, log)
	require.NoError(t, err)
	assert.NotNil(t, km2)
	assert.NotNil(t, km2.config)
	assert.NotEmpty(t, km2.config.KeystorePath)
}

// TestNewKeyManagerInvalidConfig tests error cases for KeyManager creation
func TestNewKeyManagerInvalidConfig(t *testing.T) {
	log := logger.New("info", "text", "")

	// Test with empty encryption key
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "",
	}

	km, err := NewKeyManager(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "encryption key cannot be empty")

	// Test with short encryption key
	config = &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "short",
	}

	km, err = NewKeyManager(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "encryption key must be at least 32 characters")

	// Test with empty keystore path
	config = &KeyManagerConfig{
		KeystorePath:  "",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	km, err = NewKeyManager(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "keystore path cannot be empty")
}

// TestGenerateKey tests key generation functionality
func TestGenerateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_generate",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Test generating a trading key
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(1000000000000000000), // 1 ETH
	}

	address, err := km.GenerateKey("trading", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, address)

	// Verify the key exists
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "trading", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)
	assert.WithinDuration(t, time.Now(), keyInfo.CreatedAt, time.Second)
	assert.WithinDuration(t, time.Now(), keyInfo.LastUsed, time.Second)
	assert.Equal(t, int64(0), keyInfo.UsageCount)

	// Test generating an emergency key (should have expiration)
	emergencyAddress, err := km.GenerateKey("emergency", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, emergencyAddress)

	emergencyKeyInfo, err := km.GetKeyInfo(emergencyAddress)
	require.NoError(t, err)
	assert.NotNil(t, emergencyKeyInfo.ExpiresAt)
	assert.True(t, emergencyKeyInfo.ExpiresAt.After(time.Now()))
}

// TestImportKey tests key import functionality
func TestImportKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_import",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)
	privateKeyHex := common.Bytes2Hex(crypto.FromECDSA(privateKey))

	// Import the key
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    false,
		MaxTransferWei: nil,
	}

	address, err := km.ImportKey(privateKeyHex, "test", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, address)

	// Verify the imported key
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "test", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)

	// Test importing invalid key
	_, err = km.ImportKey("invalid_private_key", "test", permissions)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid private key")

	// Test importing duplicate key
	_, err = km.ImportKey(privateKeyHex, "duplicate", permissions)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key already exists")
}

// TestListKeys tests key listing functionality
func TestListKeys(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_list",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Initially should be empty
	keys := km.ListKeys()
	assert.Empty(t, keys)

	// Generate a few keys
	permissions := KeyPermissions{CanSign: true}
	addr1, err := km.GenerateKey("test1", permissions)
	require.NoError(t, err)

	addr2, err := km.GenerateKey("test2", permissions)
	require.NoError(t, err)

	// Check that both keys are listed
	keys = km.ListKeys()
	assert.Len(t, keys, 2)
	assert.Contains(t, keys, addr1)
	assert.Contains(t, keys, addr2)
}

// TestGetKeyInfo tests key information retrieval
func TestGetKeyInfo(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_info",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate a key
	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	address, err := km.GenerateKey("test", permissions)
	require.NoError(t, err)

	// Get key info
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "test", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)
	// EncryptedKey should be nil in the returned info for security
	assert.Nil(t, keyInfo.EncryptedKey)

	// Test getting non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	_, err = km.GetKeyInfo(nonExistentAddr)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")
}

// TestEncryptDecryptPrivateKey tests the encryption/decryption functionality
func TestEncryptDecryptPrivateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_encrypt",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)

	// Test encryption
	encryptedKey, err := km.encryptPrivateKey(privateKey)
	require.NoError(t, err)
	assert.NotNil(t, encryptedKey)
	assert.NotEmpty(t, encryptedKey)

	// Test decryption
	decryptedKey, err := km.decryptPrivateKey(encryptedKey)
	require.NoError(t, err)
	assert.NotNil(t, decryptedKey)

	// Verify the keys are the same
	assert.Equal(t, crypto.PubkeyToAddress(privateKey.PublicKey), crypto.PubkeyToAddress(decryptedKey.PublicKey))
	assert.Equal(t, crypto.FromECDSA(privateKey), crypto.FromECDSA(decryptedKey))

	// Test decryption with invalid data
	_, err = km.decryptPrivateKey([]byte("invalid_encrypted_data"))
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encrypted key too short")
}

// TestRotateKey tests key rotation functionality
func TestRotateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_rotate",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate an original key
	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	originalAddr, err := km.GenerateKey("test", permissions)
	require.NoError(t, err)

	// Rotate the key
	newAddr, err := km.RotateKey(originalAddr)
	require.NoError(t, err)
	assert.NotEqual(t, originalAddr, newAddr)

	// Check that the original key still exists but has restricted permissions
	originalInfo, err := km.GetKeyInfo(originalAddr)
	require.NoError(t, err)
	assert.False(t, originalInfo.Permissions.CanSign)
	assert.False(t, originalInfo.Permissions.CanTransfer)

	// Check that the new key has the same permissions
	newInfo, err := km.GetKeyInfo(newAddr)
	require.NoError(t, err)
	assert.Equal(t, permissions, newInfo.Permissions)
	assert.True(t, newInfo.Permissions.CanSign)
	assert.True(t, newInfo.Permissions.CanTransfer)

	// Test rotating non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	_, err = km.RotateKey(nonExistentAddr)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")
}

// TestSignTransaction tests transaction signing with various scenarios
func TestSignTransaction(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_sign",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate a key with signing permissions
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(10000000000000000000), // 10 ETH
	}
	signerAddr, err := km.GenerateKey("signer", permissions)
	require.NoError(t, err)

	// Create a test transaction
	chainID := big.NewInt(1)
	tx := types.NewTransaction(0, common.Address{}, big.NewInt(1000000000000000000), 21000, big.NewInt(20000000000), nil)

	// Create signing request
	request := &SigningRequest{
		Transaction:  tx,
		ChainID:      chainID,
		From:         signerAddr,
		Purpose:      "Test transaction",
		UrgencyLevel: 3,
	}

	// Sign the transaction
	result, err := km.SignTransaction(request)
	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.NotNil(t, result.SignedTx)
	assert.NotNil(t, result.Signature)
	assert.NotEmpty(t, result.AuditID)
	assert.WithinDuration(t, time.Now(), result.SignedAt, time.Second)
	assert.Equal(t, signerAddr, result.KeyUsed)

	// Verify the signature is valid
	signedTx := result.SignedTx
	from, err := types.Sender(types.NewEIP155Signer(chainID), signedTx)
	require.NoError(t, err)
	assert.Equal(t, signerAddr, from)

	// Test signing with non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	request.From = nonExistentAddr
	_, err = km.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")

	// Test signing with key that can't sign
	km2, err := NewKeyManager(config, log)
	require.NoError(t, err)

	noSignPermissions := KeyPermissions{
		CanSign:        false,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(10000000000000000000),
	}
	noSignAddr, err := km2.GenerateKey("no_sign", noSignPermissions)
	require.NoError(t, err)

	request.From = noSignAddr
	_, err = km2.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not permitted to sign")
}

// TestSignTransactionTransferLimits tests transfer limits during signing
func TestSignTransactionTransferLimits(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_limits",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(t, err)

	// Generate a key with limited transfer permissions
	maxTransfer := big.NewInt(1000000000000000000) // 1 ETH
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: maxTransfer,
	}
	signerAddr, err := km.GenerateKey("limited_signer", permissions)
	require.NoError(t, err)

	// Create a transaction that exceeds the limit
	chainID := big.NewInt(1)
	excessiveTx := types.NewTransaction(0, common.Address{}, big.NewInt(2000000000000000000), 21000, big.NewInt(20000000000), nil) // 2 ETH

	request := &SigningRequest{
		Transaction:  excessiveTx,
		ChainID:      chainID,
		From:         signerAddr,
		Purpose:      "Excessive transfer",
		UrgencyLevel: 3,
	}

	_, err = km.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "transfer amount exceeds limit")
}

// TestDeriveEncryptionKey tests the key derivation function
func TestDeriveEncryptionKey(t *testing.T) {
	// Test with valid master key
	masterKey := "test_encryption_key_very_long_and_secure_for_testing"
	key, err := deriveEncryptionKey(masterKey)
	require.NoError(t, err)
	assert.NotNil(t, key)
	assert.Len(t, key, 32) // Should be 32 bytes for AES-256

	// Test with different master key (should produce different result)
	differentKey := "different_test_encryption_key_very_long_and_secure_for_testing"
	key2, err := deriveEncryptionKey(differentKey)
	require.NoError(t, err)
	assert.NotEqual(t, key, key2)

	// Test with empty master key
	_, err = deriveEncryptionKey("")
	assert.Error(t, err)
}

// TestValidateConfig tests the configuration validation function
func TestValidateConfig(t *testing.T) {
	// Test with valid config
	validConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}
	err := validateConfig(validConfig)
	assert.NoError(t, err)

	// Test with empty encryption key
	emptyKeyConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "",
	}
	err = validateConfig(emptyKeyConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encryption key cannot be empty")

	// Test with short encryption key
	shortKeyConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "short",
	}
	err = validateConfig(shortKeyConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encryption key must be at least 32 characters")

	// Test with empty keystore path
	emptyPathConfig := &KeyManagerConfig{
		KeystorePath:  "",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}
	err = validateConfig(emptyPathConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "keystore path cannot be empty")
}

// TestClearPrivateKey tests the private key clearing function
func TestClearPrivateKey(t *testing.T) {
	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)

	// Store original D value for comparison
	originalD := new(big.Int).Set(privateKey.D)

	// Clear the private key
	clearPrivateKey(privateKey)

	// Verify the D value has been cleared
	assert.Zero(t, privateKey.D.Sign())
	assert.NotEqual(t, originalD, privateKey.D)

	// Test with nil private key (should not panic)
	clearPrivateKey(nil)
}

// TestGenerateAuditID tests the audit ID generation function
func TestGenerateAuditID(t *testing.T) {
	id1 := generateAuditID()
	id2 := generateAuditID()

	// Both should be non-empty and different
	assert.NotEmpty(t, id1)
	assert.NotEmpty(t, id2)
	assert.NotEqual(t, id1, id2)

	// Should be a valid hex string
	_, err := common.HexToHash(id1)
	assert.NoError(t, err)

	_, err = common.HexToHash(id2)
	assert.NoError(t, err)
}

// TestCalculateRiskScore tests the risk score calculation function
func TestCalculateRiskScore(t *testing.T) {
	// Test failed operations (high risk)
	score := calculateRiskScore("TRANSACTION_SIGNED", false)
	assert.Equal(t, 8, score)

	// Test successful transaction signing (low-medium risk)
	score = calculateRiskScore("TRANSACTION_SIGNED", true)
	assert.Equal(t, 3, score)

	// Test key generation (medium risk)
	score = calculateRiskScore("KEY_GENERATED", true)
	assert.Equal(t, 5, score)

	// Test key import (medium risk)
	score = calculateRiskScore("KEY_IMPORTED", true)
	assert.Equal(t, 5, score)

	// Test key rotation (medium risk)
	score = calculateRiskScore("KEY_ROTATED", true)
	assert.Equal(t, 4, score)

	// Test default (low risk)
	score = calculateRiskScore("UNKNOWN_OPERATION", true)
	assert.Equal(t, 2, score)
}

// TestKeyPermissions tests the KeyPermissions struct
func TestKeyPermissions(t *testing.T) {
	// Test creating permissions with max transfer limit
	maxTransfer := big.NewInt(1000000000000000000) // 1 ETH
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: maxTransfer,
		AllowedContracts: []string{
			"0x1234567890123456789012345678901234567890",
			"0x0987654321098765432109876543210987654321",
		},
		RequireConfirm: true,
	}

	assert.True(t, permissions.CanSign)
	assert.True(t, permissions.CanTransfer)
	assert.Equal(t, maxTransfer, permissions.MaxTransferWei)
	assert.Len(t, permissions.AllowedContracts, 2)
	assert.True(t, permissions.RequireConfirm)
}

// BenchmarkKeyGeneration benchmarks key generation performance
func BenchmarkKeyGeneration(b *testing.B) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/benchmark_keystore",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(b, err)

	permissions := KeyPermissions{CanSign: true}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := km.GenerateKey("benchmark", permissions)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkTransactionSigning benchmarks transaction signing performance
func BenchmarkTransactionSigning(b *testing.B) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/benchmark_signing",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := NewKeyManager(config, log)
	require.NoError(b, err)

	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	signerAddr, err := km.GenerateKey("benchmark_signer", permissions)
	require.NoError(b, err)

	chainID := big.NewInt(1)
	tx := types.NewTransaction(0, common.Address{}, big.NewInt(1000000000000000000), 21000, big.NewInt(20000000000), nil)

	request := &SigningRequest{
		Transaction: tx,
		ChainID:     chainID,
		From:        signerAddr,
		Purpose:     "Benchmark transaction",
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := km.SignTransaction(request)
		if err != nil {
			b.Fatal(err)
		}
	}
}