pic/tests/test_cell_link_manager.py

#!/usr/bin/env python3
"""Unit tests for CellLinkManager (cell-to-cell VPN connections)."""

import sys
from pathlib import Path

api_dir = Path(__file__).parent.parent / 'api'
sys.path.insert(0, str(api_dir))

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

from cell_link_manager import CellLinkManager


def _make_wg_mock():
    wg = MagicMock()
    wg.get_keys.return_value = {'public_key': 'serverpubkey=', 'private_key': 'serverprivkey='}
    wg.get_server_config.return_value = {
        'endpoint': '1.2.3.4:51820', 'port': 51820,
        'dns_ip': '10.0.0.3', 'split_tunnel_ips': '10.0.0.0/24, 172.20.0.0/16',
    }
    wg._get_configured_network.return_value = '10.0.0.0/24'
    wg._get_configured_address.return_value = '10.0.0.1/24'
    wg.add_cell_peer.return_value = True
    wg.remove_peer.return_value = True
    return wg


def _make_nm_mock():
    nm = MagicMock()
    nm.add_cell_dns_forward.return_value = {'restarted': ['cell-dns (reloaded)'], 'warnings': []}
    nm.remove_cell_dns_forward.return_value = {'restarted': ['cell-dns (reloaded)'], 'warnings': []}
    return nm


SAMPLE_INVITE = {
    'cell_name': 'office',
    'public_key': 'officepubkey=',
    'endpoint': '5.6.7.8:51820',
    'vpn_subnet': '10.1.0.0/24',
    'dns_ip': '10.1.0.1',
    'domain': 'office.cell',
    'version': 1,
}


class TestCellLinkManagerInvite(unittest.TestCase):

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()
        self.mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_generate_invite_has_required_fields(self):
        invite = self.mgr.generate_invite('mycell', 'home.cell')
        for field in ('cell_name', 'public_key', 'endpoint', 'vpn_subnet', 'dns_ip', 'domain', 'version'):
            self.assertIn(field, invite, f"Missing field: {field}")

    def test_generate_invite_uses_wg_public_key(self):
        invite = self.mgr.generate_invite('mycell', 'home.cell')
        self.assertEqual(invite['public_key'], 'serverpubkey=')

    def test_generate_invite_uses_configured_network(self):
        invite = self.mgr.generate_invite('mycell', 'home.cell')
        self.assertEqual(invite['vpn_subnet'], '10.0.0.0/24')

    def test_generate_invite_dns_ip_is_server_vpn_ip(self):
        invite = self.mgr.generate_invite('mycell', 'home.cell')
        self.assertEqual(invite['dns_ip'], '10.0.0.1')

    def test_generate_invite_uses_supplied_identity(self):
        invite = self.mgr.generate_invite('myhome', 'myhome.local')
        self.assertEqual(invite['cell_name'], 'myhome')
        self.assertEqual(invite['domain'], 'myhome.local')


class TestCellLinkManagerConnections(unittest.TestCase):

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()
        self.mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_add_connection_stores_link(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        links = self.mgr.list_connections()
        self.assertEqual(len(links), 1)
        self.assertEqual(links[0]['cell_name'], 'office')

    def test_add_connection_calls_add_cell_peer(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        self.wg.add_cell_peer.assert_called_once_with(
            name='office',
            public_key='officepubkey=',
            endpoint='5.6.7.8:51820',
            vpn_subnet='10.1.0.0/24',
        )

    def test_add_connection_calls_dns_forward(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        self.nm.add_cell_dns_forward.assert_called_once_with(
            domain='office.cell', dns_ip='10.1.0.1'
        )

    def test_add_connection_duplicate_raises(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        with self.assertRaises(ValueError):
            self.mgr.add_connection(SAMPLE_INVITE)

    def test_add_connection_persists_to_disk(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        # Create a fresh manager reading same dir
        mgr2 = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)
        links = mgr2.list_connections()
        self.assertEqual(len(links), 1)
        self.assertEqual(links[0]['cell_name'], 'office')

    def test_remove_connection_calls_wg_remove_peer(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        self.mgr.remove_connection('office')
        self.wg.remove_peer.assert_called_once_with('officepubkey=')

    def test_remove_connection_calls_dns_remove(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        self.mgr.remove_connection('office')
        self.nm.remove_cell_dns_forward.assert_called_once_with('office.cell')

    def test_remove_connection_deletes_from_list(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        self.mgr.remove_connection('office')
        self.assertEqual(len(self.mgr.list_connections()), 0)

    def test_remove_nonexistent_raises(self):
        with self.assertRaises(ValueError):
            self.mgr.remove_connection('nobody')

    def test_list_connections_empty_by_default(self):
        self.assertEqual(self.mgr.list_connections(), [])

    def test_multiple_connections(self):
        self.mgr.add_connection(SAMPLE_INVITE)
        second = {**SAMPLE_INVITE, 'cell_name': 'cabin', 'public_key': 'cabinpubkey=',
                  'vpn_subnet': '10.2.0.0/24', 'dns_ip': '10.2.0.1', 'domain': 'cabin.cell'}
        self.mgr.add_connection(second)
        self.assertEqual(len(self.mgr.list_connections()), 2)


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


# ---------------------------------------------------------------------------
# TestAddConnectionAtomicity
# ---------------------------------------------------------------------------

class TestAddConnectionAtomicity(unittest.TestCase):
    """Verify that add_connection rolls back correctly when WG or DNS steps fail."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()
        self.mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_wg_fail_does_not_call_dns(self):
        """When add_cell_peer returns False, add_cell_dns_forward must NOT be called."""
        self.wg.add_cell_peer.return_value = False
        with self.assertRaises(RuntimeError):
            self.mgr.add_connection(SAMPLE_INVITE)
        self.nm.add_cell_dns_forward.assert_not_called()

    def test_wg_fail_does_not_persist_link(self):
        """When WG fails, list_connections() must still return [] (nothing persisted)."""
        self.wg.add_cell_peer.return_value = False
        with self.assertRaises(RuntimeError):
            self.mgr.add_connection(SAMPLE_INVITE)
        self.assertEqual(self.mgr.list_connections(), [])

    def test_wg_fail_raises_runtime_error(self):
        """add_connection raises RuntimeError (not some other exception) when WG fails."""
        self.wg.add_cell_peer.return_value = False
        with self.assertRaises(RuntimeError):
            self.mgr.add_connection(SAMPLE_INVITE)

    def test_dns_warning_still_persists_link(self):
        """When DNS returns warnings (not a hard failure), the link IS still saved."""
        self.nm.add_cell_dns_forward.return_value = {
            'restarted': [],
            'warnings': ['CoreDNS reload timed out'],
        }
        self.mgr.add_connection(SAMPLE_INVITE)
        links = self.mgr.list_connections()
        self.assertEqual(len(links), 1)
        self.assertEqual(links[0]['cell_name'], 'office')

    def test_dns_warning_does_not_raise(self):
        """When DNS returns warnings, add_connection completes without raising."""
        self.nm.add_cell_dns_forward.return_value = {
            'restarted': [],
            'warnings': ['CoreDNS reload timed out'],
        }
        try:
            self.mgr.add_connection(SAMPLE_INVITE)
        except Exception as e:
            self.fail(f"add_connection raised unexpectedly with DNS warnings: {e}")


# ---------------------------------------------------------------------------
# TestAddConnectionPermissions
# ---------------------------------------------------------------------------

class TestAddConnectionPermissions(unittest.TestCase):
    """Verify that inbound_services controls the permissions field on the saved link."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()
        self.mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def _get_link(self):
        links = self.mgr.list_connections()
        self.assertEqual(len(links), 1)
        return links[0]

    def test_add_with_no_inbound_defaults_all_deny(self):
        """No inbound_services arg → all inbound permissions False."""
        self.mgr.add_connection(SAMPLE_INVITE)
        link = self._get_link()
        inbound = link['permissions']['inbound']
        for service, allowed in inbound.items():
            self.assertFalse(allowed, f"Expected {service} to be False, got {allowed}")

    def test_add_with_inbound_services_sets_them(self):
        """inbound_services=['calendar','files'] → those two True, others False."""
        self.mgr.add_connection(SAMPLE_INVITE, inbound_services=['calendar', 'files'])
        link = self._get_link()
        inbound = link['permissions']['inbound']
        self.assertTrue(inbound['calendar'])
        self.assertTrue(inbound['files'])
        self.assertFalse(inbound['mail'])
        self.assertFalse(inbound['webdav'])

    def test_inbound_invalid_service_ignored(self):
        """Passing 'badservice' in inbound_services does not appear in permissions."""
        self.mgr.add_connection(SAMPLE_INVITE, inbound_services=['badservice', 'calendar'])
        link = self._get_link()
        inbound = link['permissions']['inbound']
        self.assertNotIn('badservice', inbound)
        # valid one was still applied
        self.assertTrue(inbound['calendar'])


# ---------------------------------------------------------------------------
# TestUpdatePermissions
# ---------------------------------------------------------------------------

class TestUpdatePermissions(unittest.TestCase):
    """Tests for the new update_permissions / get_permissions methods."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()
        self.mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)
        # Add a connection so there is something to update
        self.mgr.add_connection(SAMPLE_INVITE)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_update_sets_inbound_values(self):
        """update_permissions with inbound={'calendar': True} persists correctly."""
        with patch('cell_link_manager.firewall_manager', create=True) as mock_fm:
            mock_fm.apply_cell_rules = MagicMock()
            self.mgr.update_permissions('office', {'calendar': True}, {})
        # Re-read from disk to confirm persistence
        mgr2 = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)
        perms = mgr2.get_permissions('office')
        self.assertTrue(perms['inbound']['calendar'])
        self.assertFalse(perms['inbound']['files'])

    def test_update_rejects_unknown_service_by_cleaning_it_out(self):
        """update_permissions with inbound={'bad': True} — 'bad' must not appear in saved perms."""
        with patch('cell_link_manager.firewall_manager', create=True) as mock_fm:
            mock_fm.apply_cell_rules = MagicMock()
            self.mgr.update_permissions('office', {'bad': True, 'calendar': True}, {})
        perms = self.mgr.get_permissions('office')
        self.assertNotIn('bad', perms['inbound'])
        self.assertTrue(perms['inbound']['calendar'])

    def test_update_nonexistent_cell_raises(self):
        """update_permissions on an unknown cell_name raises ValueError."""
        with self.assertRaises(ValueError):
            self.mgr.update_permissions('nosuchcell', {}, {})

    def test_get_permissions_returns_correct(self):
        """get_permissions returns the dict that was saved by update_permissions."""
        with patch('cell_link_manager.firewall_manager', create=True) as mock_fm:
            mock_fm.apply_cell_rules = MagicMock()
            self.mgr.update_permissions(
                'office',
                inbound={'calendar': True, 'files': False},
                outbound={'mail': True},
            )
        perms = self.mgr.get_permissions('office')
        self.assertIn('inbound', perms)
        self.assertIn('outbound', perms)
        self.assertTrue(perms['inbound']['calendar'])
        self.assertFalse(perms['inbound']['files'])
        self.assertTrue(perms['outbound']['mail'])

    def test_get_permissions_nonexistent_cell_raises(self):
        """get_permissions on an unknown cell_name raises ValueError."""
        with self.assertRaises(ValueError):
            self.mgr.get_permissions('nosuchcell')


# ---------------------------------------------------------------------------
# TestLoadMigration
# ---------------------------------------------------------------------------

class TestLoadMigration(unittest.TestCase):
    """Verify _load() lazily injects permissions field when it is missing."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.wg = _make_wg_mock()
        self.nm = _make_nm_mock()

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_load_injects_permissions_if_missing(self):
        """Write cell_links.json without permissions; _load should add all-False defaults."""
        links_file = os.path.join(self.test_dir, 'cell_links.json')
        legacy_links = [
            {
                'cell_name': 'legacy-office',
                'public_key': 'officepubkey=',
                'vpn_subnet': '10.1.0.0/24',
                'dns_ip': '10.1.0.1',
                'domain': 'legacy-office.cell',
                # NO 'permissions' key — simulates pre-migration data
            }
        ]
        with open(links_file, 'w') as f:
            json.dump(legacy_links, f)

        mgr = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)
        links = mgr.list_connections()

        self.assertEqual(len(links), 1)
        link = links[0]
        self.assertIn('permissions', link)
        perms = link['permissions']
        self.assertIn('inbound', perms)
        self.assertIn('outbound', perms)
        for service in ('calendar', 'files', 'mail', 'webdav'):
            self.assertFalse(perms['inbound'][service])
            self.assertFalse(perms['outbound'][service])

    def test_load_migration_persists_to_disk(self):
        """After migration, re-loading the same file returns the injected permissions."""
        links_file = os.path.join(self.test_dir, 'cell_links.json')
        with open(links_file, 'w') as f:
            json.dump([{
                'cell_name': 'old-cell',
                'public_key': 'somepubkey=',
                'vpn_subnet': '10.2.0.0/24',
                'dns_ip': '10.2.0.1',
                'domain': 'old-cell.cell',
            }], f)

        mgr1 = CellLinkManager(self.test_dir, self.test_dir, self.wg, self.nm)
        mgr1.list_connections()  # triggers migration + save

        # Read the file directly and confirm permissions are now on disk
        with open(links_file) as f:
            raw = json.load(f)
        self.assertIn('permissions', raw[0])