# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 AES Encryption Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""AES Encryption Module.

This module provides a comprehensive AES encryption implementation with support
for multiple modes of operation, key derivation, and secure random generation.
Follows Google Python Style Guide and security best practices.
"""

import base64
import hashlib
import os
import secrets
from typing import Optional, Tuple, Union

from cryptography.hazmat.primitives import hashes, padding, serialization
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend


class AESEncryptionError(Exception):
    """Base exception for AES encryption operations."""
    pass


class InvalidKeyError(AESEncryptionError):
    """Raised when an invalid encryption key is provided."""
    pass


class DecryptionError(AESEncryptionError):
    """Raised when decryption fails."""
    pass


class AESEncryption:
    """AES Encryption utility class.
    
    Provides secure AES encryption and decryption with multiple modes of operation.
    Supports key derivation from passwords and secure random key generation.
    
    Attributes:
        DEFAULT_KEY_SIZE: Default AES key size in bytes (256-bit).
        DEFAULT_BLOCK_SIZE: AES block size in bytes.
        DEFAULT_SALT_SIZE: Default salt size for key derivation.
        DEFAULT_IV_SIZE: Default initialization vector size.
    """
    
    DEFAULT_KEY_SIZE = 32  # 256-bit key
    DEFAULT_BLOCK_SIZE = 16  # 128-bit blocks
    DEFAULT_SALT_SIZE = 16  # 128-bit salt
    DEFAULT_IV_SIZE = 16   # 128-bit IV
    
    def __init__(self, key: Optional[bytes] = None, mode: str = 'CBC'):
        """Initialize AES encryption instance.
        
        Args:
            key: Optional encryption key. If None, must be set later.
            mode: Encryption mode ('CBC', 'GCM', 'CTR'). Defaults to 'CBC'.
            
        Raises:
            InvalidKeyError: If provided key is invalid.
            ValueError: If mode is not supported.
        """
        self._key = None
        self._mode = mode.upper()
        
        if self._mode not in ['CBC', 'GCM', 'CTR']:
            raise ValueError(f"Unsupported mode: {mode}")
            
        if key is not None:
            self.set_key(key)
    
    def set_key(self, key: bytes) -> None:
        """Set the encryption key.
        
        Args:
            key: Encryption key (16, 24, or 32 bytes for AES-128/192/256).
            
        Raises:
            InvalidKeyError: If key length is invalid.
        """
        if not isinstance(key, bytes):
            raise InvalidKeyError("Key must be bytes")
            
        if len(key) not in [16, 24, 32]:
            raise InvalidKeyError(
                f"Invalid key length: {len(key)}. Must be 16, 24, or 32 bytes"
            )
            
        self._key = key
    
    def generate_key(self, key_size: int = DEFAULT_KEY_SIZE) -> bytes:
        """Generate a cryptographically secure random key.
        
        Args:
            key_size: Key size in bytes (16, 24, or 32).
            
        Returns:
            Randomly generated key.
            
        Raises:
            ValueError: If key_size is invalid.
        """
        if key_size not in [16, 24, 32]:
            raise ValueError(f"Invalid key size: {key_size}")
            
        key = secrets.token_bytes(key_size)
        self.set_key(key)
        return key
    
    def derive_key_from_password(
        self,
        password: str,
        salt: Optional[bytes] = None,
        iterations: int = 100000,
        key_size: int = DEFAULT_KEY_SIZE
    ) -> Tuple[bytes, bytes]:
        """Derive encryption key from password using PBKDF2.
        
        Args:
            password: Password string.
            salt: Optional salt. If None, generates random salt.
            iterations: Number of PBKDF2 iterations.
            key_size: Desired key size in bytes.
            
        Returns:
            Tuple of (derived_key, salt).
        """
        if salt is None:
            salt = secrets.token_bytes(self.DEFAULT_SALT_SIZE)
            
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=key_size,
            salt=salt,
            iterations=iterations,
            backend=default_backend()
        )
        
        key = kdf.derive(password.encode('utf-8'))
        self.set_key(key)
        return key, salt
    
    def _get_cipher(self, iv: bytes) -> Cipher:
        """Create cipher instance based on mode.
        
        Args:
            iv: Initialization vector.
            
        Returns:
            Configured cipher instance.
        """
        if self._key is None:
            raise InvalidKeyError("No key set")
            
        algorithm = algorithms.AES(self._key)
        
        if self._mode == 'CBC':
            mode = modes.CBC(iv)
        elif self._mode == 'GCM':
            mode = modes.GCM(iv)
        elif self._mode == 'CTR':
            mode = modes.CTR(iv)
        else:
            raise ValueError(f"Unsupported mode: {self._mode}")
            
        return Cipher(algorithm, mode, backend=default_backend())
    
    def encrypt(
        self,
        plaintext: Union[str, bytes],
        iv: Optional[bytes] = None
    ) -> Tuple[bytes, bytes]:
        """Encrypt plaintext data.
        
        Args:
            plaintext: Data to encrypt (string or bytes).
            iv: Optional initialization vector. If None, generates random IV.
            
        Returns:
            Tuple of (ciphertext, iv). For GCM mode, authentication tag is appended to ciphertext.
            
        Raises:
            InvalidKeyError: If no key is set.
        """
        if isinstance(plaintext, str):
            plaintext = plaintext.encode('utf-8')
            
        if iv is None:
            iv = secrets.token_bytes(self.DEFAULT_IV_SIZE)
            
        cipher = self._get_cipher(iv)
        encryptor = cipher.encryptor()
        
        if self._mode == 'CBC':
            # Apply PKCS7 padding for CBC mode
            padder = padding.PKCS7(128).padder()
            padded_data = padder.update(plaintext)
            padded_data += padder.finalize()
            ciphertext = encryptor.update(padded_data) + encryptor.finalize()
        elif self._mode == 'GCM':
            # GCM mode - append authentication tag to ciphertext
            ciphertext = encryptor.update(plaintext) + encryptor.finalize()
            ciphertext += encryptor.tag  # Append 16-byte authentication tag
        else:
            # CTR mode doesn't require padding
            ciphertext = encryptor.update(plaintext) + encryptor.finalize()
            
        return ciphertext, iv
    
    def decrypt(self, ciphertext: bytes, iv: bytes) -> bytes:
        """Decrypt ciphertext data.
        
        Args:
            ciphertext: Encrypted data. For GCM mode, should include authentication tag.
            iv: Initialization vector used for encryption.
            
        Returns:
            Decrypted plaintext.
            
        Raises:
            InvalidKeyError: If no key is set.
            DecryptionError: If decryption fails.
        """
        try:
            if self._mode == 'GCM':
                # For GCM mode, extract authentication tag from end of ciphertext
                if len(ciphertext) < 16:
                    raise DecryptionError("GCM ciphertext too short (missing authentication tag)")
                auth_tag = ciphertext[-16:]  # Last 16 bytes are the authentication tag
                actual_ciphertext = ciphertext[:-16]  # Everything except the tag
                
                # Create cipher with authentication tag
                algorithm = algorithms.AES(self._key)
                mode = modes.GCM(iv, auth_tag)
                cipher = Cipher(algorithm, mode, backend=default_backend())
                decryptor = cipher.decryptor()
                
                plaintext = decryptor.update(actual_ciphertext) + decryptor.finalize()
            else:
                # CBC and CTR modes
                cipher = self._get_cipher(iv)
                decryptor = cipher.decryptor()
                
                plaintext = decryptor.update(ciphertext) + decryptor.finalize()
                
                if self._mode == 'CBC':
                    # Remove PKCS7 padding for CBC mode
                    unpadder = padding.PKCS7(128).unpadder()
                    plaintext = unpadder.update(plaintext)
                    plaintext += unpadder.finalize()
                
            return plaintext
            
        except Exception as e:
            raise DecryptionError(f"Decryption failed: {str(e)}") from e
    
    def encrypt_to_base64(
        self,
        plaintext: Union[str, bytes],
        iv: Optional[bytes] = None
    ) -> str:
        """Encrypt data and return base64-encoded result.
        
        Args:
            plaintext: Data to encrypt.
            iv: Optional initialization vector.
            
        Returns:
            Base64-encoded string containing IV + ciphertext.
        """
        ciphertext, used_iv = self.encrypt(plaintext, iv)
        combined = used_iv + ciphertext
        return base64.b64encode(combined).decode('ascii')
    
    def decrypt_from_base64(self, encoded_data: str) -> bytes:
        """Decrypt base64-encoded data.
        
        Args:
            encoded_data: Base64-encoded string containing IV + ciphertext.
            
        Returns:
            Decrypted plaintext.
            
        Raises:
            DecryptionError: If decoding or decryption fails.
        """
        try:
            combined = base64.b64decode(encoded_data.encode('ascii'))
            iv = combined[:self.DEFAULT_IV_SIZE]
            ciphertext = combined[self.DEFAULT_IV_SIZE:]
            return self.decrypt(ciphertext, iv)
        except Exception as e:
            raise DecryptionError(f"Base64 decryption failed: {str(e)}") from e
    
    def get_key_info(self) -> dict:
        """Get information about the current key.
        
        Returns:
            Dictionary with key information.
        """
        if self._key is None:
            return {"key_set": False}
            
        return {
            "key_set": True,
            "key_size_bits": len(self._key) * 8,
            "key_size_bytes": len(self._key),
            "mode": self._mode
        }


def create_aes_encryptor(
    password: Optional[str] = None,
    key: Optional[bytes] = None,
    mode: str = 'CBC'
) -> AESEncryption:
    """Factory function to create AES encryptor.
    
    Args:
        password: Password for key derivation (mutually exclusive with key).
        key: Direct encryption key (mutually exclusive with password).
        mode: Encryption mode.
        
    Returns:
        Configured AESEncryption instance.
        
    Raises:
        ValueError: If both or neither password and key are provided.
    """
    if (password is None) == (key is None):
        raise ValueError("Provide either password or key, not both or neither")
        
    encryptor = AESEncryption(mode=mode)
    
    if password is not None:
        encryptor.derive_key_from_password(password)
    else:
        encryptor.set_key(key)
        
    return encryptor


# Convenience functions for simple use cases
def encrypt_string(plaintext: str, password: str, mode: str = 'CBC') -> str:
    """Encrypt a string with password-based key derivation.
    
    Args:
        plaintext: String to encrypt.
        password: Password for encryption.
        mode: Encryption mode.
        
    Returns:
        Base64-encoded encrypted string with embedded salt and IV.
    """
    encryptor = AESEncryption(mode=mode)
    key, salt = encryptor.derive_key_from_password(password)
    
    ciphertext, iv = encryptor.encrypt(plaintext)
    
    # Combine salt + IV + ciphertext for storage
    combined = salt + iv + ciphertext
    return base64.b64encode(combined).decode('ascii')


def decrypt_string(encrypted_data: str, password: str, mode: str = 'CBC') -> str:
    """Decrypt a string encrypted with encrypt_string.
    
    Args:
        encrypted_data: Base64-encoded encrypted data.
        password: Password for decryption.
        mode: Encryption mode.
        
    Returns:
        Decrypted string.
        
    Raises:
        DecryptionError: If decryption fails.
    """
    try:
        combined = base64.b64decode(encrypted_data.encode('ascii'))
        
        salt = combined[:AESEncryption.DEFAULT_SALT_SIZE]
        iv = combined[AESEncryption.DEFAULT_SALT_SIZE:
                     AESEncryption.DEFAULT_SALT_SIZE + AESEncryption.DEFAULT_IV_SIZE]
        ciphertext = combined[AESEncryption.DEFAULT_SALT_SIZE + AESEncryption.DEFAULT_IV_SIZE:]
        
        encryptor = AESEncryption(mode=mode)
        encryptor.derive_key_from_password(password, salt)
        
        plaintext = encryptor.decrypt(ciphertext, iv)
        return plaintext.decode('utf-8')
        
    except Exception as e:
        raise DecryptionError(f"String decryption failed: {str(e)}") from e


if __name__ == "__main__":
    # Example usage and testing
    print("AES Encryption Module - Example Usage")
    print("=" * 40)
    
    # Example 1: Simple string encryption with password
    password = "my_secure_password_123"
    message = "This is a secret message that needs to be encrypted!"
    
    print(f"Original message: {message}")
    
    # Encrypt
    encrypted = encrypt_string(message, password)
    print(f"Encrypted (base64): {encrypted[:50]}...")
    
    # Decrypt
    decrypted = decrypt_string(encrypted, password)
    print(f"Decrypted message: {decrypted}")
    print(f"Match: {message == decrypted}")
    
    print("\n" + "=" * 40)
    
    # Example 2: Using the class directly
    encryptor = AESEncryption(mode='GCM')
    key = encryptor.generate_key()
    
    print(f"Generated key info: {encryptor.get_key_info()}")
    
    test_data = b"Binary data to encrypt"
    ciphertext, iv = encryptor.encrypt(test_data)
    decrypted_data = encryptor.decrypt(ciphertext, iv)
    
    print(f"Binary encryption test: {test_data == decrypted_data}")