ops-bridge/src/bridge/manager.py

"""Tunnel lifecycle manager for OpsBridge."""
from __future__ import annotations

import logging
import os
import signal
import subprocess
import time
from pathlib import Path
from typing import List, Optional

from bridge.audit import AuditEvent, AuditLogger
from bridge.health import HealthChecker
from bridge.models import BridgeState, TunnelConfig
from bridge.state import StateManager

log = logging.getLogger(__name__)


def build_ssh_command(cfg: TunnelConfig) -> List[str]:
    """Build the SSH tunnel command (reverse -R or local -L)."""
    key = os.path.expanduser(cfg.ssh_key)
    if cfg.direction == "local":
        forward_flag = ["-L", f"{cfg.local_port}:127.0.0.1:{cfg.remote_port}"]
    else:
        forward_flag = ["-R", f"{cfg.remote_port}:127.0.0.1:{cfg.local_port}"]
    return [
        "ssh",
        "-N",
        *forward_flag,
        "-i", key,
        "-o", "ServerAliveInterval=10",
        "-o", "ServerAliveCountMax=3",
        "-o", "ExitOnForwardFailure=yes",
        "-o", "StrictHostKeyChecking=accept-new",
        f"{cfg.ssh_user}@{cfg.host}",
    ]


class TunnelManager:
    """Manages a single named SSH reverse tunnel.

    start() daemonises: forks a child that runs the reconnect loop, then the
    parent returns immediately after writing the manager PID.
    """

    def __init__(self, cfg: TunnelConfig, state_dir: Optional[Path] = None):
        self._cfg = cfg
        self._state = StateManager(state_dir=state_dir)
        self._audit = AuditLogger(state_dir=state_dir)

    def get_state(self) -> BridgeState:
        return self._state.read_state(self._cfg.name)

    def is_running(self) -> bool:
        return self._state.is_running(self._cfg.name)

    def _actor_info(self):
        return self._cfg.actor, "unknown"

    def _next_backoff(self, attempt: int) -> int:
        initial = self._cfg.reconnect.backoff_initial
        max_b = self._cfg.reconnect.backoff_max
        value = initial * (2 ** attempt)
        return min(value, max_b)

    def start(self) -> None:
        """Start the tunnel manager as a daemonised subprocess."""
        if self.is_running():
            log.info("Tunnel %s already running", self._cfg.name)
            return

        self._state.write_state(self._cfg.name, BridgeState.STARTING)
        actor, actor_class = self._actor_info()
        self._audit.log(
            tunnel=self._cfg.name,
            event=AuditEvent.BRIDGE_STARTED,
            actor=actor,
            actor_class=actor_class,
        )

        pid = os.fork()
        if pid > 0:
            # Parent: record manager PID and return
            self._state.write_pid(self._cfg.name, pid)
            return

        # Child: become a daemon
        os.setsid()

        try:
            self._run_loop()
        except Exception as e:
            log.exception("Tunnel manager loop crashed: %s", e)
        finally:
            self._state.write_state(self._cfg.name, BridgeState.STOPPED)
            self._state.clear_pid(self._cfg.name)
            self._audit.log(
                tunnel=self._cfg.name,
                event=AuditEvent.BRIDGE_STOPPED,
                actor=actor,
                actor_class=actor_class,
            )

        os._exit(0)

    def stop(self) -> None:
        """Stop the running tunnel manager."""
        pid = self._state.read_pid(self._cfg.name)
        if pid is None:
            self._state.write_state(self._cfg.name, BridgeState.STOPPED)
            return

        try:
            os.kill(pid, signal.SIGTERM)
            # Give up to 5 seconds for graceful shutdown
            for _ in range(50):
                try:
                    os.kill(pid, 0)
                    time.sleep(0.1)
                except ProcessLookupError:
                    break
            else:
                # Force kill if still running
                try:
                    os.kill(pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
        except ProcessLookupError:
            pass

        self._state.clear_pid(self._cfg.name)
        self._state.write_state(self._cfg.name, BridgeState.STOPPED)
        actor, actor_class = self._actor_info()
        self._audit.log(
            tunnel=self._cfg.name,
            event=AuditEvent.BRIDGE_STOPPED,
            actor=actor,
            actor_class=actor_class,
        )

    def _run_loop(self) -> None:
        """Reconnect loop running in daemon child."""
        import asyncio

        cfg = self._cfg
        actor, actor_class = self._actor_info()
        attempt = 0
        max_attempts = cfg.reconnect.max_attempts  # 0 = infinite

        # Setup signal handler for graceful shutdown
        _stop = [False]

        def _on_term(signum, frame):
            _stop[0] = True

        signal.signal(signal.SIGTERM, _on_term)
        signal.signal(signal.SIGINT, _on_term)

        while not _stop[0]:
            if max_attempts > 0 and attempt >= max_attempts:
                self._state.write_state(cfg.name, BridgeState.FAILED)
                break

            cmd = build_ssh_command(cfg)
            log.info("Starting SSH: %s", " ".join(cmd))
            self._state.write_state(cfg.name, BridgeState.STARTING)

            try:
                proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            except FileNotFoundError:
                self._state.write_state(cfg.name, BridgeState.FAILED)
                self._audit.log(
                    tunnel=cfg.name,
                    event=AuditEvent.BRIDGE_DISCONNECTED,
                    actor=actor,
                    actor_class=actor_class,
                    detail="ssh binary not found",
                )
                break

            # Wait briefly then assume connected if still running
            time.sleep(2)
            if proc.poll() is None:
                self._state.write_state(cfg.name, BridgeState.CONNECTED)
                self._audit.log(
                    tunnel=cfg.name,
                    event=AuditEvent.BRIDGE_CONNECTED,
                    actor=actor,
                    actor_class=actor_class,
                )
                attempt = 0

                # Health check loop
                if cfg.health_check:
                    checker = HealthChecker(
                        url=cfg.health_check.url,
                        timeout_seconds=cfg.health_check.timeout_seconds,
                    )
                    health_failing = False
                    while not _stop[0] and proc.poll() is None:
                        result = asyncio.run(checker.check())
                        if result.ok:
                            if health_failing:
                                health_failing = False
                                self._state.write_state(cfg.name, BridgeState.CONNECTED)
                                self._audit.log(
                                    tunnel=cfg.name,
                                    event=AuditEvent.HEALTH_CHECK_RECOVERED,
                                    actor=actor,
                                    actor_class=actor_class,
                                )
                        else:
                            if not health_failing:
                                health_failing = True
                                self._state.write_state(cfg.name, BridgeState.DEGRADED)
                                self._audit.log(
                                    tunnel=cfg.name,
                                    event=AuditEvent.HEALTH_CHECK_FAILED,
                                    actor=actor,
                                    actor_class=actor_class,
                                    detail=result.error or f"HTTP {result.status_code}",
                                )
                        time.sleep(cfg.health_check.interval_seconds)
                else:
                    while not _stop[0] and proc.poll() is None:
                        time.sleep(1)

            # SSH exited
            if proc.poll() is not None:
                self._audit.log(
                    tunnel=cfg.name,
                    event=AuditEvent.BRIDGE_DISCONNECTED,
                    actor=actor,
                    actor_class=actor_class,
                    detail=f"exit code {proc.returncode}",
                )

            if _stop[0]:
                if proc.poll() is None:
                    proc.terminate()
                break

            attempt += 1
            backoff = self._next_backoff(attempt - 1)
            self._state.write_state(cfg.name, BridgeState.RECONNECTING)
            self._audit.log(
                tunnel=cfg.name,
                event=AuditEvent.BRIDGE_RECONNECTING,
                actor=actor,
                actor_class=actor_class,
                detail=f"retry {attempt}, backoff {backoff}s",
            )
            log.info("Reconnecting in %ds (attempt %d)", backoff, attempt)
            time.sleep(backoff)