pic/api/firewall_manager.py

#!/usr/bin/env python3
"""
Firewall Manager for Personal Internet Cell
Manages per-peer iptables rules in the WireGuard container and DNS ACLs in CoreDNS.
"""

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

logger = logging.getLogger(__name__)

_forward_stateful_lock = threading.Lock()

# Virtual IPs assigned to Caddy per service — must match Caddyfile listeners.
# Populated at import time from the default subnet; call update_service_ips()
# whenever ip_range changes so all downstream callers see the new values.
SERVICE_IPS: Dict[str, str] = {
    'calendar': '172.20.0.21',
    'files':    '172.20.0.22',
    'mail':     '172.20.0.23',
    'webdav':   '172.20.0.24',
}


def update_service_ips(ip_range: str) -> None:
    """Recalculate SERVICE_IPS from the new subnet and update in-place."""
    from ip_utils import get_virtual_ips
    new_ips = get_virtual_ips(ip_range)
    SERVICE_IPS.clear()
    SERVICE_IPS.update(new_ips)

# Internal RFC-1918 ranges (peer traffic stays inside these = cell-only access)
PRIVATE_NETS = ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16']

WIREGUARD_CONTAINER = 'cell-wireguard'
CADDY_CONTAINER = 'cell-caddy'
COREFILE_PATH = '/app/config/dns/Corefile'
ZONE_DATA_DIR = '/data'  # inside CoreDNS container; mounted from ./data/dns


def _run(cmd: List[str], check: bool = True) -> subprocess.CompletedProcess:
    """Run a shell command and return the result."""
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
        if check and result.returncode != 0:
            logger.warning(f"Command {cmd} exited {result.returncode}: {result.stderr.strip()}")
        return result
    except Exception as e:
        logger.error(f"Command {cmd} failed: {e}")
        raise


def _wg_exec(args: List[str]) -> subprocess.CompletedProcess:
    """Run a command inside the WireGuard container via docker exec."""
    return _run(['docker', 'exec', WIREGUARD_CONTAINER] + args, check=False)


def _caddy_exec(args: List[str]) -> subprocess.CompletedProcess:
    """Run a command inside the Caddy container via docker exec."""
    return _run(['docker', 'exec', CADDY_CONTAINER] + args, check=False)


# ---------------------------------------------------------------------------
# Virtual IP management (Caddy container)
# ---------------------------------------------------------------------------

def ensure_caddy_virtual_ips() -> bool:
    """Add per-service virtual IPs to Caddy's eth0 if not already present."""
    try:
        result = _caddy_exec(['ip', 'addr', 'show', 'eth0'])
        existing = result.stdout

        for service, ip in SERVICE_IPS.items():
            if ip not in existing:
                r = _caddy_exec(['ip', 'addr', 'add', f'{ip}/16', 'dev', 'eth0'])
                if r.returncode == 0:
                    logger.info(f"Added virtual IP {ip} for {service} to Caddy eth0")
                else:
                    logger.warning(f"Failed to add virtual IP {ip}: {r.stderr.strip()}")
        return True
    except Exception as e:
        logger.error(f"ensure_caddy_virtual_ips failed: {e}")
        return False


# ---------------------------------------------------------------------------
# iptables rule helpers
# ---------------------------------------------------------------------------

def _iptables(args: List[str], check: bool = False) -> subprocess.CompletedProcess:
    return _wg_exec(['iptables'] + args)


def _rule_exists(chain: str, rule_args: List[str]) -> bool:
    result = _iptables(['-C', chain] + rule_args)
    return result.returncode == 0


def _ensure_rule(chain: str, rule_args: List[str]) -> None:
    """Insert rule at top of chain if it doesn't already exist."""
    if not _rule_exists(chain, rule_args):
        _iptables(['-I', chain] + rule_args)


def _delete_rule(chain: str, rule_args: List[str]) -> None:
    """Delete rule from chain (silently if it doesn't exist)."""
    while _rule_exists(chain, rule_args):
        _iptables(['-D', chain] + rule_args)


# ---------------------------------------------------------------------------
# Per-peer rule management
# ---------------------------------------------------------------------------

def _peer_comment(peer_ip: str) -> str:
    # SECURITY: append a non-numeric, non-dash suffix so peer comments cannot
    # be substrings of one another. Without this, the comment for 10.0.0.1
    # ('pic-peer-10-0-0-1') is a prefix of 10.0.0.10..19 and a naive
    # substring match would delete unrelated peers' rules.
    return f'pic-peer-{peer_ip.replace(".", "-")}/32'


def clear_peer_rules(peer_ip: str) -> None:
    """Remove all FORWARD rules tagged with this peer's IP via iptables-save/restore."""
    comment = _peer_comment(peer_ip)
    # SECURITY: match the comment as a complete --comment token, not a
    # substring. iptables-save renders comments as `--comment "<value>"` (or
    # occasionally without quotes), so we anchor on the surrounding quotes /
    # whitespace. Even with the unique /32 suffix in _peer_comment, we keep
    # exact-token matching so a future change to the comment format cannot
    # silently re-introduce the substring-deletion bug.
    comment_re = re.compile(
        rf'--comment\s+["\']?{re.escape(comment)}["\']?(\s|$)'
    )
    try:
        # Dump rules, strip matching lines, restore — atomic and order-stable
        save = _wg_exec(['iptables-save'])
        if save.returncode != 0:
            return
        lines = save.stdout.splitlines()
        filtered = [l for l in lines if not comment_re.search(l)]
        if len(filtered) == len(lines):
            return  # nothing to remove
        restore_input = '\n'.join(filtered) + '\n'
        restore = subprocess.run(
            ['docker', 'exec', '-i', WIREGUARD_CONTAINER, 'iptables-restore'],
            input=restore_input, capture_output=True, text=True, timeout=10
        )
        if restore.returncode != 0:
            logger.warning(f"iptables-restore failed: {restore.stderr.strip()}")
    except Exception as e:
        logger.error(f"clear_peer_rules({peer_ip}): {e}")


def apply_peer_rules(peer_ip: str, settings: Dict[str, Any],
                     wg_subnet: str = '10.0.0.0/24',
                     cell_subnets: Optional[List[str]] = None) -> bool:
    """
    Apply iptables FORWARD rules for a peer based on their access settings.

    wg_subnet: the local cell's WireGuard VPN subnet (e.g. '10.0.2.0/24').
               Used for the peer-to-peer ACCEPT/DROP rule. Defaults to the
               legacy hardcoded value so callers that don't yet pass it are safe.
    cell_subnets: list of connected cells' vpn_subnet strings. When provided,
                  explicit ACCEPT rules are added so split-tunnel peers can reach
                  connected cells regardless of the internet_access setting.

    Each rule is inserted at position 1 (-I), so the LAST call ends up at the TOP.
    We insert in reverse-priority order: lowest-priority rules first, highest last.

    Desired final chain order (top = highest priority):
      1. Connected-cell subnet ACCEPTs (explicit cross-cell routing)
      2. Peer-to-peer ACCEPT/DROP      (local VPN subnet)
      3. Private-net ACCEPTs           (for no-internet peers to reach local resources)
      4. Internet DROP or ACCEPT       (lowest priority catch-all)
    """
    try:
        comment = _peer_comment(peer_ip)
        clear_peer_rules(peer_ip)

        internet_access = settings.get('internet_access', True)
        service_access  = settings.get('service_access', list(SERVICE_IPS.keys()))
        peer_access     = settings.get('peer_access', True)

        # --- Step 1 (inserted first → ends up at bottom before default ACCEPT) ---
        # Internet catch-all: allow or block
        if internet_access:
            _iptables(['-I', 'FORWARD', '-s', peer_ip,
                       '-m', 'comment', '--comment', comment, '-j', 'ACCEPT'])
        else:
            # Block non-private, allow private nets
            _iptables(['-I', 'FORWARD', '-s', peer_ip,
                       '-m', 'comment', '--comment', comment, '-j', 'DROP'])
            for net in reversed(PRIVATE_NETS):
                _iptables(['-I', 'FORWARD', '-s', peer_ip, '-d', net,
                           '-m', 'comment', '--comment', comment, '-j', 'ACCEPT'])

        # --- Step 2 --- Peer-to-peer: use the actual local VPN subnet
        target = 'ACCEPT' if peer_access else 'DROP'
        _iptables(['-I', 'FORWARD', '-s', peer_ip, '-d', wg_subnet,
                   '-m', 'comment', '--comment', comment, '-j', target])

        # --- Step 3 --- Explicit ACCEPT for each connected cell's subnet so
        # split-tunnel peers can route to connected cells (wg0 → wg0 forwarding).
        if cell_subnets:
            for subnet in reversed(cell_subnets):
                _iptables(['-I', 'FORWARD', '-s', peer_ip, '-d', subnet,
                           '-m', 'comment', '--comment', comment, '-j', 'ACCEPT'])

        # Service access restriction is done entirely by CoreDNS ACL.
        # No per-peer iptables rule for Caddy:80 — blocking it would also
        # prevent the peer from reaching the PIC web UI and API.

        logger.info(f"Applied rules for {peer_ip}: internet={internet_access} "
                    f"services={service_access} peers={peer_access} "
                    f"wg_subnet={wg_subnet} cell_subnets={cell_subnets}")
        return True
    except Exception as e:
        logger.error(f"apply_peer_rules({peer_ip}): {e}")
        return False


def apply_all_peer_rules(peers: List[Dict[str, Any]],
                         wg_subnet: str = '10.0.0.0/24',
                         cell_subnets: Optional[List[str]] = None) -> None:
    """Re-apply rules for all peers (called on startup)."""
    ensure_caddy_virtual_ips()
    for peer in peers:
        ip = peer.get('ip')
        if not ip:
            continue
        apply_peer_rules(ip, {
            'internet_access': peer.get('internet_access', True),
            'service_access':  peer.get('service_access', list(SERVICE_IPS.keys())),
            'peer_access':     peer.get('peer_access', True),
        }, wg_subnet=wg_subnet, cell_subnets=cell_subnets)


def reconcile_stale_peer_rules(peers: List[Dict[str, Any]]) -> int:
    """Remove iptables rules for peer IPs that are no longer in the registry.

    Returns the number of stale IPs cleaned up.
    """
    known_ips = set()
    for peer in peers:
        raw = peer.get('ip', '')
        ip = raw.split('/')[0] if raw else ''
        if ip:
            known_ips.add(ip)

    # Parse pic-peer-* comments from iptables-save to find IPs with live rules
    save = _wg_exec(['iptables-save'])
    if save.returncode != 0:
        return 0

    # Comment format: pic-peer-A-B-C-D/32  (dots replaced with dashes)
    comment_re = re.compile(r'pic-peer-([\d]+-[\d]+-[\d]+-[\d]+)/32')
    stale_ips: set = set()
    for line in save.stdout.splitlines():
        m = comment_re.search(line)
        if m:
            ip = m.group(1).replace('-', '.')
            if ip not in known_ips:
                stale_ips.add(ip)

    for ip in stale_ips:
        logger.warning(f"Removing stale iptables rules for deleted peer {ip}")
        clear_peer_rules(ip)

    if stale_ips:
        logger.info(f"Reconciled {len(stale_ips)} stale peer(s): {sorted(stale_ips)}")
    return len(stale_ips)


# ---------------------------------------------------------------------------
# Cell-to-cell firewall rules
# ---------------------------------------------------------------------------

def _cell_tag(cell_name: str) -> str:
    """iptables comment tag for cell rules — distinct prefix from pic-peer-* to prevent collision."""
    safe = re.sub(r'[^a-z0-9]', '-', cell_name.lower())
    return f'pic-cell-{safe}'


def clear_cell_rules(cell_name: str) -> None:
    """Remove all FORWARD rules tagged for this cell (atomic save/restore)."""
    tag = _cell_tag(cell_name)
    comment_re = re.compile(rf'--comment\s+["\']?{re.escape(tag)}["\']?(\s|$)')
    try:
        save = _wg_exec(['iptables-save'])
        if save.returncode != 0:
            return
        lines = save.stdout.splitlines()
        filtered = [l for l in lines if not comment_re.search(l)]
        if len(filtered) == len(lines):
            return
        restore_input = '\n'.join(filtered) + '\n'
        restore = subprocess.run(
            ['docker', 'exec', '-i', WIREGUARD_CONTAINER, 'iptables-restore'],
            input=restore_input, capture_output=True, text=True, timeout=10
        )
        if restore.returncode != 0:
            logger.warning(f"clear_cell_rules iptables-restore failed: {restore.stderr.strip()}")
    except Exception as e:
        logger.error(f"clear_cell_rules({cell_name}): {e}")


def _get_cell_api_ip() -> Optional[str]:
    """Return cell-api's Docker bridge IP. Returns empty string on failure."""
    r = _run(['docker', 'inspect', '--format',
              '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}',
              'cell-api'], check=False)
    return r.stdout.strip()


def _get_dns_container_ip() -> str:
    """Return cell-dns container's Docker bridge IP. Falls back to 172.20.0.3."""
    try:
        r = _run(['docker', 'inspect', '--format',
                  '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}',
                  'cell-dns'], check=False)
        return r.stdout.strip() or '172.20.0.3'
    except Exception:
        return '172.20.0.3'


def _get_wg_server_ip() -> Optional[str]:
    """Return the WireGuard server's VPN IP from wg0.conf (e.g. '10.0.0.1')."""
    import ipaddress as _ipaddress
    wg_conf_path = '/app/config/wireguard/wg_confs/wg0.conf'
    try:
        with open(wg_conf_path) as f:
            for line in f:
                line = line.strip()
                if line.startswith('Address') and '=' in line:
                    addr = line.split('=', 1)[1].strip()
                    return str(_ipaddress.ip_interface(addr).ip)
    except Exception:
        pass
    return None


def _get_caddy_container_ip() -> str:
    """Return cell-caddy container's Docker bridge IP. Falls back to 172.20.0.2."""
    try:
        r = _run(['docker', 'inspect', '--format',
                  '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}',
                  'cell-caddy'], check=False)
        return r.stdout.strip() or '172.20.0.2'
    except Exception:
        return '172.20.0.2'


def apply_cell_rules(cell_name: str, vpn_subnet: str, inbound_services: List[str],
                     exit_relay: bool = False) -> bool:
    """Apply FORWARD rules for a cell-to-cell peer.

    Traffic from vpn_subnet is allowed only to service VIPs listed in
    inbound_services; all other cell traffic is DROPped.  Cells get no
    internet or peer access — only explicit service access via Caddy on
    port 80, plus the cell-api port (3000) for permission-sync pushes.

    DNS (port 53) is always allowed so cell peers can resolve service names.
    Service names resolve to the WG server IP; ensure_service_dnat() routes
    wg0:80 to Caddy, which routes by Host header.

    When exit_relay=True, the remote cell's peers can route internet
    traffic through this cell (Phase 3).

    Rule insertion order (first inserted = bottom, last inserted = top):
      1. Catch-all DROP for the subnet           (inserted first → bottom)
      2. Exit relay ACCEPT (-o eth0)             (if exit_relay, above catch-all)
      3. Service ACCEPT to Caddy port 80         (if any inbound_services)
      4. DNS ACCEPT to cell-dns port 53          (UDP + TCP)
      5. API-sync ACCEPT                         (inserted last → top)
    """
    try:
        tag = _cell_tag(cell_name)
        clear_cell_rules(cell_name)

        # Catch-all DROP for new connections only — inserted first so it ends up at the bottom.
        # Using state=NEW,INVALID preserves ESTABLISHED/RELATED packets (ICMP replies,
        # TCP ACKs) for connections initiated by local peers to this cell, which would
        # otherwise be dropped before reaching the stateful ACCEPT rule.
        _iptables(['-I', 'FORWARD', '-s', vpn_subnet,
                   '-m', 'state', '--state', 'NEW,INVALID',
                   '-m', 'comment', '--comment', tag, '-j', 'DROP'])

        # Exit relay ACCEPT — allow internet-bound traffic from this cell's peers.
        if exit_relay:
            _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-o', 'eth0',
                       '-m', 'comment', '--comment', tag, '-j', 'ACCEPT'])

        # Service access via Caddy — DNAT wg0:80 → Caddy; Host header routes to service.
        # Only add ACCEPT if this cell has any inbound services granted.
        if inbound_services:
            caddy_ip = _get_caddy_container_ip()
            if caddy_ip:
                _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-d', caddy_ip,
                           '-p', 'tcp', '--dport', '80',
                           '-m', 'comment', '--comment', tag, '-j', 'ACCEPT'])

        # DNS ACCEPT — allow cross-cell peers to query CoreDNS via the WG server IP.
        # ensure_dns_dnat() routes wg0:53 to cell-dns; FORWARD must allow it.
        dns_ip = _get_dns_container_ip()
        if dns_ip:
            for proto in ('udp', 'tcp'):
                _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-d', dns_ip,
                           '-p', proto, '--dport', '53',
                           '-m', 'comment', '--comment', tag, '-j', 'ACCEPT'])

        # API permission-sync ACCEPT — inserted LAST so it goes to position 1 (above
        # the catch-all DROP). Remote cells push permissions to our cell-api via the
        # WG tunnel; iptables sees source=cell_subnet dst=api_ip after DNAT.
        api_ip = _get_cell_api_ip()
        if api_ip:
            _iptables(['-I', 'FORWARD', '-s', vpn_subnet, '-d', api_ip,
                       '-p', 'tcp', '--dport', '3000',
                       '-m', 'comment', '--comment', tag, '-j', 'ACCEPT'])

        # Ensure reply traffic (e.g. ICMP, TCP ACKs) for connections initiated
        # by local peers to this cell is not dropped by the cell's catch-all DROP.
        ensure_forward_stateful()

        logger.info(
            f"Applied cell rules for {cell_name} ({vpn_subnet}): "
            f"inbound={inbound_services} exit_relay={exit_relay}"
        )
        return True
    except Exception as e:
        logger.error(f"apply_cell_rules({cell_name}): {e}")
        return False


def apply_all_cell_rules(cell_links: List[Dict[str, Any]]) -> None:
    """Re-apply firewall rules for all cell connections (called on startup)."""
    for link in cell_links:
        name = link.get('cell_name')
        subnet = link.get('vpn_subnet')
        if not name or not subnet:
            continue
        perms = link.get('permissions', {})
        inbound = [s for s, v in perms.get('inbound', {}).items() if v]
        exit_relay = bool(link.get('remote_exit_relay_active', False))
        apply_cell_rules(name, subnet, inbound, exit_relay=exit_relay)


def ensure_forward_stateful() -> bool:
    """Ensure ESTABLISHED/RELATED ACCEPT is at position 1 (top) of FORWARD.

    Cell rules DROP all traffic from a connected cell's subnet except specific
    service ports. Without conntrack, ICMP replies and TCP ACKs for connections
    initiated BY local peers to the connected cell are also dropped, making
    cross-cell routing (peer → cell → remote cell) broken.

    This function always deletes any existing instance and re-inserts at position 1.
    That re-anchoring is necessary because wg0 PostUp uses -I FORWARD (insert at top),
    which pushes this rule down every time wg0 restarts — causing ICMP to hit the
    per-peer DROP rule before reaching the stateful ACCEPT.
    """
    with _forward_stateful_lock:
        try:
            # Remove all existing instances so we can re-anchor at position 1.
            # PostUp -I FORWARD rules drift this rule down on every wg0 restart.
            while _wg_exec(['iptables', '-D', 'FORWARD', '-m', 'state',
                            '--state', 'ESTABLISHED,RELATED', '-j', 'ACCEPT']).returncode == 0:
                pass
            _wg_exec(['iptables', '-I', 'FORWARD', '1', '-m', 'state',
                      '--state', 'ESTABLISHED,RELATED', '-j', 'ACCEPT'])
            logger.info('ensure_forward_stateful: ESTABLISHED,RELATED anchored at FORWARD position 1')
            return True
        except Exception as e:
            logger.error(f'ensure_forward_stateful: {e}')
            return False


def ensure_cell_api_dnat() -> bool:
    """DNAT wg0:3000 (scoped to WG server IP) → cell-api:3000 inside cell-wireguard.

    Remote cells push permission updates over the WireGuard tunnel to our
    wg0 interface on port 3000. The DNAT is scoped to -d {server_ip} so that
    cross-cell traffic destined for another cell's API (different WG IP) is
    not intercepted. Called on every startup so rules survive container restarts.
    """
    try:
        server_ip = _get_wg_server_ip()
        if not server_ip:
            logger.warning('ensure_cell_api_dnat: could not determine WG server IP')
            return False

        r = _run(['docker', 'inspect', '--format',
                  '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}',
                  'cell-api'], check=False)
        api_ip = r.stdout.strip()
        if not api_ip:
            logger.warning('ensure_cell_api_dnat: cell-api container not found or no IP')
            return False

        dnat_check = ['-t', 'nat', '-C', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                      '-p', 'tcp', '--dport', '3000',
                      '-j', 'DNAT', '--to-destination', f'{api_ip}:3000']
        dnat_add   = ['-t', 'nat', '-A', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                      '-p', 'tcp', '--dport', '3000',
                      '-j', 'DNAT', '--to-destination', f'{api_ip}:3000']
        if _wg_exec(['iptables'] + dnat_check).returncode != 0:
            _wg_exec(['iptables'] + dnat_add)

        masq_check = ['-t', 'nat', '-C', 'POSTROUTING', '-o', 'eth0', '-d', api_ip,
                      '-p', 'tcp', '--dport', '3000', '-j', 'MASQUERADE']
        masq_add   = ['-t', 'nat', '-A', 'POSTROUTING', '-o', 'eth0', '-d', api_ip,
                      '-p', 'tcp', '--dport', '3000', '-j', 'MASQUERADE']
        if _wg_exec(['iptables'] + masq_check).returncode != 0:
            _wg_exec(['iptables'] + masq_add)

        fwd_check = ['-C', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                     '-p', 'tcp', '--dport', '3000', '-j', 'ACCEPT']
        fwd_add   = ['-I', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                     '-p', 'tcp', '--dport', '3000', '-j', 'ACCEPT']
        if _wg_exec(['iptables'] + fwd_check).returncode != 0:
            _wg_exec(['iptables'] + fwd_add)

        logger.info(f'ensure_cell_api_dnat: wg0:3000 → {api_ip}:3000')
        return True
    except Exception as e:
        logger.error(f'ensure_cell_api_dnat: {e}')
        return False


def ensure_dns_dnat() -> bool:
    """DNAT wg0:53 (scoped to WG server IP) → cell-dns:53.

    Peers send DNS queries to the WG server IP. DNAT is scoped with -d {server_ip}
    so cross-cell DNS traffic destined for another cell is forwarded, not hijacked.
    """
    try:
        server_ip = _get_wg_server_ip()
        if not server_ip:
            logger.warning('ensure_dns_dnat: could not determine WG server IP')
            return False
        dns_ip = _get_dns_container_ip()
        if not dns_ip:
            logger.warning('ensure_dns_dnat: cell-dns not found')
            return False
        for proto in ('udp', 'tcp'):
            dnat_check = ['-t', 'nat', '-C', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                          '-p', proto, '--dport', '53',
                          '-j', 'DNAT', '--to-destination', f'{dns_ip}:53']
            dnat_add   = ['-t', 'nat', '-A', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                          '-p', proto, '--dport', '53',
                          '-j', 'DNAT', '--to-destination', f'{dns_ip}:53']
            if _wg_exec(['iptables'] + dnat_check).returncode != 0:
                _wg_exec(['iptables'] + dnat_add)
        for proto in ('udp', 'tcp'):
            fwd_check = ['-C', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                         '-p', proto, '--dport', '53', '-j', 'ACCEPT']
            fwd_add   = ['-I', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                         '-p', proto, '--dport', '53', '-j', 'ACCEPT']
            if _wg_exec(['iptables'] + fwd_check).returncode != 0:
                _wg_exec(['iptables'] + fwd_add)
        logger.info(f'ensure_dns_dnat: wg0:{server_ip}:53 → {dns_ip}:53')
        return True
    except Exception as e:
        logger.error(f'ensure_dns_dnat: {e}')
        return False


def ensure_service_dnat() -> bool:
    """DNAT wg0:80 (scoped to WG server IP) → cell-caddy:80.

    Service DNS names resolve to the WG server IP. DNAT is scoped with -d {server_ip}
    so that cross-cell HTTP traffic destined for another cell passes through unmodified.
    """
    try:
        server_ip = _get_wg_server_ip()
        if not server_ip:
            logger.warning('ensure_service_dnat: could not determine WG server IP')
            return False
        caddy_ip = _get_caddy_container_ip()
        if not caddy_ip:
            logger.warning('ensure_service_dnat: cell-caddy not found')
            return False
        dnat_check = ['-t', 'nat', '-C', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                      '-p', 'tcp', '--dport', '80',
                      '-j', 'DNAT', '--to-destination', f'{caddy_ip}:80']
        dnat_add   = ['-t', 'nat', '-A', 'PREROUTING', '-i', 'wg0', '-d', server_ip,
                      '-p', 'tcp', '--dport', '80',
                      '-j', 'DNAT', '--to-destination', f'{caddy_ip}:80']
        if _wg_exec(['iptables'] + dnat_check).returncode != 0:
            _wg_exec(['iptables'] + dnat_add)
        fwd_check = ['-C', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                     '-p', 'tcp', '--dport', '80', '-j', 'ACCEPT']
        fwd_add   = ['-I', 'FORWARD', '-i', 'wg0', '-o', 'eth0',
                     '-p', 'tcp', '--dport', '80', '-j', 'ACCEPT']
        if _wg_exec(['iptables'] + fwd_check).returncode != 0:
            _wg_exec(['iptables'] + fwd_add)
        logger.info(f'ensure_service_dnat: wg0:{server_ip}:80 → {caddy_ip}:80')
        return True
    except Exception as e:
        logger.error(f'ensure_service_dnat: {e}')
        return False


def ensure_wg_masquerade() -> bool:
    """MASQUERADE Docker bridge traffic leaving via wg0, and allow it through FORWARD.

    cell-dns and other Docker containers need to reach remote cell subnets via
    cell-wireguard's wg0. Without MASQUERADE the source IP (172.20.x.x) can't be
    routed back over the WireGuard tunnel (WireGuard only accepts 10.0.x.x sources
    from peers). MASQUERADE rewrites the source to wg0's IP so replies can return.
    """
    try:
        masq_check = ['-t', 'nat', '-C', 'POSTROUTING', '-o', 'wg0',
                      '-s', '172.20.0.0/16', '-j', 'MASQUERADE']
        masq_add   = ['-t', 'nat', '-A', 'POSTROUTING', '-o', 'wg0',
                      '-s', '172.20.0.0/16', '-j', 'MASQUERADE']
        if _wg_exec(['iptables'] + masq_check).returncode != 0:
            _wg_exec(['iptables'] + masq_add)

        fwd_check = ['-C', 'FORWARD', '-i', 'eth0', '-o', 'wg0',
                     '-s', '172.20.0.0/16', '-j', 'ACCEPT']
        fwd_add   = ['-I', 'FORWARD', '-i', 'eth0', '-o', 'wg0',
                     '-s', '172.20.0.0/16', '-j', 'ACCEPT']
        if _wg_exec(['iptables'] + fwd_check).returncode != 0:
            _wg_exec(['iptables'] + fwd_add)

        logger.info('ensure_wg_masquerade: Docker→wg0 MASQUERADE+FORWARD configured')
        return True
    except Exception as e:
        logger.error(f'ensure_wg_masquerade: {e}')
        return False


def ensure_cell_subnet_routes(cell_links: List[Dict[str, Any]]) -> None:
    """Add host-namespace routes for remote cell VPN subnets via cell-wireguard.

    Docker containers (cell-dns, etc.) use the host's routing table to reach
    non-bridge destinations. Without a route, packets to 10.0.x.0/24 subnets
    of connected cells hit the host's default gateway instead of cell-wireguard.

    Uses a temporary '--network host --rm' container to run ip route replace in
    the host network namespace. cell-api has docker.sock so this works without
    privileged mode or nsenter namespace tricks.
    """
    if not cell_links:
        return
    WG_BRIDGE_IP = '172.20.0.9'  # cell-wireguard's fixed Docker IP (docker-compose.yml)
    for link in cell_links:
        subnet = link.get('vpn_subnet', '')
        if not subnet:
            continue
        try:
            result = _run(
                ['docker', 'run', '--rm',
                 '--network', 'host',
                 '--cap-add', 'NET_ADMIN',
                 'alpine',
                 'ip', 'route', 'replace', subnet, 'via', WG_BRIDGE_IP],
                check=False
            )
            if result.returncode == 0:
                logger.info(f'ensure_cell_subnet_routes: {subnet} via {WG_BRIDGE_IP}')
            else:
                logger.warning(
                    f'ensure_cell_subnet_routes: {subnet} failed: {result.stderr.strip()}'
                )
        except Exception as e:
            logger.warning(f'ensure_cell_subnet_routes: {subnet}: {e}')


# ---------------------------------------------------------------------------
# DNS ACL (CoreDNS Corefile generation)
# ---------------------------------------------------------------------------

# Service subdomains that get per-peer ACL rules in the CoreDNS zone block
_ACL_SERVICES = ('calendar', 'files', 'mail', 'webdav')


def _build_acl_block(blocked_peers_by_service: Dict[str, List[str]],
                     domain: str = 'cell') -> str:
    """
    Build CoreDNS ACL plugin stanzas.

    blocked_peers_by_service: { 'calendar': ['10.0.0.2', '10.0.0.3'], ... }
    Returns a string to embed in the primary zone block.
    """
    if not blocked_peers_by_service:
        return ''

    lines = []
    for service in _ACL_SERVICES:
        peer_ips = blocked_peers_by_service.get(service, [])
        if not peer_ips:
            continue
        host = f'{service}.{domain}.'
        # All blocked IPs for this service in ONE block — separate blocks would
        # cause the first block's allow-all to match before the second block's
        # block rule, silently granting access to all but the first blocked peer.
        lines.append(f'    acl {host} {{')
        for ip in peer_ips:
            lines.append(f'        block net {ip}/32')
        lines.append(f'        allow net 0.0.0.0/0')
        lines.append(f'        allow net ::/0')
        lines.append(f'    }}')
    return '\n'.join(lines)


def generate_corefile(peers: List[Dict[str, Any]], corefile_path: str = COREFILE_PATH,
                      domain: str = 'cell',
                      cell_links: Optional[List[Dict[str, Any]]] = None) -> bool:
    """
    Rewrite the CoreDNS Corefile with per-peer ACL rules and reload plugin.
    The file is written to corefile_path (API-side path mapped into CoreDNS container).
    domain: the configured cell domain (e.g. 'cell', 'dev') — must match zone file names.
    cell_links: optional list of cell-to-cell DNS forwarding entries, each a dict with
                'domain' and 'dns_ip' keys (same shape as CellLinkManager.list_connections()).
                When non-empty, a forwarding stanza is appended for each entry.
    """
    try:
        # Collect which peers block which services
        blocked: Dict[str, List[str]] = {s: [] for s in SERVICE_IPS}
        for peer in peers:
            ip = peer.get('ip')
            if not ip:
                continue
            allowed_services = peer.get('service_access', list(SERVICE_IPS.keys()))
            for service in SERVICE_IPS:
                if service not in allowed_services:
                    blocked[service].append(ip)

        acl_block = _build_acl_block(blocked, domain)

        primary_zone_block = f'{domain} {{\n    file /data/{domain}.zone\n    log\n'
        if acl_block:
            primary_zone_block += acl_block + '\n'
        primary_zone_block += '}\n'

        corefile = f""". {{
    forward . 8.8.8.8 1.1.1.1
    cache
    log
    health
    reload
}}

{primary_zone_block}"""

        # Append cell-to-cell DNS forwarding stanzas if provided
        if cell_links:
            for link in cell_links:
                link_domain = link.get('domain', '')
                link_dns_ip = link.get('dns_ip', '')
                if not link_domain or not link_dns_ip:
                    continue
                corefile += (
                    f'\n{link_domain} {{\n'
                    f'    forward . {link_dns_ip}\n'
                    f'    cache\n'
                    f'    log\n'
                    f'}}\n'
                )
        else:
            corefile += '\n'

        # local.{domain} block intentionally omitted: /data/local.zone does not exist
        # and CoreDNS logs errors on every reload for a missing zone file.
        os.makedirs(os.path.dirname(corefile_path), exist_ok=True)
        tmp_path = corefile_path + '.tmp'
        with open(tmp_path, 'w') as f:
            f.write(corefile)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp_path, corefile_path)

        logger.info(f"Wrote Corefile to {corefile_path}")
        return True
    except Exception as e:
        logger.error(f"generate_corefile: {e}")
        return False


def reload_coredns() -> bool:
    """Signal CoreDNS to reload its config. SIGUSR1 triggers the reload plugin; SIGHUP kills the process."""
    try:
        result = _run(['docker', 'kill', '--signal=SIGUSR1', 'cell-dns'], check=False)
        if result.returncode == 0:
            logger.info("Sent SIGUSR1 to cell-dns (reload)")
            return True
        logger.warning(f"SIGUSR1 to cell-dns failed: {result.stderr.strip()}")
        return False
    except Exception as e:
        logger.error(f"reload_coredns: {e}")
        return False


def apply_all_dns_rules(peers: List[Dict[str, Any]], corefile_path: str = COREFILE_PATH,
                        domain: str = 'cell',
                        cell_links: Optional[List[Dict[str, Any]]] = None) -> bool:
    """Regenerate Corefile (including any cell-to-cell forwarding stanzas) and reload CoreDNS."""
    ok = generate_corefile(peers, corefile_path, domain, cell_links)
    if ok:
        reload_coredns()
    return ok


# ---------------------------------------------------------------------------
# Service store firewall rules
# ---------------------------------------------------------------------------

def _service_tag(service_id: str) -> str:
    safe = re.sub(r'[^a-z0-9]', '-', service_id.lower())
    return f'pic-svc-{safe}'


def apply_service_rules(service_id: str, service_ip: str, rules: list) -> bool:
    """Apply manifest-declared ACCEPT rules for an installed service."""
    tag = _service_tag(service_id)
    clear_service_rules(service_id)
    for r in rules:
        if r.get('type') != 'ACCEPT':
            continue
        dest_ip = r['dest_ip'].replace('${SERVICE_IP}', service_ip)
        dport = str(r['dest_port'])
        proto = r.get('proto', 'tcp')
        _iptables(['-I', 'FORWARD',
                   '-d', dest_ip, '-p', proto, '--dport', dport,
                   '-m', 'comment', '--comment', tag,
                   '-j', 'ACCEPT'])
    return True


def clear_service_rules(service_id: str) -> None:
    """Remove all iptables rules tagged for this service using save/restore."""
    tag = _service_tag(service_id)
    comment_re = re.compile(rf'--comment\s+["\']?{re.escape(tag)}["\']?(\s|$)')
    try:
        save = _wg_exec(['iptables-save'])
        if save.returncode != 0:
            return
        lines = save.stdout.splitlines()
        filtered = [l for l in lines if not comment_re.search(l)]
        if len(filtered) == len(lines):
            return
        restore_input = '\n'.join(filtered) + '\n'
        restore = subprocess.run(
            ['docker', 'exec', '-i', WIREGUARD_CONTAINER, 'iptables-restore'],
            input=restore_input, capture_output=True, text=True, timeout=10
        )
        if restore.returncode != 0:
            logger.warning(f'clear_service_rules iptables-restore failed: {restore.stderr.strip()}')
    except Exception as e:
        logger.error(f'clear_service_rules({service_id}): {e}')