crypto_simple.py

"""
PhishGuard M3 - Cryptography Module
AES-256-GCM encryption with scrypt key derivation
Production-grade security implementation
"""

import os
import json
import base64
from typing import Dict, Any
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.backends import default_backend


# Constants for cryptographic parameters
SALT_SIZE = 16  # 128 bits
NONCE_SIZE = 12  # 96 bits (recommended for GCM)
KEY_SIZE = 32  # 256 bits for AES-256
SCRYPT_N = 2**14  # CPU/memory cost parameter (16384)
SCRYPT_R = 8  # Block size parameter
SCRYPT_P = 1  # Parallelization parameter


def derive_key(passphrase: str, salt: bytes) -> bytes:
    """
    Derive a 256-bit encryption key from passphrase using scrypt.
    
    Args:
        passphrase: User-provided passphrase
        salt: Random salt (16 bytes)
        
    Returns:
        32-byte encryption key
    """
    kdf = Scrypt(
        salt=salt,
        length=KEY_SIZE,
        n=SCRYPT_N,
        r=SCRYPT_R,
        p=SCRYPT_P,
        backend=default_backend()
    )
    key = kdf.derive(passphrase.encode('utf-8'))
    return key


def encrypt_data(data: Dict[str, Any], passphrase: str) -> bytes:
    """
    Encrypt a dictionary using AES-256-GCM.
    
    Format: salt (16) + nonce (12) + ciphertext+tag
    
    Args:
        data: Dictionary to encrypt
        passphrase: Encryption passphrase
        
    Returns:
        Encrypted blob as bytes
        
    Raises:
        ValueError: If passphrase is empty
        TypeError: If data cannot be serialized to JSON
    """
    if not passphrase:
        raise ValueError("Passphrase cannot be empty")
    
    # Generate random salt and nonce
    salt = os.urandom(SALT_SIZE)
    nonce = os.urandom(NONCE_SIZE)
    
    # Derive encryption key from passphrase
    key = derive_key(passphrase, salt)
    
    # Serialize data to JSON
    try:
        plaintext = json.dumps(data, indent=2).encode('utf-8')
    except (TypeError, ValueError) as e:
        raise TypeError(f"Data cannot be serialized to JSON: {e}")
    
    # Encrypt using AES-256-GCM
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext, None)
    
    # Combine: salt + nonce + ciphertext+tag
    encrypted_blob = salt + nonce + ciphertext
    
    return encrypted_blob


def decrypt_data(blob: bytes, passphrase: str) -> Dict[str, Any]:
    """
    Decrypt an encrypted blob back to dictionary.
    
    Args:
        blob: Encrypted data (salt + nonce + ciphertext+tag)
        passphrase: Decryption passphrase
        
    Returns:
        Decrypted dictionary
        
    Raises:
        ValueError: If blob format is invalid or passphrase is wrong
    """
    if not passphrase:
        raise ValueError("Passphrase cannot be empty")
    
    # Validate minimum blob size
    min_size = SALT_SIZE + NONCE_SIZE + 16  # 16 = minimum GCM tag size
    if len(blob) < min_size:
        raise ValueError("Invalid encrypted blob: too short")
    
    # Extract components
    salt = blob[:SALT_SIZE]
    nonce = blob[SALT_SIZE:SALT_SIZE + NONCE_SIZE]
    ciphertext = blob[SALT_SIZE + NONCE_SIZE:]
    
    # Derive key from passphrase
    key = derive_key(passphrase, salt)
    
    # Decrypt using AES-256-GCM
    aesgcm = AESGCM(key)
    try:
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    except Exception as e:
        raise ValueError(f"Decryption failed: incorrect passphrase or corrupted data ({e})")
    
    # Deserialize JSON
    try:
        data = json.loads(plaintext.decode('utf-8'))
    except (json.JSONDecodeError, UnicodeDecodeError) as e:
        raise ValueError(f"Invalid decrypted data format: {e}")
    
    return data


def save_encrypted(data: Dict[str, Any], filepath: str, passphrase: str) -> None:
    """
    Encrypt and save data to a file.
    
    Args:
        data: Dictionary to encrypt and save
        filepath: Path to output file
        passphrase: Encryption passphrase
        
    Raises:
        IOError: If file cannot be written
        ValueError: If encryption fails
    """
    encrypted_blob = encrypt_data(data, passphrase)
    
    try:
        with open(filepath, 'wb') as f:
            f.write(encrypted_blob)
    except IOError as e:
        raise IOError(f"Failed to write encrypted file: {e}")


def load_encrypted(filepath: str, passphrase: str) -> Dict[str, Any]:
    """
    Load and decrypt data from a file.
    
    Args:
        filepath: Path to encrypted file
        passphrase: Decryption passphrase
        
    Returns:
        Decrypted dictionary
        
    Raises:
        IOError: If file cannot be read
        ValueError: If decryption fails
    """
    try:
        with open(filepath, 'rb') as f:
            encrypted_blob = f.read()
    except IOError as e:
        raise IOError(f"Failed to read encrypted file: {e}")
    
    return decrypt_data(encrypted_blob, passphrase)


# Optional: Base64 encoding helpers for text-based storage/transmission

def encrypt_to_base64(data: Dict[str, Any], passphrase: str) -> str:
    """
    Encrypt data and return as base64 string.
    
    Args:
        data: Dictionary to encrypt
        passphrase: Encryption passphrase
        
    Returns:
        Base64-encoded encrypted blob
    """
    encrypted_blob = encrypt_data(data, passphrase)
    return base64.b64encode(encrypted_blob).decode('ascii')


def decrypt_from_base64(b64_string: str, passphrase: str) -> Dict[str, Any]:
    """
    Decrypt base64-encoded encrypted data.
    
    Args:
        b64_string: Base64-encoded encrypted blob
        passphrase: Decryption passphrase
        
    Returns:
        Decrypted dictionary
        
    Raises:
        ValueError: If base64 decoding or decryption fails
    """
    try:
        encrypted_blob = base64.b64decode(b64_string)
    except Exception as e:
        raise ValueError(f"Invalid base64 string: {e}")
    
    return decrypt_data(encrypted_blob, passphrase)


# Security note: Always use strong passphrases (12+ characters, mixed case, numbers, symbols)
# The scrypt parameters are tuned for security vs. performance balance
# Increase SCRYPT_N for higher security at the cost of slower key derivation