diff --git a/pkg/cryptography/aes.go b/pkg/cryptography/aes.go index 10297a2..97b08f3 100644 --- a/pkg/cryptography/aes.go +++ b/pkg/cryptography/aes.go @@ -8,7 +8,57 @@ import ( "io" ) +const ( + // AES key sizes in bytes + AES128KeySize = 16 // 128 bits + AES192KeySize = 24 // 192 bits + AES256KeySize = 32 // 256 bits + + // Default to AES-256 + DefaultKeySize = AES256KeySize +) + +// GenerateAESKey generates a random AES key of the specified size +func GenerateAESKey(keySize int) ([]byte, error) { + if keySize != AES128KeySize && keySize != AES192KeySize && keySize != AES256KeySize { + return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes") + } + + key := make([]byte, keySize) + if _, err := io.ReadFull(rand.Reader, key); err != nil { + return nil, err + } + return key, nil +} + +// GenerateAES256Key generates a random AES-256 key (default) +func GenerateAES256Key() ([]byte, error) { + return GenerateAESKey(AES256KeySize) +} + +// EncryptAES256CBC encrypts data using AES-256 in CBC mode +func EncryptAES256CBC(key, plaintext []byte) ([]byte, error) { + if len(key) != AES256KeySize { + return nil, errors.New("key must be 32 bytes for AES-256") + } + return EncryptAESCBC(key, plaintext) +} + +// DecryptAES256CBC decrypts data using AES-256 in CBC mode +func DecryptAES256CBC(key, ciphertext []byte) ([]byte, error) { + if len(key) != AES256KeySize { + return nil, errors.New("key must be 32 bytes for AES-256") + } + return DecryptAESCBC(key, ciphertext) +} + +// EncryptAESCBC encrypts data using AES in CBC mode (accepts any valid AES key size) func EncryptAESCBC(key, plaintext []byte) ([]byte, error) { + // Validate key size + if len(key) != AES128KeySize && len(key) != AES192KeySize && len(key) != AES256KeySize { + return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes") + } + block, err := aes.NewCipher(key) if err != nil { return nil, err @@ -36,7 +86,13 @@ func EncryptAESCBC(key, plaintext []byte) ([]byte, error) { return append(iv, ciphertext...), nil } +// DecryptAESCBC decrypts data using AES in CBC mode (accepts any valid AES key size) func DecryptAESCBC(key, ciphertext []byte) ([]byte, error) { + // Validate key size + if len(key) != AES128KeySize && len(key) != AES192KeySize && len(key) != AES256KeySize { + return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes") + } + block, err := aes.NewCipher(key) if err != nil { return nil, err @@ -58,10 +114,22 @@ func DecryptAESCBC(key, ciphertext []byte) ([]byte, error) { mode.CryptBlocks(plaintext, ciphertext) // Remove PKCS7 padding + if len(plaintext) == 0 { + return nil, errors.New("invalid padding: empty plaintext") + } + padding := int(plaintext[len(plaintext)-1]) - if padding == 0 || padding > len(plaintext) { + if padding == 0 || padding > aes.BlockSize || padding > len(plaintext) { return nil, errors.New("invalid PKCS7 padding") } - // TODO: Add check to ensure all padding bytes are correct? + + // Verify all padding bytes are correct + for i := len(plaintext) - padding; i < len(plaintext); i++ { + if plaintext[i] != byte(padding) { + return nil, errors.New("invalid PKCS7 padding") + } + } + return plaintext[:len(plaintext)-padding], nil } + diff --git a/pkg/cryptography/aes_test.go b/pkg/cryptography/aes_test.go index 1656899..a685b12 100644 --- a/pkg/cryptography/aes_test.go +++ b/pkg/cryptography/aes_test.go @@ -4,14 +4,49 @@ import ( "bytes" "crypto/aes" "crypto/rand" + "fmt" "testing" ) -func TestAESCBCEncryptionDecryption(t *testing.T) { - // Generate a random key (AES-256) - key := make([]byte, 32) - if _, err := rand.Read(key); err != nil { - t.Fatalf("Failed to generate random key: %v", err) +func TestGenerateAESKeys(t *testing.T) { + t.Run("GenerateAES256Key", func(t *testing.T) { + key, err := GenerateAES256Key() + if err != nil { + t.Fatalf("GenerateAES256Key failed: %v", err) + } + if len(key) != AES256KeySize { + t.Errorf("Expected key size %d, got %d", AES256KeySize, len(key)) + } + }) + + t.Run("GenerateAESKey_AllSizes", func(t *testing.T) { + sizes := []int{AES128KeySize, AES192KeySize, AES256KeySize} + for _, size := range sizes { + key, err := GenerateAESKey(size) + if err != nil { + t.Fatalf("GenerateAESKey(%d) failed: %v", size, err) + } + if len(key) != size { + t.Errorf("Expected key size %d, got %d", size, len(key)) + } + } + }) + + t.Run("GenerateAESKey_InvalidSize", func(t *testing.T) { + invalidSizes := []int{8, 15, 17, 23, 25, 31, 33, 64} + for _, size := range invalidSizes { + _, err := GenerateAESKey(size) + if err == nil { + t.Errorf("GenerateAESKey(%d) should have failed but didn't", size) + } + } + }) +} + +func TestAES256CBCEncryptionDecryption(t *testing.T) { + key, err := GenerateAES256Key() + if err != nil { + t.Fatalf("Failed to generate AES-256 key: %v", err) } testCases := []struct { @@ -20,67 +55,208 @@ func TestAESCBCEncryptionDecryption(t *testing.T) { }{ {"ShortMessage", []byte("Hello")}, {"BlockSizeMessage", []byte("This is 16 bytes")}, - {"LongMessage", []byte("This is a longer message that spans multiple AES blocks.")}, + {"LongMessage", []byte("This is a longer message that spans multiple AES blocks and tests the padding.")}, {"EmptyMessage", []byte("")}, + {"SingleByte", []byte("A")}, + {"ExactlyTwoBlocks", []byte("This is exactly 32 bytes long!!!")}, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - ciphertext, err := EncryptAESCBC(key, tc.plaintext) + ciphertext, err := EncryptAES256CBC(key, tc.plaintext) if err != nil { - t.Fatalf("EncryptAESCBC failed: %v", err) + t.Fatalf("EncryptAES256CBC failed: %v", err) } - decrypted, err := DecryptAESCBC(key, ciphertext) + decrypted, err := DecryptAES256CBC(key, ciphertext) if err != nil { - t.Fatalf("DecryptAESCBC failed: %v", err) + t.Fatalf("DecryptAES256CBC failed: %v", err) } if !bytes.Equal(tc.plaintext, decrypted) { - t.Errorf("Decrypted text does not match original plaintext. Got %q, want %q", decrypted, tc.plaintext) + t.Errorf("Decrypted text does not match original plaintext.\nGot: %q (%x)\nWant: %q (%x)", + decrypted, decrypted, tc.plaintext, tc.plaintext) + } + }) + } +} + +func TestAES256CBC_InvalidKeySize(t *testing.T) { + plaintext := []byte("test message") + + invalidKeys := [][]byte{ + make([]byte, 16), // AES-128 + make([]byte, 24), // AES-192 + make([]byte, 15), // Too short + make([]byte, 33), // Too long + nil, // Nil key + } + + for i, key := range invalidKeys { + t.Run(fmt.Sprintf("InvalidKey_%d", i), func(t *testing.T) { + _, err := EncryptAES256CBC(key, plaintext) + if err == nil { + t.Error("EncryptAES256CBC should have failed with invalid key size") + } + + // Test with some dummy ciphertext + dummyCiphertext := make([]byte, 32) // Just enough for IV + one block + rand.Read(dummyCiphertext) + _, err = DecryptAES256CBC(key, dummyCiphertext) + if err == nil { + t.Error("DecryptAES256CBC should have failed with invalid key size") + } + }) + } +} + +func TestAESCBCEncryptionDecryption(t *testing.T) { + keySizes := []int{AES128KeySize, AES192KeySize, AES256KeySize} + + for _, keySize := range keySizes { + t.Run(fmt.Sprintf("AES_%d", keySize*8), func(t *testing.T) { + key, err := GenerateAESKey(keySize) + if err != nil { + t.Fatalf("Failed to generate AES-%d key: %v", keySize*8, err) + } + + testCases := []struct { + name string + plaintext []byte + }{ + {"ShortMessage", []byte("Hello")}, + {"BlockSizeMessage", []byte("This is 16 bytes")}, + {"LongMessage", []byte("This is a longer message that spans multiple AES blocks.")}, + {"EmptyMessage", []byte("")}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ciphertext, err := EncryptAESCBC(key, tc.plaintext) + if err != nil { + t.Fatalf("EncryptAESCBC failed: %v", err) + } + + decrypted, err := DecryptAESCBC(key, ciphertext) + if err != nil { + t.Fatalf("DecryptAESCBC failed: %v", err) + } + + if !bytes.Equal(tc.plaintext, decrypted) { + t.Errorf("Decrypted text does not match original plaintext. Got %q, want %q", decrypted, tc.plaintext) + } + }) } }) } } func TestDecryptAESCBCErrorCases(t *testing.T) { - key := make([]byte, 32) - _, _ = rand.Read(key) + key, err := GenerateAES256Key() + if err != nil { + t.Fatalf("Failed to generate key: %v", err) + } t.Run("CiphertextTooShort", func(t *testing.T) { shortCiphertext := []byte{0x01, 0x02, 0x03} // Less than AES block size _, err := DecryptAESCBC(key, shortCiphertext) if err == nil { - t.Error("DecryptAESCBC should have failed for ciphertext shorter than block size, but it didn't") + t.Error("DecryptAESCBC should have failed for ciphertext shorter than block size") } }) - t.Run("InvalidPadding", func(t *testing.T) { - // Encrypt something valid first - plaintext := []byte("valid data") - ciphertext, _ := EncryptAESCBC(key, plaintext) - - // Tamper with the ciphertext (specifically the part that would affect padding) - if len(ciphertext) > aes.BlockSize { - ciphertext[len(ciphertext)-1] = ^ciphertext[len(ciphertext)-1] // Flip bits of last byte - } - - _, err := DecryptAESCBC(key, ciphertext) + t.Run("InvalidKeySize", func(t *testing.T) { + invalidKey := make([]byte, 17) // Invalid key size + validCiphertext := make([]byte, 32) // IV + one block + rand.Read(validCiphertext) + + _, err := DecryptAESCBC(invalidKey, validCiphertext) if err == nil { - // Note: Depending on the padding implementation and the nature of the tampering, - // CBC decryption might not always error out on bad padding. It might return garbage data. - // A more robust test might check the decrypted content, but error checking is a start. - t.Logf("DecryptAESCBC did not error on potentially invalid padding (this might be expected)") + t.Error("DecryptAESCBC should have failed for invalid key size") } }) t.Run("CiphertextNotMultipleOfBlockSize", func(t *testing.T) { iv := make([]byte, aes.BlockSize) - _, _ = rand.Read(iv) + rand.Read(iv) invalidCiphertext := append(iv, []byte{0x01, 0x02, 0x03}...) // IV + data not multiple of block size _, err := DecryptAESCBC(key, invalidCiphertext) if err == nil { - t.Error("DecryptAESCBC should have failed for ciphertext not multiple of block size, but it didn't") + t.Error("DecryptAESCBC should have failed for ciphertext not multiple of block size") + } + }) + + t.Run("InvalidPadding", func(t *testing.T) { + // Create a valid ciphertext first + plaintext := []byte("valid data") + ciphertext, err := EncryptAESCBC(key, plaintext) + if err != nil { + t.Fatalf("Failed to create test ciphertext: %v", err) + } + + // Corrupt the last byte (which affects padding) + corruptedCiphertext := make([]byte, len(ciphertext)) + copy(corruptedCiphertext, ciphertext) + corruptedCiphertext[len(corruptedCiphertext)-1] ^= 0xFF + + _, err = DecryptAESCBC(key, corruptedCiphertext) + if err == nil { + t.Error("DecryptAESCBC should have failed for corrupted padding") + } + }) + + t.Run("EmptyPlaintext", func(t *testing.T) { + // Create a ciphertext that would result in empty plaintext + invalidCiphertext := make([]byte, aes.BlockSize) // Only IV, no data + _, err := DecryptAESCBC(key, invalidCiphertext) + if err == nil { + t.Error("DecryptAESCBC should have failed for empty ciphertext data") + } + }) +} + +func TestConstants(t *testing.T) { + if AES128KeySize != 16 { + t.Errorf("AES128KeySize should be 16, got %d", AES128KeySize) + } + if AES192KeySize != 24 { + t.Errorf("AES192KeySize should be 24, got %d", AES192KeySize) + } + if AES256KeySize != 32 { + t.Errorf("AES256KeySize should be 32, got %d", AES256KeySize) + } + if DefaultKeySize != AES256KeySize { + t.Errorf("DefaultKeySize should be AES256KeySize (%d), got %d", AES256KeySize, DefaultKeySize) + } +} + +func BenchmarkAES256CBC(b *testing.B) { + key, err := GenerateAES256Key() + if err != nil { + b.Fatalf("Failed to generate key: %v", err) + } + + data := make([]byte, 1024) // 1KB of data + rand.Read(data) + + b.Run("Encrypt", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := EncryptAES256CBC(key, data) + if err != nil { + b.Fatal(err) + } + } + }) + + ciphertext, _ := EncryptAES256CBC(key, data) + b.Run("Decrypt", func(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := DecryptAES256CBC(key, ciphertext) + if err != nil { + b.Fatal(err) + } } }) }