#!/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])