teleport/pkg/encryption/encryption_test.go

package encryption

import (
	"bytes"
	"encoding/hex"
	"testing"
	"time"
)

func TestDeriveKey(t *testing.T) {
	password := "test-password"
	key := DeriveKey(password)

	if len(key) != 32 {
		t.Errorf("Expected key length 32, got %d", len(key))
	}

	// Test that same password produces same key
	key2 := DeriveKey(password)
	if string(key) != string(key2) {
		t.Error("Same password should produce same key")
	}

	// Test that different passwords produce different keys
	key3 := DeriveKey("different-password")
	if string(key) == string(key3) {
		t.Error("Different passwords should produce different keys")
	}
}

func TestEncryptDecrypt(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("Hello, World! This is a test message.")

	// Test encryption
	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	if len(encryptedData) <= len(originalData) {
		t.Error("Encrypted data should be longer than original data (due to nonce)")
	}

	// Test decryption
	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption failed: %v", err)
	}

	if string(decryptedData) != string(originalData) {
		t.Errorf("Decrypted data doesn't match original. Expected: %s, Got: %s",
			string(originalData), string(decryptedData))
	}
}

func TestEncryptDecryptEmptyData(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("")

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption of empty data failed: %v", err)
	}

	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption of empty data failed: %v", err)
	}

	if len(decryptedData) != 0 {
		t.Error("Decrypted empty data should be empty")
	}
}

func TestEncryptDecryptLargeData(t *testing.T) {
	key := DeriveKey("test-key")

	// Create a large data block (1MB)
	originalData := make([]byte, 1024*1024)
	for i := range originalData {
		originalData[i] = byte(i % 256)
	}

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption of large data failed: %v", err)
	}

	decryptedData, err := DecryptData(encryptedData, key)
	if err != nil {
		t.Fatalf("Decryption of large data failed: %v", err)
	}

	if len(decryptedData) != len(originalData) {
		t.Errorf("Decrypted data length mismatch. Expected: %d, Got: %d",
			len(originalData), len(decryptedData))
	}

	for i := range originalData {
		if decryptedData[i] != originalData[i] {
			t.Errorf("Data mismatch at position %d", i)
			break
		}
	}
}

func TestWrongKeyDecryption(t *testing.T) {
	key1 := DeriveKey("key1")
	key2 := DeriveKey("key2")
	originalData := []byte("test data")

	encryptedData, err := EncryptData(originalData, key1)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	// Try to decrypt with wrong key
	_, err = DecryptData(encryptedData, key2)
	if err == nil {
		t.Error("Decryption with wrong key should fail")
	}
}

func TestCorruptedDataDecryption(t *testing.T) {
	key := DeriveKey("test-key")
	originalData := []byte("test data")

	encryptedData, err := EncryptData(originalData, key)
	if err != nil {
		t.Fatalf("Encryption failed: %v", err)
	}

	// Corrupt the data
	encryptedData[0] ^= 0xFF

	// Try to decrypt corrupted data
	_, err = DecryptData(encryptedData, key)
	if err == nil {
		t.Error("Decryption of corrupted data should fail")
	}
}

func TestValidateEncryptionKey(t *testing.T) {
	// Test valid key
	validKey := "a0e3dd20a761b118ca234160dd8b87230a001e332a97c9cfe3b8b9c99efaae03"
	if err := ValidateEncryptionKey(validKey); err != nil {
		t.Errorf("Valid key should pass validation: %v", err)
	}

	// Test short key
	shortKey := "short"
	if err := ValidateEncryptionKey(shortKey); err == nil {
		t.Error("Short key should fail validation")
	}

	// Test weak keys
	weakKeys := []string{
		"password", "123456", "admin", "test", "default",
		"your-secure-encryption-key-change-this-to-something-random",
		"test-encryption-key-12345", "teleport-key", "secret",
	}

	for _, weakKey := range weakKeys {
		if err := ValidateEncryptionKey(weakKey); err == nil {
			t.Errorf("Weak key '%s' should fail validation", weakKey)
		}
	}

	// Test low entropy key
	lowEntropyKey := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
	if err := ValidateEncryptionKey(lowEntropyKey); err == nil {
		t.Error("Low entropy key should fail validation")
	}
}

func TestReplayProtection(t *testing.T) {
	rp := NewReplayProtection()
	nonce := uint64(12345)
	timestamp := time.Now().Unix()

	// First use should be valid
	if !rp.IsValidNonce(nonce, timestamp) {
		t.Error("First use of nonce should be valid")
	}

	// Replay should be invalid
	if rp.IsValidNonce(nonce, timestamp) {
		t.Error("Replay of nonce should be invalid")
	}

	// Different nonce should be valid
	if !rp.IsValidNonce(nonce+1, timestamp) {
		t.Error("Different nonce should be valid")
	}
}

func TestValidatePacketTimestamp(t *testing.T) {
	now := time.Now().Unix()

	// Current timestamp should be valid
	if !ValidatePacketTimestamp(now) {
		t.Error("Current timestamp should be valid")
	}

	// Recent timestamp should be valid
	recent := now - 60 // 1 minute ago
	if !ValidatePacketTimestamp(recent) {
		t.Error("Recent timestamp should be valid")
	}

	// Old timestamp should be invalid
	old := now - int64(MaxPacketAge.Seconds()) - 1
	if ValidatePacketTimestamp(old) {
		t.Error("Old timestamp should be invalid")
	}

	// Future timestamp should be invalid
	future := now + 3600 // 1 hour in future
	if ValidatePacketTimestamp(future) {
		t.Error("Future timestamp should be invalid")
	}
}

func TestConstantTimeCompare(t *testing.T) {
	a := []byte("test")
	b := []byte("test")
	c := []byte("different")

	if !ConstantTimeCompare(a, b) {
		t.Error("Identical byte slices should compare equal")
	}

	if ConstantTimeCompare(a, c) {
		t.Error("Different byte slices should not compare equal")
	}

	// Test with empty slices
	empty1 := []byte{}
	empty2 := []byte{}
	if !ConstantTimeCompare(empty1, empty2) {
		t.Error("Empty slices should compare equal")
	}
}

func TestResolveKeyRaw(t *testing.T) {
	raw, err := GenerateRawKey()
	if err != nil {
		t.Fatalf("GenerateRawKey: %v", err)
	}
	if !IsRawKey(raw) {
		t.Fatalf("generated key is not raw: %q", raw[:4])
	}
	key, err := ResolveKey(raw)
	if err != nil {
		t.Fatalf("ResolveKey raw: %v", err)
	}
	if len(key) != 32 {
		t.Fatalf("raw key length %d", len(key))
	}
	decoded, err := hex.DecodeString(raw[len(RawKeyPrefix):])
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(key, decoded) {
		t.Fatal("raw key was not hex-decoded as-is")
	}
	if bytes.Equal(key, DeriveKey(raw)) {
		t.Fatal("raw key must not go through PBKDF2")
	}
}

func TestResolveKeyLegacyHexStillPBKDF2(t *testing.T) {
	legacy := "a0e3dd20a761b118ca234160dd8b87230a001e332a97c9cfe3b8b9c99efaae03"
	decoded, err := hex.DecodeString(legacy)
	if err != nil {
		t.Fatal(err)
	}
	got, err := ResolveKey(legacy)
	if err != nil {
		t.Fatalf("ResolveKey legacy hex: %v", err)
	}
	want := DeriveKey(legacy)
	if !bytes.Equal(got, want) {
		t.Fatal("unprefixed 64-hex must still use PBKDF2 (old --generate-key configs)")
	}
	if bytes.Equal(got, decoded) {
		t.Fatal("unprefixed 64-hex must not be treated as a raw AES key")
	}
}

func TestResolveKeyPassphrase(t *testing.T) {
	pw := "test-passphrase-not-a-hex-key-value"
	got, err := ResolveKey(pw)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, DeriveKey(pw)) {
		t.Fatal("passphrase should use PBKDF2")
	}
}

func TestDeriveKeyCached(t *testing.T) {
	pw := "cache-me-please-this-is-long-enough"
	start := time.Now()
	k1 := DeriveKey(pw)
	first := time.Since(start)
	start = time.Now()
	k2 := DeriveKey(pw)
	second := time.Since(start)
	if !bytes.Equal(k1, k2) {
		t.Fatal("cached key mismatch")
	}
	if first < 10*time.Millisecond {
		t.Logf("first PBKDF2 unexpectedly fast: %v", first)
	}
	if second > 5*time.Millisecond {
		t.Fatalf("cached DeriveKey too slow: first=%v second=%v", first, second)
	}
}

func TestValidateRawEncryptionKey(t *testing.T) {
	raw, err := GenerateRawKey()
	if err != nil {
		t.Fatal(err)
	}
	if err := ValidateEncryptionKey(raw); err != nil {
		t.Fatalf("valid raw key rejected: %v", err)
	}
	if err := ValidateEncryptionKey("raw:not-hex"); err == nil {
		t.Fatal("invalid raw hex should fail")
	}
	if err := ValidateEncryptionKey("raw:abcd"); err == nil {
		t.Fatal("short raw key should fail")
	}
}