crypto

"""
Cryptographic operations for password manager.
"""

import os
import base64
import hashlib
import hmac
import secrets
from typing import Tuple

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend

# Constants
SALT_SIZE = 32  # bytes
NONCE_SIZE = 12  # bytes for AES-GCM
KEY_SIZE = 32  # 256 bits
ITERATIONS = 100_000  # PBKDF2 iterations


def generate_salt() -> bytes:
    """Generate a random salt."""
    return secrets.token_bytes(SALT_SIZE)


def generate_nonce() -> bytes:
    """Generate a random nonce for AES-GCM."""
    return secrets.token_bytes(NONCE_SIZE)


def derive_key(password: str, salt: bytes, iterations: int = ITERATIONS) -> bytes:
    """
    Derive an encryption key from a password using PBKDF2.
    
    Args:
        password: Master password
        salt: Random salt
        iterations: Number of PBKDF2 iterations
    
    Returns:
        256-bit key
    """
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=iterations,
        backend=default_backend()
    )
    
    return kdf.derive(password.encode('utf-8'))


def encrypt(plaintext: bytes, key: bytes) -> bytes:
    """
    Encrypt data using AES-256-GCM.
    
    Args:
        plaintext: Data to encrypt (bytes)
        key: 256-bit encryption key
    
    Returns:
        nonce + ciphertext with auth tag
    """
    nonce = generate_nonce()
    aesgcm = AESGCM(key)
    
    ciphertext = aesgcm.encrypt(nonce, plaintext, None)
    
    # Prepend nonce to ciphertext
    return nonce + ciphertext


def decrypt(ciphertext: bytes, key: bytes) -> bytes:
    """
    Decrypt data using AES-256-GCM.
    
    Args:
        ciphertext: nonce + encrypted data with auth tag
        key: 256-bit encryption key
    
    Returns:
        Decrypted plaintext
    
    Raises:
        cryptography.exceptions.InvalidTag: If authentication fails
    """
    # Extract nonce from beginning
    nonce = ciphertext[:NONCE_SIZE]
    actual_ciphertext = ciphertext[NONCE_SIZE:]
    
    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, actual_ciphertext, None)


def encrypt_data(plaintext: str, key: bytes) -> Tuple[bytes, bytes]:
    """
    Encrypt data using AES-256-GCM.
    
    Args:
        plaintext: Data to encrypt
        key: 256-bit encryption key
    
    Returns:
        Tuple of (nonce, ciphertext with auth tag)
    """
    nonce = generate_nonce()
    aesgcm = AESGCM(key)
    
    ciphertext = aesgcm.encrypt(
        nonce,
        plaintext.encode('utf-8'),
        None  # Additional authenticated data (not used)
    )
    
    return nonce, ciphertext


def decrypt_data(nonce: bytes, ciphertext: bytes, key: bytes) -> str:
    """
    Decrypt data using AES-256-GCM.
    
    Args:
        nonce: Nonce used for encryption
        ciphertext: Encrypted data with auth tag
        key: 256-bit encryption key
    
    Returns:
        Decrypted plaintext
    
    Raises:
        cryptography.exceptions.InvalidTag: If authentication fails
    """
    aesgcm = AESGCM(key)
    
    plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    
    return plaintext.decode('utf-8')


def secure_compare(a: bytes, b: bytes) -> bool:
    """
    Constant-time comparison of two byte strings.
    Prevents timing attacks.
    
    Args:
        a: First byte string
        b: Second byte string
    
    Returns:
        True if equal
    """
    return secrets.compare_digest(a, b)


def generate_hmac(data: bytes, key: bytes) -> bytes:
    """
    Generate HMAC-SHA256 for data authentication.
    
    Args:
        data: Data to authenticate
        key: Secret key
    
    Returns:
        HMAC digest
    """
    return hmac.new(key, data, hashlib.sha256).digest()


def verify_hmac(data: bytes, expected_hmac: bytes, key: bytes) -> bool:
    """
    Verify HMAC-SHA256 for data authentication.
    
    Args:
        data: Data to verify
        expected_hmac: Expected HMAC value
        key: Secret key
    
    Returns:
        True if HMAC is valid
    """
    actual_hmac = generate_hmac(data, key)
    return secrets.compare_digest(actual_hmac, expected_hmac)


def encode_base64(data: bytes) -> str:
    """Encode bytes to base64 string."""
    return base64.b64encode(data).decode('utf-8')


def decode_base64(data: str) -> bytes:
    """Decode base64 string to bytes."""
    return base64.b64decode(data.encode('utf-8'))


def hash_password(password: str) -> str:
    """
    Hash a password for storage/comparison.
    Uses SHA-256.
    
    Args:
        password: Password to hash
    
    Returns:
        Hex digest of hash
    """
    return hashlib.sha256(password.encode('utf-8')).hexdigest()


def secure_random_bytes(length: int) -> bytes:
    """Generate cryptographically secure random bytes."""
    return secrets.token_bytes(length)


def secure_random_string(length: int, alphabet: str) -> str:
    """
    Generate a secure random string from given alphabet.
    
    Args:
        length: Length of string
        alphabet: Characters to choose from
    
    Returns:
        Random string
    """
    return ''.join(secrets.choice(alphabet) for _ in range(length))


class SecureString:
    """
    A string class that attempts to securely wipe memory.
    Note: Python's memory management makes this imperfect.
    """
    
    def __init__(self, value: str):
        self._value = value
    
    def get(self) -> str:
        """Get the string value."""
        return self._value
    
    def clear(self):
        """Attempt to clear the string from memory."""
        # Overwrite with random data
        if self._value:
            length = len(self._value)
            self._value = '\x00' * length
            self._value = ''
    
    def __del__(self):
        """Destructor attempts to clear memory."""
        self.clear()
    
    def __str__(self):
        return '[SECURE STRING]'
    
    def __repr__(self):
        return 'SecureString([HIDDEN])'