encryption

"""
File encryption utilities using AES.
"""

import os
import base64
from pathlib import Path
from typing import Optional
import logging

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend

logger = logging.getLogger(__name__)

# AES block size
BLOCK_SIZE = 128  # bits
IV_SIZE = 16  # bytes


def _get_key_bytes(key: str) -> bytes:
    """Convert string key to 32 bytes for AES-256."""
    if len(key) == 32:
        return key.encode('utf-8')
    
    # If key is base64 encoded
    try:
        decoded = base64.b64decode(key)
        if len(decoded) == 32:
            return decoded
    except:
        pass
    
    # Pad or truncate key
    key_bytes = key.encode('utf-8')
    if len(key_bytes) < 32:
        key_bytes = key_bytes.ljust(32, b'\0')
    return key_bytes[:32]


def encrypt_file(
    input_path: Path,
    key: str,
    output_path: Optional[Path] = None,
    chunk_size: int = 8192
) -> Path:
    """
    Encrypt a file using AES-256-CBC.
    
    Args:
        input_path: Path to input file
        key: Encryption key (string)
        output_path: Path for encrypted file (default: input + .enc)
        chunk_size: Size of chunks to read
    
    Returns:
        Path to encrypted file
    """
    if output_path is None:
        output_path = input_path.with_suffix(input_path.suffix + ".enc")
    
    key_bytes = _get_key_bytes(key)
    iv = os.urandom(IV_SIZE)
    
    cipher = Cipher(
        algorithms.AES(key_bytes),
        modes.CBC(iv),
        backend=default_backend()
    )
    encryptor = cipher.encryptor()
    padder = padding.PKCS7(BLOCK_SIZE).padder()
    
    with open(input_path, 'rb') as f_in:
        with open(output_path, 'wb') as f_out:
            # Write IV first (needed for decryption)
            f_out.write(iv)
            
            while True:
                chunk = f_in.read(chunk_size)
                
                if chunk:
                    padded_chunk = padder.update(chunk)
                    encrypted_chunk = encryptor.update(padded_chunk)
                    f_out.write(encrypted_chunk)
                else:
                    # Final block
                    padded_final = padder.finalize()
                    encrypted_final = encryptor.update(padded_final) + encryptor.finalize()
                    f_out.write(encrypted_final)
                    break
    
    logger.debug(f"Encrypted: {input_path} -> {output_path}")
    return output_path


def decrypt_file(
    input_path: Path,
    key: str,
    output_path: Optional[Path] = None,
    chunk_size: int = 8192
) -> Path:
    """
    Decrypt a file encrypted with AES-256-CBC.
    
    Args:
        input_path: Path to encrypted file
        key: Encryption key (string)
        output_path: Path for decrypted file (default: remove .enc)
        chunk_size: Size of chunks to read
    
    Returns:
        Path to decrypted file
    """
    if output_path is None:
        if input_path.suffix == ".enc":
            output_path = input_path.with_suffix("")
        else:
            output_path = input_path.with_suffix(".decrypted")
    
    key_bytes = _get_key_bytes(key)
    
    with open(input_path, 'rb') as f_in:
        # Read IV from file
        iv = f_in.read(IV_SIZE)
        
        cipher = Cipher(
            algorithms.AES(key_bytes),
            modes.CBC(iv),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        unpadder = padding.PKCS7(BLOCK_SIZE).unpadder()
        
        with open(output_path, 'wb') as f_out:
            encrypted_data = f_in.read()
            decrypted_padded = decryptor.update(encrypted_data) + decryptor.finalize()
            decrypted_data = unpadder.update(decrypted_padded) + unpadder.finalize()
            f_out.write(decrypted_data)
    
    logger.debug(f"Decrypted: {input_path} -> {output_path}")
    return output_path


def generate_key() -> str:
    """
    Generate a random 32-byte key for AES-256.
    
    Returns:
        Base64-encoded key string
    """
    key_bytes = os.urandom(32)
    return base64.b64encode(key_bytes).decode('utf-8')


class FileEncryptor:
    """Convenience class for file encryption."""
    
    def __init__(self, key: str):
        self.key = key
    
    def encrypt(self, input_path: Path, output_path: Optional[Path] = None) -> Path:
        """Encrypt a file."""
        return encrypt_file(input_path, self.key, output_path)
    
    def decrypt(self, input_path: Path, output_path: Optional[Path] = None) -> Path:
        """Decrypt a file."""
        return decrypt_file(input_path, self.key, output_path)