ai-agent/security/encrypt.py

import ssl

def create_tls_context(purpose=ssl.Purpose.CLIENT_AUTH):
    """
    Creates an SSL context configured for TLS 1.3.

    Args:
        purpose: The SSL purpose (e.g., CLIENT_AUTH, SERVER_AUTH).

    Returns:
        An SSLContext object configured for TLS 1.3.
    """
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT if purpose == ssl.Purpose.CLIENT_AUTH else ssl.PROTOCOL_TLS_SERVER)

    # Require TLS 1.3
    context.minimum_version = ssl.TLSVersion.TLSv1_3

    # TLS 1.3 cipher suites are handled automatically by the underlying SSL library
    # when minimum_version is set to TLSv1_3. Explicitly setting them via
    # set_ciphers can cause issues. The required suites (AES-256-GCM, CHACHA20)
    # are typically included and preferred by default in modern OpenSSL.
    # context.set_ciphers('...') # Removed

    # Configure certificate loading and verification
    if purpose == ssl.Purpose.SERVER_AUTH:
        # Server context: Load server cert/key and require client certs for RBAC
        # context.load_cert_chain(certfile="path/to/server_cert.pem", keyfile="path/to/server_key.pem") # Placeholder: Needs actual paths
        context.verify_mode = ssl.CERT_REQUIRED
        # context.load_verify_locations(cafile="path/to/trusted_client_ca.pem") # Placeholder: Needs actual CA path for client cert validation
    elif purpose == ssl.Purpose.CLIENT_AUTH:
        # Client context: Load client cert/key and verify server cert against CA
        # context.load_cert_chain(certfile="path/to/client_cert.pem", keyfile="path/to/client_key.pem") # Placeholder: Needs actual paths
        # context.load_verify_locations(cafile="path/to/trusted_server_ca.pem") # Placeholder: Needs actual CA path
        context.verify_mode = ssl.CERT_REQUIRED # Verify the server certificate

    # minimum_version = TLSv1_3 implicitly disables older protocols.
    # Explicit OP_NO flags are redundant but harmless; removed for clarity.
    # context.options |= ssl.OP_NO_SSLv2 # Redundant
    # context.options |= ssl.OP_NO_SSLv3 # Redundant
    # context.options |= ssl.OP_NO_TLSv1 # Redundant
    # context.options |= ssl.OP_NO_TLSv1_1 # Redundant
    # context.options |= ssl.OP_NO_TLSv1_2 # Redundant

    # Options for perfect forward secrecy are generally enabled by default with TLS 1.3 ciphers
    # context.options |= ssl.OP_SINGLE_DH_USE
    # context.options |= ssl.OP_SINGLE_ECDH_USE

    # Ensure TLS 1.3 is the minimum (already set, but good to be explicit)
    context.minimum_version = ssl.TLSVersion.TLSv1_3

    return context

# Example usage (can be removed or kept for demonstration)
if __name__ == '__main__':
    client_context = create_tls_context(ssl.Purpose.CLIENT_AUTH)
    print(f"Client Context Minimum TLS Version: {client_context.minimum_version}")
    # print(f"Client Context Ciphers: {client_context.get_ciphers()}") # Requires OpenSSL 1.1.1+

    server_context = create_tls_context(ssl.Purpose.SERVER_AUTH)
    print(f"Server Context Minimum TLS Version: {server_context.minimum_version}")
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import os

class AES256Cipher:
    """AES-256-GCM encryption/decryption wrapper class."""
    
    def __init__(self, key: bytes = None):
        """
        Initialize cipher with optional key.
        
        Args:
            key: Optional 32-byte AES-256 key. If None, generates new key.
        """
        self.key = key if key is not None else self.generate_key()
    
    @staticmethod
    def generate_key() -> bytes:
        """Generate a secure 256-bit AES key for encryption/decryption.
        
        Returns:
            bytes: 32-byte AES-256 key
        """
        return os.urandom(32)
    
    def encrypt(self, plaintext: bytes) -> bytes:
        """Encrypt data using AES-256-GCM.
        
        Args:
            plaintext: Data to encrypt
            
        Returns:
            bytes: Encrypted data in format (nonce + ciphertext + tag)
        """
        return encrypt_data(plaintext, self.key)
    
    def decrypt(self, encrypted_data: bytes) -> bytes:
        """Decrypt data using AES-256-GCM.
        
        Args:
            encrypted_data: Data in format (nonce + ciphertext + tag)
            
        Returns:
            bytes: Decrypted plaintext
        """
        return decrypt_data(encrypted_data, self.key)

def generate_key():
    """Generate a secure 256-bit AES key for encryption/decryption.
    
    Returns:
        bytes: 32-byte AES-256 key
    """
    return os.urandom(32)

def encrypt_data(plaintext: bytes, key: bytes) -> bytes:
    """Encrypt data using AES-256-GCM.
    
    Args:
        plaintext: Data to encrypt
        key: 32-byte AES-256 key
        
    Returns:
        bytes: Encrypted data in format (nonce + ciphertext + tag)
        
    Raises:
        ValueError: If key length is invalid
    """
    if len(key) != 32:
        raise ValueError("Key must be 32 bytes for AES-256")
        
    # Generate random 96-bit nonce
    nonce = os.urandom(12)
    
    # Create cipher and encrypt
    cipher = Cipher(
        algorithms.AES(key),
        modes.GCM(nonce),
        backend=default_backend()
    )
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
    
    # Return nonce + ciphertext + tag
    return nonce + ciphertext + encryptor.tag

def decrypt_data(encrypted_data: bytes, key: bytes) -> bytes:
    """Decrypt data using AES-256-GCM.
    
    Args:
        encrypted_data: Data in format (nonce + ciphertext + tag)
        key: 32-byte AES-256 key
        
    Returns:
        bytes: Decrypted plaintext
        
    Raises:
        ValueError: If key length is invalid or data is malformed
    """
    if len(key) != 32:
        raise ValueError("Key must be 32 bytes for AES-256")
    if len(encrypted_data) < 28:  # Minimum: 12 nonce + 16 tag
        raise ValueError("Encrypted data too short")
        
    # Split into components
    nonce = encrypted_data[:12]
    ciphertext = encrypted_data[12:-16]
    tag = encrypted_data[-16:]
    
    # Create cipher and decrypt
    cipher = Cipher(
        algorithms.AES(key),
        modes.GCM(nonce, tag),
        backend=default_backend()
    )
    decryptor = cipher.decryptor()
    return decryptor.update(ciphertext) + decryptor.finalize()

# Example usage for AES-256-GCM functions
if __name__ == '__main__':
    # Generate key
    key = generate_key()
    print(f"Generated AES-256 key: {key.hex()}")
    
    # Encrypt test data
    plaintext = b"Test message for AES-256-GCM implementation"
    encrypted = encrypt_data(plaintext, key)
    print(f"Encrypted data (hex): {encrypted.hex()}")
    
    # Decrypt test data
    decrypted = decrypt_data(encrypted, key)
    print(f"Decrypted data: {decrypted.decode()}")
    # print(f"Server Context Ciphers: {server_context.get_ciphers()}") # Requires OpenSSL 1.1.1+