"""
Tests for EgressManager — per-service egress enforcement via host iptables.

Connectivity v2: a service routes through a connection *instance* (by id),
sharing the connection's fwmark / routing table / redirect port. The egress
override map is service_id → connection_id, and (mark, table, redirect_port)
are resolved from ConnectivityManager.get_connection(id). EgressManager no
longer owns its own per-type MARKS/TABLES.

All subprocess calls (iptables, iptables-save, iptables-restore, ip rule,
docker inspect) and manager state are mocked so these tests run without any
live infrastructure or root privileges.
"""

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

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

import egress_manager as em_module
from egress_manager import EgressManager, EGRESS_CHAIN


# ---------------------------------------------------------------------------
# Connection fixtures — mirror the v2 allocator output.
# ---------------------------------------------------------------------------

CONNECTIONS = {
    'conn_wg': {
        'id': 'conn_wg', 'type': 'wireguard_ext', 'name': 'Work VPN',
        'enabled': True, 'mark': 0x1000, 'table': 1000,
        'iface': 'wgext_aaaa', 'redirect_port': None,
    },
    'conn_ovpn': {
        'id': 'conn_ovpn', 'type': 'openvpn', 'name': 'OVPN',
        'enabled': True, 'mark': 0x1010, 'table': 1001,
        'iface': 'ovpn_bbbb', 'redirect_port': None,
    },
    'conn_tor': {
        'id': 'conn_tor', 'type': 'tor', 'name': 'Tor',
        'enabled': True, 'mark': 0x1020, 'table': 1002,
        'iface': None, 'redirect_port': 9100,
    },
}


def _make_manager(installed=None, overrides=None, connections=None):
    """Build an EgressManager backed by mock config + connectivity managers."""
    cm = MagicMock()
    cm.get_installed_services.return_value = installed or {}
    cm.configs = {'egress_overrides': overrides or {}}
    cm._save_all_configs = MagicMock()

    conns = connections if connections is not None else CONNECTIONS
    conn_list = list(conns.values())
    cm.list_connections.return_value = conn_list
    cm.get_connection.side_effect = lambda cid: conns.get(cid)

    connectivity = MagicMock()
    connectivity.list_connections.return_value = conn_list
    connectivity.get_connection.side_effect = lambda cid: conns.get(cid)

    mgr = EgressManager(config_manager=cm, connectivity_manager=connectivity)
    return mgr, cm


def _subprocess_ok(stdout=''):
    return MagicMock(returncode=0, stdout=stdout, stderr='')


def _subprocess_fail(stderr='error', stdout=''):
    return MagicMock(returncode=1, stdout=stdout, stderr=stderr)


def _make_manifest(has_egress=True, egress_default='conn_wg',
                   allowed=None, container_name='cell-myapp'):
    """Return a minimal manifest dict with optional egress configuration.

    `allowed` is a list of connection *types* (manifests are type-scoped).
    """
    m = {
        'id': 'myapp',
        'name': 'My App',
        'container_name': container_name,
    }
    if has_egress:
        m['has_egress'] = True
        m['egress'] = {
            'default': egress_default,
            'allowed': allowed if allowed is not None
            else ['wireguard_ext', 'openvpn', 'tor', 'sshuttle', 'proxy'],
        }
    else:
        m['has_egress'] = False
    return m


def _installed_with_manifest(manifest, service_id='myapp'):
    return {service_id: {'id': service_id, 'manifest': manifest}}


# ---------------------------------------------------------------------------
# 1. default exit → no iptables rule-adding
# ---------------------------------------------------------------------------

class TestApplyServiceDefaultExit(unittest.TestCase):

    def test_apply_service_default_exit_no_iptables_calls(self):
        manifest = _make_manifest(egress_default='default')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = _subprocess_ok(stdout='172.20.0.50\n')
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'])
        self.assertEqual(result.get('exit_via'), 'default')

        rule_add_calls = [
            c for c in mock_run.call_args_list
            if c.args and c.args[0][:1] == ['iptables']
            and any(a in c.args[0] for a in ('-A', '-I', 'MARK', 'REDIRECT'))
        ]
        self.assertEqual(rule_add_calls, [])


# ---------------------------------------------------------------------------
# 2. wireguard_ext connection → mark rule with the connection's own mark
# ---------------------------------------------------------------------------

class TestApplyServiceWireguardExt(unittest.TestCase):

    def test_apply_service_wireguard_ext_adds_mark_rule(self):
        manifest = _make_manifest(egress_default='conn_wg')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        calls_made = []

        def fake_run(cmd, **kwargs):
            calls_made.append(cmd)
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='172.20.0.50\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            if cmd[:3] == ['ip', 'rule', 'del']:
                return _subprocess_fail()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'], result)
        self.assertEqual(result['exit_via'], 'conn_wg')

        mark_calls = [
            c for c in calls_made
            if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
        ]
        self.assertGreater(len(mark_calls), 0, 'No MARK rule was added')
        mark_cmd = ' '.join(mark_calls[0])
        self.assertIn('0x1000', mark_cmd)  # the connection's mark
        self.assertIn('pic-egr-myapp', mark_cmd)
        self.assertIn('mangle', mark_cmd)

        # The ip rule must point fwmark→the connection's table.
        ip_rule_add = [
            c for c in calls_made
            if c[:3] == ['ip', 'rule', 'add']
        ]
        self.assertGreater(len(ip_rule_add), 0)
        self.assertIn('1000', ' '.join(ip_rule_add[0]))  # table


# ---------------------------------------------------------------------------
# 3. openvpn connection → its own mark
# ---------------------------------------------------------------------------

class TestApplyServiceOpenVPN(unittest.TestCase):

    def test_apply_service_openvpn_adds_mark_rule(self):
        manifest = _make_manifest(egress_default='conn_ovpn')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        calls_made = []

        def fake_run(cmd, **kwargs):
            calls_made.append(cmd)
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='172.20.0.51\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            if cmd[:3] == ['ip', 'rule', 'del']:
                return _subprocess_fail()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'], result)
        self.assertEqual(result['exit_via'], 'conn_ovpn')

        mark_calls = [
            c for c in calls_made
            if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
        ]
        self.assertGreater(len(mark_calls), 0)
        self.assertIn('0x1010', ' '.join(mark_calls[0]))


# ---------------------------------------------------------------------------
# 4. tor (redirect-style) connection → mark + REDIRECT to its redirect_port
# ---------------------------------------------------------------------------

class TestApplyServiceTor(unittest.TestCase):

    def test_apply_service_tor_adds_mark_and_redirect(self):
        manifest = _make_manifest(egress_default='conn_tor')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        calls_made = []

        def fake_run(cmd, **kwargs):
            calls_made.append(cmd)
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='172.20.0.52\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            if cmd[:3] == ['ip', 'rule', 'del']:
                return _subprocess_fail()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'], result)
        self.assertEqual(result['exit_via'], 'conn_tor')

        mark_calls = [
            c for c in calls_made
            if 'iptables' in str(c) and 'MARK' in c and '--set-mark' in c
        ]
        self.assertGreater(len(mark_calls), 0, 'No MARK rule found')
        self.assertIn('0x1020', ' '.join(mark_calls[0]))

        redirect_calls = [
            c for c in calls_made
            if 'iptables' in str(c) and 'REDIRECT' in c
        ]
        self.assertGreater(len(redirect_calls), 0, 'No REDIRECT rule found')
        redirect_cmd = ' '.join(redirect_calls[0])
        self.assertIn('9100', redirect_cmd)  # the connection's redirect_port
        self.assertIn('nat', redirect_cmd)


# ---------------------------------------------------------------------------
# 5. no container IP → error
# ---------------------------------------------------------------------------

class TestApplyServiceNoContainerIP(unittest.TestCase):

    def test_apply_service_no_container_ip_returns_error(self):
        manifest = _make_manifest(egress_default='conn_wg')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        def fake_run(cmd, **kwargs):
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.apply_service('myapp')

        self.assertFalse(result['ok'])
        self.assertIn('container IP not discoverable', result.get('error', ''))


# ---------------------------------------------------------------------------
# 6. container IP retries
# ---------------------------------------------------------------------------

class TestApplyServiceRetries(unittest.TestCase):

    def test_apply_service_container_ip_retries(self):
        manifest = _make_manifest(egress_default='conn_wg')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        inspect_count = [0]

        def fake_run(cmd, **kwargs):
            if 'docker' in cmd and 'inspect' in cmd:
                inspect_count[0] += 1
                if inspect_count[0] == 1:
                    return _subprocess_ok(stdout='\n')
                return _subprocess_ok(stdout='172.20.0.50\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            if cmd[:3] == ['ip', 'rule', 'del']:
                return _subprocess_fail()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            with patch('time.sleep'):
                result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'], result)
        self.assertGreaterEqual(inspect_count[0], 2)


# ---------------------------------------------------------------------------
# 7. has_egress False → skipped
# ---------------------------------------------------------------------------

class TestHasEgressFalse(unittest.TestCase):

    def test_has_egress_false_skips_rules(self):
        manifest = _make_manifest(has_egress=False)
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = _subprocess_ok(stdout='')
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'])
        self.assertTrue(result.get('skipped'))

        rule_add_calls = [
            c for c in mock_run.call_args_list
            if c.args and c.args[0][:1] == ['iptables']
            and any(a in c.args[0] for a in ('-A', '-I', 'MARK', 'REDIRECT'))
        ]
        self.assertEqual(rule_add_calls, [])


# ---------------------------------------------------------------------------
# 8. has_egress True but no egress block → skipped
# ---------------------------------------------------------------------------

class TestHasEgressMissingBlock(unittest.TestCase):

    def test_has_egress_missing_egress_block_skips(self):
        manifest = {
            'id': 'myapp',
            'container_name': 'cell-myapp',
            'has_egress': True,
        }
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        with patch('subprocess.run') as mock_run:
            mock_run.return_value = _subprocess_ok(stdout='')
            result = mgr.apply_service('myapp')

        self.assertTrue(result['ok'])
        self.assertTrue(result.get('skipped'))


# ---------------------------------------------------------------------------
# 9. clear_service removes only the tagged rules
# ---------------------------------------------------------------------------

class TestClearService(unittest.TestCase):

    def test_clear_service_removes_tagged_rules(self):
        mgr, _ = _make_manager()

        mangle_rules = (
            '-A PIC_EGRESS -s 172.20.0.50 -j MARK --set-mark 0x1000 '
            '-m comment --comment "pic-egr-myapp"\n'
            '-A PIC_EGRESS -s 172.20.0.99 -j MARK --set-mark 0x1000 '
            '-m comment --comment "pic-egr-otherapp"\n'
        )
        nat_rules = ''
        restore_inputs = {}

        def fake_run(cmd, input=None, **kwargs):
            if cmd == ['iptables-save', '-t', 'mangle']:
                return _subprocess_ok(stdout=mangle_rules)
            if cmd == ['iptables-save', '-t', 'nat']:
                return _subprocess_ok(stdout=nat_rules)
            if cmd == ['iptables-restore', '-T', 'mangle']:
                restore_inputs['mangle'] = input
                return _subprocess_ok()
            if cmd == ['iptables-restore', '-T', 'nat']:
                restore_inputs['nat'] = input
                return _subprocess_ok()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.clear_service('myapp')

        self.assertTrue(result['ok'])
        restored = restore_inputs.get('mangle', '')
        self.assertNotIn('pic-egr-myapp', restored)
        self.assertIn('pic-egr-otherapp', restored)


# ---------------------------------------------------------------------------
# 10. set_service_exit rejects a connection whose type is not in allowed
# ---------------------------------------------------------------------------

class TestSetServiceExitRejectNotAllowed(unittest.TestCase):

    def test_set_service_exit_rejects_type_not_in_allowed(self):
        manifest = _make_manifest(
            egress_default='default',
            allowed=['tor'],  # wireguard_ext not allowed
        )
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        result = mgr.set_service_exit('myapp', 'conn_wg')

        self.assertFalse(result['ok'])
        self.assertIn('error', result)
        self.assertIn('allowed', result['error'])


# ---------------------------------------------------------------------------
# 11. set_service_exit persists the connection id and applies
# ---------------------------------------------------------------------------

class TestSetServiceExitPersistsAndApplies(unittest.TestCase):

    def test_set_service_exit_persists_and_applies(self):
        manifest = _make_manifest(egress_default='default')
        mgr, cm = _make_manager(installed=_installed_with_manifest(manifest))

        apply_calls = []
        mgr.apply_service = lambda sid: apply_calls.append(sid) or {
            'ok': True, 'exit_via': 'conn_tor'}

        result = mgr.set_service_exit('myapp', 'conn_tor')

        self.assertTrue(result['ok'], result)
        self.assertIn('myapp', apply_calls)
        cm._save_all_configs.assert_called()
        self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'conn_tor')

    def test_set_service_exit_default_clears_override(self):
        manifest = _make_manifest(egress_default='conn_wg')
        mgr, cm = _make_manager(
            installed=_installed_with_manifest(manifest),
            overrides={'myapp': 'conn_tor'},
        )
        mgr.apply_service = lambda sid: {'ok': True, 'exit_via': 'default'}

        result = mgr.set_service_exit('myapp', 'default')

        self.assertTrue(result['ok'])
        self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'default')

    def test_set_service_exit_unknown_connection_rejected(self):
        manifest = _make_manifest(egress_default='default')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        result = mgr.set_service_exit('myapp', 'conn_ghost')

        self.assertFalse(result['ok'])
        self.assertIn('unknown connection', result['error'])

    def test_set_service_exit_legacy_type_resolves_to_single_instance(self):
        """Back-compat shim: a legacy type resolves to the one instance of it."""
        manifest = _make_manifest(egress_default='default', allowed=['tor'])
        mgr, cm = _make_manager(installed=_installed_with_manifest(manifest))
        mgr.apply_service = lambda sid: {'ok': True}

        result = mgr.set_service_exit('myapp', 'tor')

        self.assertTrue(result['ok'], result)
        self.assertEqual(cm.configs['egress_overrides'].get('myapp'), 'conn_tor')


# ---------------------------------------------------------------------------
# 12. apply_all iterates installed services
# ---------------------------------------------------------------------------

class TestApplyAll(unittest.TestCase):

    def test_apply_all_iterates_installed_services(self):
        manifests = {
            'svc1': _make_manifest(egress_default='conn_wg'),
            'svc2': _make_manifest(egress_default='conn_ovpn'),
            'svc3': _make_manifest(egress_default='conn_tor'),
        }
        installed = {
            sid: {'id': sid, 'manifest': m}
            for sid, m in manifests.items()
        }
        mgr, _ = _make_manager(installed=installed)

        applied = []
        mgr.apply_service = lambda sid: applied.append(sid) or {'ok': True}

        result = mgr.apply_all()

        self.assertTrue(result['ok'])
        self.assertEqual(sorted(applied), ['svc1', 'svc2', 'svc3'])


# ---------------------------------------------------------------------------
# 13. service + peer on the same connection share the same mark
# ---------------------------------------------------------------------------

class TestSharedMarkWithConnection(unittest.TestCase):

    def test_service_uses_connection_mark_not_a_private_table(self):
        """A service routed via a connection uses that connection's mark/table,
        i.e. the SAME resources a peer on that connection would use."""
        manifest = _make_manifest(egress_default='conn_wg')
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        captured = {}

        def fake_run(cmd, **kwargs):
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='172.20.0.50\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            if cmd[:3] == ['ip', 'rule', 'del']:
                return _subprocess_fail()
            if cmd[:3] == ['ip', 'rule', 'add']:
                captured['ip_rule'] = ' '.join(cmd)
            if 'MARK' in cmd:
                captured['mark'] = ' '.join(cmd)
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            mgr.apply_service('myapp')

        # conn_wg has mark 0x1000 / table 1000 — both must appear, proving the
        # service inherits the connection's shared routing resources.
        self.assertIn('0x1000', captured.get('mark', ''))
        self.assertIn('1000', captured.get('ip_rule', ''))


# ---------------------------------------------------------------------------
# 14. apply_service with an unknown connection id → error
# ---------------------------------------------------------------------------

class TestApplyServiceUnknownExit(unittest.TestCase):

    def test_apply_service_unknown_connection_rejected(self):
        manifest = {
            'id': 'myapp',
            'container_name': 'cell-myapp',
            'has_egress': True,
            'egress': {
                'default': 'conn_ghost',  # no such connection
                'allowed': ['wireguard_ext'],
            },
        }
        mgr, _ = _make_manager(installed=_installed_with_manifest(manifest))

        def fake_run(cmd, **kwargs):
            if 'docker' in cmd and 'inspect' in cmd:
                return _subprocess_ok(stdout='172.20.0.50\n')
            if 'iptables-save' in cmd:
                return _subprocess_ok(stdout='')
            if 'iptables-restore' in cmd:
                return _subprocess_ok()
            return _subprocess_ok()

        with patch('subprocess.run', side_effect=fake_run):
            result = mgr.apply_service('myapp')

        self.assertFalse(result['ok'])
        self.assertIn('error', result)


# ---------------------------------------------------------------------------
# _has_egress edge cases
# ---------------------------------------------------------------------------

class TestHasEgressLogic(unittest.TestCase):

    def setUp(self):
        self.mgr, _ = _make_manager()

    def test_has_egress_both_required(self):
        m = {'has_egress': True, 'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
        self.assertTrue(self.mgr._has_egress(m))

    def test_has_egress_false_field(self):
        m = {'has_egress': False, 'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
        self.assertFalse(self.mgr._has_egress(m))

    def test_has_egress_missing_has_egress_key(self):
        m = {'egress': {'default': 'conn_tor', 'allowed': ['tor']}}
        self.assertFalse(self.mgr._has_egress(m))

    def test_has_egress_empty_egress_dict(self):
        m = {'has_egress': True, 'egress': {}}
        self.assertFalse(self.mgr._has_egress(m))


# ---------------------------------------------------------------------------
# _resolve_exit (now returns connection ids)
# ---------------------------------------------------------------------------

class TestResolveExit(unittest.TestCase):

    def test_override_takes_precedence(self):
        mgr, _ = _make_manager(overrides={'myapp': 'conn_ovpn'})
        manifest = _make_manifest(egress_default='conn_wg')
        self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_ovpn')

    def test_manifest_default_used_when_no_override(self):
        mgr, _ = _make_manager(overrides={})
        manifest = _make_manifest(egress_default='conn_tor')
        self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_tor')

    def test_fallback_to_default_when_no_egress_block(self):
        mgr, _ = _make_manager(overrides={})
        manifest = {'id': 'myapp'}
        self.assertEqual(mgr._resolve_exit('myapp', manifest), 'default')

    def test_legacy_type_override_migrates_to_connection_id(self):
        """An old override holding a type string resolves to the migrated id."""
        mgr, _ = _make_manager(overrides={'myapp': 'wireguard_ext'})
        manifest = _make_manifest(egress_default='default')
        self.assertEqual(mgr._resolve_exit('myapp', manifest), 'conn_wg')

    def test_legacy_type_default_with_no_instance_falls_back(self):
        """A legacy type with no matching instance falls back to 'default'."""
        mgr, _ = _make_manager(connections={})
        manifest = _make_manifest(egress_default='tor')
        self.assertEqual(mgr._resolve_exit('myapp', manifest), 'default')


# ---------------------------------------------------------------------------
# apply_service with missing manifest
# ---------------------------------------------------------------------------

class TestApplyServiceMissingManifest(unittest.TestCase):

    def test_apply_service_missing_manifest_returns_error(self):
        mgr, _ = _make_manager(installed={})
        result = mgr.apply_service('ghost')
        self.assertFalse(result['ok'])
        self.assertIn('error', result)


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