pic/tests/test_connectivity_cell_relay.py

"""
Phase 7 tests for ConnectivityManager — cell_relay connection type.

A cell_relay connection unifies the cell-to-cell exit-relay concept into the
connection model: a connected cell that OFFERS its internet surfaces as an
assignable connection, and peers route through it via the existing cell-exit
mechanism (wireguard_manager.apply_peer_route_via + firewall cell rules) rather
than a local exit container.

Covers:
  - reconcile derives one cell_relay per offering cell link (no dups) and
    removes it when the offer is withdrawn (unless still referenced)
  - a cell_relay allocates mark+table but NO iface/redirect_port/container
  - assigning a peer to a cell_relay drives the cell-routing path (route_via +
    apply_peer_route_via + cell firewall rules), NOT a local container/redirect
  - loop detection rejects a cycle (A→B→A)
  - health derives from the cell tunnel handshake + offer

All subprocess/docker/iptables access is mocked: the cell_link_manager,
wireguard_manager and firewall_manager are MagicMocks / patched modules, and
ConnectivityManager's own iptables/ip-rule helpers are patched out.
"""

import os
import sys
import shutil
import tempfile
import unittest
from unittest.mock import MagicMock, patch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'api'))

from config_manager import ConfigManager
from connectivity_manager import ConnectivityManager


def _link(cell_name='peercell', **overrides):
    link = {
        'cell_name': cell_name,
        'public_key': 'PUBKEY_' + cell_name,
        'vpn_subnet': '10.9.0.0/24',
        'dns_ip': '10.9.0.1',
        'domain': f'{cell_name}.example',
        'permissions': {'inbound': {}, 'outbound': {}},
        'exit_offered': False,
        'remote_exit_offered': True,
        'exit_relay_active': False,
        'remote_exit_relay_active': False,
    }
    link.update(overrides)
    return link


class _Base(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.cfg_file = os.path.join(self.tmp, 'cell_config.json')
        self.data_dir = os.path.join(self.tmp, 'data')
        os.makedirs(self.data_dir, exist_ok=True)
        self.cm = ConfigManager(self.cfg_file, self.data_dir)

        self.peer_registry = MagicMock()
        self.peer_registry.list_peers.return_value = []

        self.cell_link = MagicMock()
        self.cell_link.list_connections.return_value = []

        self.wg = MagicMock()
        self.wg.apply_peer_route_via.return_value = True

        with patch.object(ConnectivityManager, '_subscribe_to_events',
                          lambda self: None):
            self.mgr = ConnectivityManager(
                config_manager=self.cm,
                peer_registry=self.peer_registry,
                vault_manager=MagicMock(),
                data_dir=self.data_dir,
                config_dir=self.tmp,
            )
        self.mgr.cell_link_manager = self.cell_link
        self.mgr.wireguard_manager = self.wg
        # Empty v2 slate (no auto-migrated Tor).
        self.cm.configs['connectivity'] = {'version': 2, 'connections': []}
        self.cm._save_all_configs()

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def _raw_relays(self):
        return [c for c in self.cm.list_connections()
                if c.get('type') == 'cell_relay']


# ---------------------------------------------------------------------------
# Reconcile from cell links
# ---------------------------------------------------------------------------

class TestReconcile(_Base):

    def test_creates_relay_for_offering_link(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        res = self.mgr.reconcile_cell_relays()
        self.assertEqual(len(res['created']), 1)
        relays = self._raw_relays()
        self.assertEqual(len(relays), 1)
        self.assertEqual(relays[0]['cell_name'], 'alpha')
        self.assertEqual(relays[0]['name'], 'Cell: alpha')

    def test_no_relay_when_no_offer(self):
        self.cell_link.list_connections.return_value = [
            _link('alpha', remote_exit_offered=False, exit_offered=False)]
        self.mgr.reconcile_cell_relays()
        self.assertEqual(len(self._raw_relays()), 0)

    def test_local_exit_offered_also_surfaces(self):
        self.cell_link.list_connections.return_value = [
            _link('alpha', remote_exit_offered=False, exit_offered=True)]
        self.mgr.reconcile_cell_relays()
        self.assertEqual(len(self._raw_relays()), 1)

    def test_idempotent_no_duplicates(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        r2 = self.mgr.reconcile_cell_relays()
        self.assertEqual(r2['created'], [])
        self.assertEqual(len(self._raw_relays()), 1)

    def test_removes_relay_when_offer_withdrawn(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        self.assertEqual(len(self._raw_relays()), 1)
        # Offer withdrawn.
        self.cell_link.list_connections.return_value = [
            _link('alpha', remote_exit_offered=False)]
        res = self.mgr.reconcile_cell_relays()
        self.assertEqual(len(res['removed']), 1)
        self.assertEqual(len(self._raw_relays()), 0)

    def test_removes_relay_when_link_gone(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        self.cell_link.list_connections.return_value = []
        self.mgr.reconcile_cell_relays()
        self.assertEqual(len(self._raw_relays()), 0)

    def test_kept_when_offer_gone_but_still_referenced(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        relay_id = self._raw_relays()[0]['id']
        # A peer references this relay → must not be auto-removed.
        self.peer_registry.list_peers.return_value = [
            {'peer': 'laptop', 'exit_via': relay_id}]
        self.cell_link.list_connections.return_value = []
        self.mgr.reconcile_cell_relays()
        self.assertEqual(len(self._raw_relays()), 1)

    def test_list_connections_reconciles(self):
        # list_connections runs reconcile, so the relay appears without an
        # explicit reconcile call.
        self.cell_link.list_connections.return_value = [_link('alpha')]
        listed = self.mgr.list_connections()
        relays = [c for c in listed if c['type'] == 'cell_relay']
        self.assertEqual(len(relays), 1)


# ---------------------------------------------------------------------------
# Resource allocation
# ---------------------------------------------------------------------------

class TestAllocation(_Base):

    def test_relay_has_mark_table_no_iface_port(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        relay = self._raw_relays()[0]
        self.assertIsInstance(relay['mark'], int)
        self.assertIsInstance(relay['table'], int)
        self.assertIsNone(relay['iface'])
        self.assertIsNone(relay['redirect_port'])

    def test_relay_has_no_container(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        relay = self._raw_relays()[0]
        self.assertIsNone(self.mgr.instance_container_name(relay))

    def test_manual_create_rejected(self):
        res = self.mgr.create_connection('cell_relay', 'Cell: alpha')
        self.assertFalse(res['ok'])
        self.assertIn('invalid type', res['error'])


# ---------------------------------------------------------------------------
# Routing drives the cell path, not a local container
# ---------------------------------------------------------------------------

class TestRouting(_Base):

    def _setup_assigned_peer(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        relay = self._raw_relays()[0]
        self.peer = {'peer': 'laptop', 'ip': '10.0.0.5/32',
                     'exit_via': relay['id'], 'route_via': None}
        self.peer_registry.list_peers.return_value = [self.peer]
        self.peer_registry.get_peer.return_value = self.peer
        return relay

    def test_apply_routes_drives_cell_path(self):
        self._setup_assigned_peer()
        with patch('firewall_manager.apply_cell_rules') as fw, \
                patch.object(self.mgr, '_ensure_chains'), \
                patch.object(self.mgr, '_flush_chain'), \
                patch.object(self.mgr, '_add_mark_rule') as add_mark, \
                patch.object(self.mgr, '_add_redirect') as add_redirect, \
                patch.object(self.mgr, '_add_ip_rule'), \
                patch.object(self.mgr, '_remove_ip_rule'):
            self.mgr.apply_routes()

        # Cell path was driven.
        self.wg.apply_peer_route_via.assert_called_once()
        args, kwargs = self.wg.apply_peer_route_via.call_args
        self.assertEqual(args[0], '10.0.0.5')
        self.assertEqual(kwargs.get('via_wg_ip'), '10.9.0.1')
        fw.assert_called_once()
        self.assertTrue(fw.call_args.kwargs.get('exit_relay'))
        self.cell_link.set_exit_relay_active.assert_called_once_with('alpha', True)
        # route_via kept in sync for startup-replay parity.
        self.peer_registry.set_route_via.assert_called_once_with('laptop', 'alpha')

        # NOT the local exit container / redirect path.
        add_mark.assert_not_called()
        add_redirect.assert_not_called()

    def test_apply_routes_skips_when_offer_withdrawn(self):
        self._setup_assigned_peer()
        # Offer withdrawn after assignment: routing must skip (peer falls back).
        self.cell_link.list_connections.return_value = [
            _link('alpha', remote_exit_offered=False)]
        with patch('firewall_manager.apply_cell_rules') as fw, \
                patch.object(self.mgr, '_ensure_chains'), \
                patch.object(self.mgr, '_flush_chain'), \
                patch.object(self.mgr, '_add_ip_rule'), \
                patch.object(self.mgr, '_remove_ip_rule'):
            self.mgr.apply_routes()
        self.wg.apply_peer_route_via.assert_not_called()
        fw.assert_not_called()


# ---------------------------------------------------------------------------
# Loop detection
# ---------------------------------------------------------------------------

class TestLoopDetection(_Base):

    def test_set_peer_exit_rejects_loop(self):
        # We already act as alpha's exit relay (exit_offered True) → routing a
        # peer THROUGH alpha would close the loop A→B→A.
        self.cell_link.list_connections.return_value = [
            _link('alpha', exit_offered=True)]
        self.mgr.reconcile_cell_relays()
        relay = self._raw_relays()[0]
        peer = {'peer': 'laptop', 'ip': '10.0.0.5/32', 'route_via': None}
        self.peer_registry.get_peer.return_value = peer

        res = self.mgr.set_peer_exit('laptop', relay['id'])
        self.assertFalse(res['ok'])
        self.assertIn('loop', res['error'])
        self.peer_registry.set_peer_exit_via.assert_not_called()

    def test_apply_routes_refuses_loop(self):
        self.cell_link.list_connections.return_value = [
            _link('alpha', exit_offered=True)]
        self.mgr.reconcile_cell_relays()
        relay = self._raw_relays()[0]
        peer = {'peer': 'laptop', 'ip': '10.0.0.5/32',
                'exit_via': relay['id'], 'route_via': None}
        self.peer_registry.list_peers.return_value = [peer]
        self.peer_registry.get_peer.return_value = peer
        with patch('firewall_manager.apply_cell_rules'), \
                patch.object(self.mgr, '_ensure_chains'), \
                patch.object(self.mgr, '_flush_chain'), \
                patch.object(self.mgr, '_add_ip_rule'), \
                patch.object(self.mgr, '_remove_ip_rule'):
            self.mgr.apply_routes()
        self.wg.apply_peer_route_via.assert_not_called()


# ---------------------------------------------------------------------------
# Health derives from cell handshake + offer
# ---------------------------------------------------------------------------

class TestHealth(_Base):

    def _relay(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        return self._raw_relays()[0]

    def test_working_when_online_and_offered(self):
        relay = self._relay()
        self.cell_link.get_connection_status.return_value = {
            **_link('alpha'), 'online': True}
        health, _ = self.mgr.probe_health(relay)
        self.assertEqual(health, 'working')

    def test_down_when_handshake_stale(self):
        relay = self._relay()
        self.cell_link.get_connection_status.return_value = {
            **_link('alpha'), 'online': False}
        health, _ = self.mgr.probe_health(relay)
        self.assertEqual(health, 'down')

    def test_down_when_offer_withdrawn(self):
        relay = self._relay()
        self.cell_link.get_connection_status.return_value = {
            **_link('alpha', remote_exit_offered=False), 'online': True}
        health, _ = self.mgr.probe_health(relay)
        self.assertEqual(health, 'down')


# ---------------------------------------------------------------------------
# Teardown cleanup — regression for the confirmed cell_relay routing leak.
#
# A cell_relay policy-routes a peer with `ip rule from <peer> lookup <table>`
# plus a shared `default via <cell-ip>` route in that table, inside
# cell-wireguard. Before the fix, detaching/deleting the peer left the rule
# (remove_peer_route_via used the wrong default table) and nothing ever flushed
# the table's default route — both leaked, confirmed on hardware.
# ---------------------------------------------------------------------------

class TestTeardownCleanup(_Base):

    def _relay(self):
        self.cell_link.list_connections.return_value = [_link('alpha')]
        self.mgr.reconcile_cell_relays()
        return self._raw_relays()[0]

    def test_detach_removes_peer_ip_rule(self):
        relay = self._relay()
        peer = {'peer': 'laptop', 'ip': '10.0.0.5/32',
                'exit_via': relay['id'], 'route_via': 'alpha'}
        self.peer_registry.get_peer.return_value = peer
        self.peer_registry.set_peer_exit_via.return_value = True
        with patch.object(self.mgr, 'apply_routes'):
            res = self.mgr.set_peer_exit('laptop', 'default')
        self.assertTrue(res['ok'])
        # The peer's source ip rule is cleared by source (table-agnostic), so it
        # matches the relay's allocated table rather than the old default 100.
        self.wg.remove_peer_route_via.assert_called_once_with('10.0.0.5')

    def test_delete_connection_flushes_relay_route_table(self):
        relay = self._relay()
        # Not referenced by any peer (detached) → deletable.
        self.peer_registry.list_peers.return_value = []
        res = self.mgr.delete_connection(relay['id'])
        self.assertTrue(res['ok'])
        self.wg.teardown_route_table.assert_called_once_with(relay['table'])

    def test_reconcile_removal_flushes_relay_route_table(self):
        relay = self._relay()
        table = relay['table']
        # Offer withdrawn and not referenced → reconcile removes the relay and
        # must flush its routing table (this path bypasses delete_connection).
        self.cell_link.list_connections.return_value = [
            _link('alpha', remote_exit_offered=False)]
        self.peer_registry.list_peers.return_value = []
        out = self.mgr.reconcile_cell_relays()
        self.assertIn(relay['id'], out['removed'])
        self.wg.teardown_route_table.assert_called_once_with(table)


if __name__ == '__main__':
    unittest.main()