pic/api/cell_link_manager.py

#!/usr/bin/env python3
"""
CellLinkManager — manages site-to-site connections between PIC cells.

Each connection is stored in data/cell_links.json and manifests as:
  - A WireGuard [Peer] block (AllowedIPs = remote cell's VPN subnet)
  - A CoreDNS forwarding block (remote domain → remote cell's DNS IP)
  - An iptables FORWARD rule set (service-level access control)
"""

import os
import json
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

VALID_SERVICES = ('calendar', 'files', 'mail', 'webdav')

_DEFAULT_PERMISSIONS = {
    'inbound': {s: False for s in VALID_SERVICES},
    'outbound': {s: False for s in VALID_SERVICES},
}


def _default_perms() -> Dict[str, Any]:
    return {
        'inbound':  {s: False for s in VALID_SERVICES},
        'outbound': {s: False for s in VALID_SERVICES},
    }


class CellLinkManager:
    def __init__(self, data_dir: str, config_dir: str, wireguard_manager, network_manager):
        self.data_dir = data_dir
        self.config_dir = config_dir
        self.wireguard_manager = wireguard_manager
        self.network_manager = network_manager
        self.links_file = os.path.join(data_dir, 'cell_links.json')

    # ── Storage ───────────────────────────────────────────────────────────────

    def _load(self) -> List[Dict[str, Any]]:
        if os.path.exists(self.links_file):
            try:
                with open(self.links_file) as f:
                    links = json.load(f)
                # Lazy migration: inject permissions field if missing
                changed = False
                for link in links:
                    if 'permissions' not in link:
                        link['permissions'] = _default_perms()
                        changed = True
                if changed:
                    self._save(links)
                return links
            except Exception:
                return []
        return []

    def _save(self, links: List[Dict[str, Any]]):
        with open(self.links_file, 'w') as f:
            json.dump(links, f, indent=2)

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

    def generate_invite(self, cell_name: str, domain: str) -> Dict[str, Any]:
        """Return an invite package describing this cell for another cell to import."""
        keys = self.wireguard_manager.get_keys()
        srv = self.wireguard_manager.get_server_config()
        server_vpn_ip = self.wireguard_manager._get_configured_address().split('/')[0]
        return {
            'cell_name': cell_name,
            'public_key': keys['public_key'],
            'endpoint': srv.get('endpoint'),
            'vpn_subnet': self.wireguard_manager._get_configured_network(),
            'dns_ip': server_vpn_ip,
            'domain': domain,
            'version': 1,
        }

    def list_connections(self) -> List[Dict[str, Any]]:
        return self._load()

    def add_connection(self, invite: Dict[str, Any],
                       inbound_services: Optional[List[str]] = None) -> Dict[str, Any]:
        """Import a remote cell's invite and establish the connection.

        inbound_services: which of THIS cell's services to share with the remote
                          cell immediately. Defaults to none (all-deny).
        """
        links = self._load()
        name = invite['cell_name']
        if any(l['cell_name'] == name for l in links):
            raise ValueError(f"Cell '{name}' is already connected")

        ok = self.wireguard_manager.add_cell_peer(
            name=name,
            public_key=invite['public_key'],
            endpoint=invite.get('endpoint', ''),
            vpn_subnet=invite['vpn_subnet'],
        )
        if not ok:
            raise RuntimeError(f"Failed to add WireGuard peer for cell '{name}'")

        dns_result = self.network_manager.add_cell_dns_forward(
            domain=invite['domain'],
            dns_ip=invite['dns_ip'],
        )
        if dns_result.get('warnings'):
            logger.warning('DNS forward warnings for %s: %s', name, dns_result['warnings'])

        inbound = [s for s in (inbound_services or []) if s in VALID_SERVICES]
        perms = _default_perms()
        for s in inbound:
            perms['inbound'][s] = True

        link = {
            'cell_name': name,
            'public_key': invite['public_key'],
            'endpoint': invite.get('endpoint'),
            'vpn_subnet': invite['vpn_subnet'],
            'dns_ip': invite['dns_ip'],
            'domain': invite['domain'],
            'connected_at': datetime.utcnow().isoformat(),
            'permissions': perms,
        }
        links.append(link)
        self._save(links)

        # Apply iptables rules for the new cell (non-fatal if it fails)
        try:
            import firewall_manager as _fm
            _fm.apply_cell_rules(name, invite['vpn_subnet'], inbound)
        except Exception as e:
            logger.warning(f"apply_cell_rules for {name} failed (non-fatal): {e}")

        return link

    def remove_connection(self, cell_name: str):
        """Tear down a cell connection by name."""
        links = self._load()
        link = next((l for l in links if l['cell_name'] == cell_name), None)
        if not link:
            raise ValueError(f"Cell '{cell_name}' not found")

        # Clear firewall rules first (non-fatal)
        try:
            import firewall_manager as _fm
            _fm.clear_cell_rules(cell_name)
        except Exception as e:
            logger.warning(f"clear_cell_rules for {cell_name} failed (non-fatal): {e}")

        self.wireguard_manager.remove_peer(link['public_key'])
        self.network_manager.remove_cell_dns_forward(link['domain'])

        links = [l for l in links if l['cell_name'] != cell_name]
        self._save(links)

    def update_permissions(self, cell_name: str,
                           inbound: Dict[str, bool],
                           outbound: Dict[str, bool]) -> Dict[str, Any]:
        """Update service sharing permissions for a cell connection.

        Validates service names, persists, and re-applies iptables rules.
        Returns the updated link record.
        """
        links = self._load()
        link = next((l for l in links if l['cell_name'] == cell_name), None)
        if not link:
            raise ValueError(f"Cell '{cell_name}' not found")

        # Validate and normalise — only known services, boolean values
        clean_inbound = {s: bool(inbound.get(s, False)) for s in VALID_SERVICES}
        clean_outbound = {s: bool(outbound.get(s, False)) for s in VALID_SERVICES}
        link['permissions'] = {'inbound': clean_inbound, 'outbound': clean_outbound}
        self._save(links)

        # Re-apply firewall rules
        inbound_list = [s for s, v in clean_inbound.items() if v]
        try:
            import firewall_manager as _fm
            _fm.apply_cell_rules(cell_name, link['vpn_subnet'], inbound_list)
        except Exception as e:
            logger.warning(f"apply_cell_rules for {cell_name} failed (non-fatal): {e}")

        return link

    def get_permissions(self, cell_name: str) -> Dict[str, Any]:
        """Return the permissions dict for a connected cell."""
        links = self._load()
        link = next((l for l in links if l['cell_name'] == cell_name), None)
        if not link:
            raise ValueError(f"Cell '{cell_name}' not found")
        return link.get('permissions', _default_perms())

    def get_connection_status(self, cell_name: str) -> Dict[str, Any]:
        """Return link record enriched with live WireGuard handshake status."""
        links = self._load()
        link = next((l for l in links if l['cell_name'] == cell_name), None)
        if not link:
            raise ValueError(f"Cell '{cell_name}' not found")
        try:
            st = self.wireguard_manager.get_peer_status(link['public_key'])
            return {**link, 'online': st.get('online', False),
                    'last_handshake': st.get('last_handshake')}
        except Exception:
            return {**link, 'online': False, 'last_handshake': None}