bigbrother/utils/crypto.py

#!/usr/bin/env python3
"""AES-256-GCM encryption, key derivation, and LUKS container management."""

import os
import struct
import subprocess
import hashlib
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

try:
    from argon2.low_level import hash_secret_raw, Type
    HAS_ARGON2 = True
except ImportError:
    HAS_ARGON2 = False

NONCE_SIZE = 12
KEY_SIZE = 32
SALT_SIZE = 16
HEADER_MAGIC = b"BB01"
HEADER_FMT = f"!4s{SALT_SIZE}s{NONCE_SIZE}s"
HEADER_SIZE = struct.calcsize(HEADER_FMT)


def derive_key(password: bytes, salt: bytes, method: str = "argon2id") -> bytes:
    if method == "argon2id" and HAS_ARGON2:
        return hash_secret_raw(
            secret=password,
            salt=salt,
            time_cost=3,
            memory_cost=65536,
            parallelism=4,
            hash_len=KEY_SIZE,
            type=Type.ID,
        )
    return PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=600_000,
    ).derive(password)


def derive_network_key(network_secret: bytes, device_serial: bytes) -> bytes:
    return HKDF(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=device_serial,
        info=b"bigbrother-luks-unlock",
    ).derive(network_secret)


class CryptoEngine:
    def __init__(self, key: bytes):
        if len(key) != KEY_SIZE:
            raise ValueError(f"Key must be {KEY_SIZE} bytes")
        self._aesgcm = AESGCM(key)
        self._key = key

    def encrypt(self, plaintext: bytes, aad: Optional[bytes] = None) -> bytes:
        nonce = os.urandom(NONCE_SIZE)
        ct = self._aesgcm.encrypt(nonce, plaintext, aad)
        return nonce + ct

    def decrypt(self, data: bytes, aad: Optional[bytes] = None) -> bytes:
        nonce = data[:NONCE_SIZE]
        ct = data[NONCE_SIZE:]
        return self._aesgcm.decrypt(nonce, ct, aad)




def encrypt_credential_dict(payload: dict, key: bytes) -> dict:
    """Encrypt sensitive fields in credential event payload before bus emission.
    
    Sensitive fields: username, password, hash, domain, token, secret, api_key
    Returns a new dict with encrypted fields as bytes wrapped in __encrypted__ markers.
    """
    if len(key) != KEY_SIZE:
        raise ValueError(f"Key must be {KEY_SIZE} bytes")
    
    engine = CryptoEngine(key)
    encrypted = {}
    
    sensitive_fields = {
        "username", "password", "hash", "hash_value", "domain",
        "token", "secret", "api_key", "access_token", "refresh_token",
        "credential_value", "nt_hash", "lm_hash", "response",
        "auth_string", "bearer_token"
    }
    
    for field, value in payload.items():
        if field.lower() in sensitive_fields and isinstance(value, str) and value:
            try:
                # Encrypt the string value
                plaintext = value.encode("utf-8")
                ciphertext = engine.encrypt(plaintext)
                # Encode to base64 for transport over JSON
                import base64
                encrypted[field] = f"__encrypted__{base64.b64encode(ciphertext).decode('ascii')}"
            except Exception:
                # If encryption fails, include plaintext (log warning in caller)
                encrypted[field] = value
        else:
            # Pass through unencrypted
            encrypted[field] = value
    
    return encrypted


def decrypt_credential_field(value: str, key: bytes) -> str:
    """Decrypt a single encrypted credential field."""
    if not isinstance(value, str) or not value.startswith("__encrypted__"):
        return value
    
    if len(key) != KEY_SIZE:
        raise ValueError(f"Key must be {KEY_SIZE} bytes")
    
    try:
        import base64
        engine = CryptoEngine(key)
        encrypted_part = value[len("__encrypted__"):]
        ciphertext = base64.b64decode(encrypted_part)
        plaintext = engine.decrypt(ciphertext)
        return plaintext.decode("utf-8")
    except Exception:
        # Return original if decryption fails
        return value

def encrypt_file(
    src: str,
    dst: str,
    password: bytes,
    method: str = "argon2id",
    chunk_size: int = 64 * 1024,
) -> None:
    salt = os.urandom(SALT_SIZE)
    key = derive_key(password, salt, method)
    engine = CryptoEngine(key)

    with open(src, "rb") as fin:
        plaintext = fin.read()

    nonce = os.urandom(NONCE_SIZE)
    aesgcm = AESGCM(key)
    ct = aesgcm.encrypt(nonce, plaintext, None)

    with open(dst, "wb") as fout:
        fout.write(struct.pack(HEADER_FMT, HEADER_MAGIC, salt, nonce))
        fout.write(ct)


def decrypt_file(
    src: str,
    dst: str,
    password: bytes,
    method: str = "argon2id",
) -> None:
    with open(src, "rb") as fin:
        header = fin.read(HEADER_SIZE)
        magic, salt, nonce = struct.unpack(HEADER_FMT, header)
        if magic != HEADER_MAGIC:
            raise ValueError("Invalid encrypted file header")
        ct = fin.read()

    key = derive_key(password, salt, method)
    aesgcm = AESGCM(key)
    plaintext = aesgcm.decrypt(nonce, ct, None)

    with open(dst, "wb") as fout:
        fout.write(plaintext)


class LUKSContainer:
    def __init__(self, container_path: str, mount_point: str, mapper_name: str = "bb_storage"):
        self.container_path = container_path
        self.mount_point = mount_point
        self.mapper_name = mapper_name
        self.device_path = f"/dev/mapper/{mapper_name}"

    def create(self, size_mb: int, passphrase: bytes) -> bool:
        try:
            subprocess.run(
                ["dd", "if=/dev/zero", f"of={self.container_path}",
                 "bs=1M", f"count={size_mb}"],
                check=True, capture_output=True,
            )
            proc = subprocess.run(
                ["cryptsetup", "luksFormat", "--type", "luks2",
                 "--cipher", "aes-xts-plain64", "--key-size", "512",
                 "--hash", "sha256", "--batch-mode", self.container_path],
                input=passphrase + b"\n",
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def open(self, passphrase: bytes) -> bool:
        try:
            subprocess.run(
                ["cryptsetup", "open", self.container_path,
                 self.mapper_name, "--type", "luks2"],
                input=passphrase + b"\n",
                check=True, capture_output=True,
            )
            Path(self.mount_point).mkdir(parents=True, exist_ok=True)
            subprocess.run(
                ["mount", self.device_path, self.mount_point],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def close(self) -> bool:
        try:
            subprocess.run(
                ["umount", self.mount_point],
                check=True, capture_output=True,
            )
            subprocess.run(
                ["cryptsetup", "close", self.mapper_name],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    def destroy_header(self) -> bool:
        try:
            subprocess.run(
                ["dd", "if=/dev/urandom", f"of={self.container_path}",
                 "bs=1M", "count=2", "conv=notrunc"],
                check=True, capture_output=True,
            )
            return True
        except subprocess.CalledProcessError:
            return False

    @property
    def is_open(self) -> bool:
        return Path(self.device_path).exists()


def secure_wipe_file(filepath: str, passes: int = 1) -> bool:
    """Securely wipe a file by overwriting with random data before deletion.
    
    Args:
        filepath: Path to file to wipe
        passes: Number of overwrite passes (default 1, use 3 for stronger security)
        
    Returns:
        True if file was successfully wiped, False otherwise
    """
    try:
        file_size = os.path.getsize(filepath)
        if file_size == 0:
            # Empty file, just delete
            os.unlink(filepath)
            return True
        
        # Overwrite passes
        for _ in range(passes):
            with open(filepath, "wb") as f:
                f.write(os.urandom(file_size))
                f.flush()
                os.fsync(f.fileno())
        
        # Delete the file
        os.unlink(filepath)
        return True
    except Exception as e:
        # Fallback: try regular delete if secure wipe fails
        try:
            os.unlink(filepath)
            return False
        except Exception:
            return False