#!/usr/bin/env python3
"""
Connectivity Manager for Personal Internet Cell — Phase 5 Extended Connectivity.

Provides per-peer egress routing through alternate exits (WireGuard external,
OpenVPN, Tor) via Linux policy routing (fwmark + ip rule + dedicated routing
tables) and dedicated iptables chains.

Architecture
------------
- A peer's `exit_via` field selects the egress path: "default", "wireguard_ext",
  "openvpn", or "tor".
- Each non-default exit type is assigned a unique fwmark and a dedicated routing
  table:

      wireguard_ext  mark 0x10  table 110  iface wg_ext0
      openvpn        mark 0x20  table 120  iface tun0
      tor            mark 0x30  table 130  (transparent proxy → 9040)

- All rules live in dedicated PIC_CONNECTIVITY chains in the `mangle` and `nat`
  tables so they can be flushed/rebuilt without touching firewall_manager rules.
- A kill-switch FORWARD DROP rule prevents leaks if the exit interface is down.

Container model
---------------
Each exit type runs in its own separate container; this manager only programs
policy routing rules in the WireGuard server container (cell-wireguard) where
peer traffic flows through.

Config files
------------
- WireGuard external:  {config_dir}/connectivity/wireguard_ext/wg_ext0.conf
- OpenVPN:             {config_dir}/connectivity/openvpn/<name>.ovpn

Both are validated to strip / reject hook directives that could execute
arbitrary commands on the host.
"""

import logging
import os
import re
import subprocess
from typing import Any, Dict, List, Optional

from base_service_manager import BaseServiceManager

logger = logging.getLogger(__name__)

WIREGUARD_CONTAINER = 'cell-wireguard'

# Lines we strip from uploaded WireGuard configs — these can run arbitrary
# host commands when wg-quick brings the interface up/down.
_WG_FORBIDDEN_PREFIXES = ('PostUp', 'PostDown', 'PreUp', 'PreDown')

# Lines we strip from uploaded OpenVPN configs — these execute external
# scripts/binaries on connect/disconnect.
_OVPN_FORBIDDEN_DIRECTIVES = (
    'up', 'down', 'script-security', 'plugin',
    'route-up', 'route-pre-down',
)

_NAME_RE = re.compile(r'^[a-z0-9_-]{1,32}$')


class ConnectivityManager(BaseServiceManager):
    """Manages alternate egress paths (extended connectivity) for peers."""

    EXIT_TYPES = ("default", "wireguard_ext", "openvpn", "tor")
    MARKS = {"wireguard_ext": 0x10, "openvpn": 0x20, "tor": 0x30}
    TABLES = {"wireguard_ext": 110, "openvpn": 120, "tor": 130}
    IFACES = {"wireguard_ext": "wg_ext0", "openvpn": "tun0"}
    TOR_TRANS_PORT = 9040
    TOR_DNS_PORT = 5353

    CONNECTIVITY_CHAIN = 'PIC_CONNECTIVITY'

    def __init__(self, config_manager=None, peer_registry=None,
                 data_dir: str = '/app/data', config_dir: str = '/app/config'):
        super().__init__('connectivity', data_dir, config_dir)
        self.config_manager = config_manager
        self.peer_registry = peer_registry

        # Config file directories
        self.connectivity_config_dir = os.path.join(config_dir, 'connectivity')
        self.wireguard_ext_dir = os.path.join(self.connectivity_config_dir, 'wireguard_ext')
        self.openvpn_dir = os.path.join(self.connectivity_config_dir, 'openvpn')

        for d in (self.connectivity_config_dir, self.wireguard_ext_dir, self.openvpn_dir):
            self.safe_makedirs(d)

        # Subscribe to ServiceBus CONFIG_CHANGED events so routes are
        # reapplied if the underlying network changes.  Done lazily —
        # service_bus is a singleton imported at app startup.
        self._subscribe_to_events()

    # ── Event wiring ──────────────────────────────────────────────────────

    def _subscribe_to_events(self) -> None:
        """Subscribe to network change events so routes auto-reapply."""
        try:
            from managers import service_bus, EventType
            service_bus.subscribe_to_event(
                EventType.CONFIG_CHANGED, self._on_network_changed
            )
        except Exception as e:
            # Non-fatal: subscription is best-effort, manual apply still works.
            logger.debug(f"connectivity: event subscribe skipped: {e}")

    def _on_network_changed(self, event) -> None:
        """ServiceBus handler: re-apply routes when network config changes."""
        try:
            source = getattr(event, 'source', '')
            if source not in ('network', 'wireguard', 'connectivity'):
                return
            logger.info(f"connectivity: re-applying routes due to {source} change")
            self.apply_routes()
        except Exception as e:
            logger.warning(f"connectivity: on_network_changed failed (non-fatal): {e}")

    # ── BaseServiceManager required ───────────────────────────────────────

    def get_status(self) -> Dict[str, Any]:
        """Return status summary including configured exits and peer count."""
        try:
            exits_status: Dict[str, Dict[str, Any]] = {}
            for exit_type in self.EXIT_TYPES:
                if exit_type == "default":
                    continue
                exits_status[exit_type] = self._exit_status(exit_type)

            peers_with_exit = 0
            if self.peer_registry is not None:
                try:
                    for peer in self.peer_registry.list_peers():
                        if peer.get('exit_via', 'default') != 'default':
                            peers_with_exit += 1
                except Exception as e:
                    logger.warning(f"get_status: peer count failed: {e}")

            return {
                'service': 'connectivity',
                'running': True,
                'exits': exits_status,
                'peers_with_exit': peers_with_exit,
            }
        except Exception as e:
            return self.handle_error(e, 'get_status')

    def test_connectivity(self) -> Dict[str, Any]:
        """Minimal connectivity self-test."""
        return {'success': True}

    def get_config(self) -> Dict[str, Any]:
        """Return current connectivity config from config_manager."""
        try:
            if self.config_manager is not None and hasattr(
                self.config_manager, 'get_connectivity_config'
            ):
                return self.config_manager.get_connectivity_config()
        except Exception as e:
            logger.warning(f"get_config: config_manager lookup failed: {e}")
        return {'exits': {}, 'peer_exit_map': {}}

    # ── Public API ────────────────────────────────────────────────────────

    def list_exits(self) -> List[Dict[str, Any]]:
        """List configured exits with current status."""
        result: List[Dict[str, Any]] = []
        for exit_type in self.EXIT_TYPES:
            if exit_type == "default":
                continue
            entry = {'type': exit_type}
            entry.update(self._exit_status(exit_type))
            result.append(entry)
        return result

    def get_peer_exits(self) -> Dict[str, str]:
        """Return {peer_name: exit_type} for all peers."""
        out: Dict[str, str] = {}
        if self.peer_registry is None:
            return out
        try:
            for peer in self.peer_registry.list_peers():
                name = peer.get('peer')
                if name:
                    out[name] = peer.get('exit_via', 'default')
        except Exception as e:
            logger.warning(f"get_peer_exits: {e}")
        return out

    def set_peer_exit(self, peer_name: str, exit_type: str) -> Dict[str, Any]:
        """Assign a peer to an egress path and apply the rule changes."""
        if exit_type not in self.EXIT_TYPES:
            return {
                'ok': False,
                'error': f"invalid exit_type {exit_type!r}; "
                         f"must be one of {self.EXIT_TYPES}",
            }
        if not isinstance(peer_name, str) or not re.match(r'^[A-Za-z0-9_.-]{1,64}$', peer_name):
            return {'ok': False, 'error': f'invalid peer_name {peer_name!r}'}

        if self.peer_registry is None:
            return {'ok': False, 'error': 'peer_registry not available'}

        try:
            ok = self.peer_registry.set_peer_exit_via(peer_name, exit_type)
        except Exception as e:
            logger.error(f"set_peer_exit: registry update failed: {e}")
            return {'ok': False, 'error': str(e)}
        if not ok:
            return {'ok': False, 'error': f'peer {peer_name!r} not found'}

        try:
            self.apply_routes()
        except Exception as e:
            logger.warning(f"set_peer_exit: apply_routes failed (non-fatal): {e}")

        return {'ok': True, 'peer': peer_name, 'exit_via': exit_type}

    def upload_wireguard_ext(self, conf_text: str) -> Dict[str, Any]:
        """Validate and store an external WireGuard config."""
        try:
            cleaned = self._validate_wg_conf(conf_text)
        except ValueError as e:
            return {'ok': False, 'error': str(e)}

        path = os.path.join(self.wireguard_ext_dir, 'wg_ext0.conf')
        try:
            self._write_secure(path, cleaned)
        except Exception as e:
            logger.error(f"upload_wireguard_ext: write failed: {e}")
            return {'ok': False, 'error': str(e)}

        logger.info(f"connectivity: stored wg_ext0.conf ({len(cleaned)} bytes)")
        return {'ok': True}

    def upload_openvpn(self, ovpn_text: str, name: str = 'default') -> Dict[str, Any]:
        """Validate and store an OpenVPN profile."""
        if not isinstance(name, str) or not _NAME_RE.match(name):
            return {
                'ok': False,
                'error': f'invalid name {name!r}; must match [a-z0-9_-]{{1,32}}',
            }
        try:
            cleaned = self._validate_ovpn(ovpn_text)
        except ValueError as e:
            return {'ok': False, 'error': str(e)}

        path = os.path.join(self.openvpn_dir, f'{name}.ovpn')
        try:
            self._write_secure(path, cleaned)
        except Exception as e:
            logger.error(f"upload_openvpn: write failed: {e}")
            return {'ok': False, 'error': str(e)}

        logger.info(f"connectivity: stored {name}.ovpn ({len(cleaned)} bytes)")
        return {'ok': True}

    # ── Routing application ───────────────────────────────────────────────

    def apply_routes(self) -> Dict[str, Any]:
        """Idempotently rebuild all connectivity rules and policy routing."""
        rules_applied = 0

        try:
            self._ensure_chains()
        except Exception as e:
            logger.warning(f"apply_routes: _ensure_chains failed: {e}")

        # Flush our dedicated chains (without deleting them)
        for table, chain in (('mangle', self.CONNECTIVITY_CHAIN),
                             ('nat', self.CONNECTIVITY_CHAIN)):
            try:
                self._flush_chain(table, chain)
            except Exception as e:
                logger.warning(f"apply_routes: flush {table}/{chain} failed: {e}")

        # Idempotent ip rule registration for each non-default exit
        for exit_type in ('wireguard_ext', 'openvpn', 'tor'):
            mark = self.MARKS[exit_type]
            table = self.TABLES[exit_type]
            try:
                self._remove_ip_rule(mark, table)
                self._add_ip_rule(mark, table)
                rules_applied += 1
            except Exception as e:
                logger.warning(f"apply_routes: ip rule {exit_type} failed: {e}")

        # Per-peer marking + nat redirect (Tor only)
        if self.peer_registry is not None:
            try:
                peers = self.peer_registry.list_peers()
            except Exception as e:
                logger.warning(f"apply_routes: list_peers failed: {e}")
                peers = []

            for peer in peers:
                exit_via = peer.get('exit_via', 'default')
                if exit_via == 'default' or exit_via not in self.MARKS:
                    continue
                src_ip = self._peer_source_ip(peer.get('peer', ''))
                if not src_ip:
                    continue
                mark = self.MARKS[exit_via]
                try:
                    self._add_mark_rule(src_ip, mark)
                    rules_applied += 1
                except Exception as e:
                    logger.warning(
                        f"apply_routes: mark rule for {src_ip}/{exit_via}: {e}"
                    )

                # Tor: redirect TCP to local transparent proxy
                if exit_via == 'tor':
                    try:
                        self._add_tor_redirect(src_ip)
                        rules_applied += 1
                    except Exception as e:
                        logger.warning(
                            f"apply_routes: tor redirect for {src_ip}: {e}"
                        )

        # Kill-switch: drop marked packets that would otherwise leak via the
        # default route if the exit interface is down.
        for exit_type, iface in self.IFACES.items():
            mark = self.MARKS[exit_type]
            try:
                self._add_killswitch(mark, iface)
                rules_applied += 1
            except Exception as e:
                logger.warning(f"apply_routes: killswitch {exit_type}: {e}")

        return {'ok': True, 'rules_applied': rules_applied}

    # ── iptables / ip rule helpers ────────────────────────────────────────

    def _wg_iptables(self, args: List[str], timeout: int = 10) -> subprocess.CompletedProcess:
        """Run iptables inside the WireGuard container (where peer traffic forwards)."""
        cmd = ['docker', 'exec', WIREGUARD_CONTAINER, 'iptables'] + args
        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)

    def _wg_ip(self, args: List[str], timeout: int = 10) -> subprocess.CompletedProcess:
        """Run `ip` inside the WireGuard container."""
        cmd = ['docker', 'exec', WIREGUARD_CONTAINER, 'ip'] + args
        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)

    def _ensure_chains(self) -> None:
        """Create PIC_CONNECTIVITY chains in mangle and nat (idempotent)."""
        for table, parent_chain in (
            ('mangle', 'PREROUTING'),
            ('nat', 'PREROUTING'),
        ):
            # Create chain if it doesn't already exist
            check = self._wg_iptables(
                ['-t', table, '-L', self.CONNECTIVITY_CHAIN, '-n']
            )
            if check.returncode != 0:
                create = self._wg_iptables(
                    ['-t', table, '-N', self.CONNECTIVITY_CHAIN]
                )
                if create.returncode != 0 and 'exists' not in (create.stderr or ''):
                    logger.warning(
                        f"_ensure_chains: cannot create {table}/{self.CONNECTIVITY_CHAIN}: "
                        f"{create.stderr.strip()}"
                    )

            # Insert jump from parent chain at position 1, idempotent.
            jump_args = ['-t', table, '-C', parent_chain, '-j', self.CONNECTIVITY_CHAIN]
            exists = self._wg_iptables(jump_args)
            if exists.returncode != 0:
                self._wg_iptables(
                    ['-t', table, '-I', parent_chain, '1',
                     '-j', self.CONNECTIVITY_CHAIN]
                )

    def _flush_chain(self, table: str, chain: str) -> None:
        """Flush a chain in-place (`iptables -F`) without deleting it."""
        self._wg_iptables(['-t', table, '-F', chain])

    def _add_ip_rule(self, mark: int, table: int) -> None:
        """Add `ip rule fwmark <mark> lookup <table>`."""
        self._wg_ip(['rule', 'add', 'fwmark', hex(mark), 'lookup', str(table)])

    def _remove_ip_rule(self, mark: int, table: int) -> None:
        """Remove all matching `ip rule fwmark <mark> lookup <table>` (idempotent)."""
        # `ip rule del` returns nonzero when no matching rule exists; loop
        # until it fails to drain duplicates.
        for _ in range(8):
            r = self._wg_ip(['rule', 'del', 'fwmark', hex(mark), 'lookup', str(table)])
            if r.returncode != 0:
                break

    def _add_mark_rule(self, src_ip: str, mark: int) -> None:
        """Mark packets from src_ip with mark in the mangle PIC_CONNECTIVITY chain."""
        self._wg_iptables([
            '-t', 'mangle', '-A', self.CONNECTIVITY_CHAIN,
            '-s', src_ip,
            '-j', 'MARK', '--set-mark', hex(mark),
        ])

    def _add_tor_redirect(self, src_ip: str) -> None:
        """Redirect peer's TCP traffic to local Tor TransPort."""
        self._wg_iptables([
            '-t', 'nat', '-A', self.CONNECTIVITY_CHAIN,
            '-s', src_ip, '-p', 'tcp',
            '-j', 'REDIRECT', '--to-ports', str(self.TOR_TRANS_PORT),
        ])

    def _add_killswitch(self, mark: int, iface: Optional[str]) -> None:
        """Drop marked packets that would egress via any interface other than iface.

        For Tor (no exit iface), skip — Tor traffic is fully redirected at
        nat/REDIRECT and never reaches FORWARD.
        """
        if not iface:
            return
        # Use -C to test, -A to add — idempotent.
        check_args = ['-C', 'FORWARD',
                      '-m', 'mark', '--mark', hex(mark),
                      '!', '-o', iface, '-j', 'DROP']
        exists = self._wg_iptables(check_args)
        if exists.returncode != 0:
            self._wg_iptables(['-A', 'FORWARD',
                               '-m', 'mark', '--mark', hex(mark),
                               '!', '-o', iface, '-j', 'DROP'])

    def _exit_status(self, exit_type: str) -> Dict[str, Any]:
        """Return per-exit status (config presence + interface up/down)."""
        info: Dict[str, Any] = {'configured': False, 'iface_up': False}
        if exit_type == 'wireguard_ext':
            path = os.path.join(self.wireguard_ext_dir, 'wg_ext0.conf')
            info['configured'] = os.path.isfile(path)
        elif exit_type == 'openvpn':
            try:
                info['configured'] = any(
                    f.endswith('.ovpn')
                    for f in os.listdir(self.openvpn_dir)
                )
            except OSError:
                info['configured'] = False
        elif exit_type == 'tor':
            info['configured'] = True  # Tor uses defaults; no per-cell config

        iface = self.IFACES.get(exit_type)
        if iface:
            try:
                r = self._wg_ip(['link', 'show', iface], timeout=5)
                info['iface_up'] = r.returncode == 0 and 'UP' in (r.stdout or '')
            except Exception:
                info['iface_up'] = False
        return info

    def _peer_source_ip(self, peer_name: str) -> Optional[str]:
        """Return a peer's WireGuard IP (no /CIDR suffix)."""
        if not peer_name or self.peer_registry is None:
            return None
        try:
            peer = self.peer_registry.get_peer(peer_name)
        except Exception as e:
            logger.warning(f"_peer_source_ip({peer_name}): {e}")
            return None
        if not peer:
            return None
        ip = peer.get('ip', '')
        if not ip:
            return None
        return ip.split('/')[0]

    # ── Config validation ─────────────────────────────────────────────────

    def _validate_wg_conf(self, text: str) -> str:
        """Strip Pre/Post-Up/Down hooks and reject conflicting wg0 interface.

        Raises ValueError if the config tries to define `Interface = wg0`
        (which would clash with the existing peer-server interface).
        """
        if not isinstance(text, str):
            raise ValueError('wg conf must be a string')

        cleaned: List[str] = []
        for raw_line in text.splitlines():
            stripped = raw_line.strip()

            # Reject wg0 interface declaration that would conflict with the
            # existing WireGuard server interface.
            if stripped.lower().startswith('interface'):
                # Look ahead in subsequent lines for `= wg0` would be hard;
                # the [Interface] section header itself is fine.  We only
                # reject explicit Name/Interface = wg0 directives.
                pass
            # Match assignments like `PostUp = ...`
            if '=' in stripped:
                key = stripped.split('=', 1)[0].strip()
                if key in _WG_FORBIDDEN_PREFIXES:
                    logger.info(f"_validate_wg_conf: dropped {key} hook")
                    continue
                # Detect Name = wg0 or Interface = wg0 inside Interface section
                if key.lower() in ('name', 'interface') and \
                        stripped.split('=', 1)[1].strip().lower() == 'wg0':
                    raise ValueError(
                        "config defines interface 'wg0' which conflicts "
                        "with the peer-server interface"
                    )
            cleaned.append(raw_line)
        return '\n'.join(cleaned).rstrip() + '\n'

    def _validate_ovpn(self, text: str) -> str:
        """Strip directives that execute external scripts/binaries."""
        if not isinstance(text, str):
            raise ValueError('ovpn conf must be a string')

        cleaned: List[str] = []
        for raw_line in text.splitlines():
            stripped = raw_line.strip()
            # Match the directive name (first whitespace-delimited token).
            if stripped and not stripped.startswith('#'):
                first = stripped.split(None, 1)[0]
                if first in _OVPN_FORBIDDEN_DIRECTIVES:
                    logger.info(f"_validate_ovpn: dropped {first} directive")
                    continue
            cleaned.append(raw_line)
        return '\n'.join(cleaned).rstrip() + '\n'

    # ── Filesystem helpers ────────────────────────────────────────────────

    @staticmethod
    def _write_secure(path: str, text: str) -> None:
        """Atomic 0o600 write — secrets in these configs must not be world-readable."""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp = path + '.tmp'
        fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        try:
            with os.fdopen(fd, 'w') as f:
                f.write(text)
        except Exception:
            try:
                os.unlink(tmp)
            except OSError:
                pass
            raise
        os.chmod(tmp, 0o600)
        os.replace(tmp, path)
        os.chmod(path, 0o600)