pic/tests/test_peer_wg_integration.py

#!/usr/bin/env python3
"""
Tests for peer add/remove flow — ensures server-side WireGuard AllowedIPs
are always the peer's /32 VPN IP, never the client tunnel AllowedIPs.
"""

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

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

from wireguard_manager import WireGuardManager
from peer_registry import PeerRegistry


class TestServerSideAllowedIPs(unittest.TestCase):
    """Server-side peer AllowedIPs must always be peer_ip/32."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.data_dir = os.path.join(self.tmp, 'data')
        self.config_dir = os.path.join(self.tmp, 'config')
        os.makedirs(self.data_dir)
        os.makedirs(self.config_dir)
        # Patch syncconf so tests don't need docker
        patcher = patch.object(WireGuardManager, '_syncconf', return_value=None)
        self.mock_sync = patcher.start()
        self.addCleanup(patcher.stop)
        self.wg = WireGuardManager(self.data_dir, self.config_dir)

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

    def _config(self):
        with open(self.wg._config_file()) as f:
            return f.read()

    def test_add_peer_uses_host_slash32(self):
        """Peer added with /32 stays as /32 in config."""
        self.wg.add_peer('alice', 'YWxpY2VfdGVzdF93Z19wZWVyX2tleV8xMjM0NTY3OCE=', '', allowed_ips='10.0.0.2/32')
        cfg = self._config()
        self.assertIn('AllowedIPs = 10.0.0.2/32', cfg)

    def test_full_tunnel_client_ips_rejected(self):
        """add_peer must refuse 0.0.0.0/0 — it would route all internet traffic to that peer."""
        result = self.wg.add_peer('bob', 'Ym9iX3Rlc3Rfd2dfcGVlcl9rZXlfMTIzNDU2Nzg5MCE=', '', allowed_ips='0.0.0.0/0, ::/0')
        self.assertFalse(result,
            "0.0.0.0/0 in server peer AllowedIPs routes ALL traffic to that peer, breaking internet")

    def test_split_tunnel_client_ips_rejected(self):
        """add_peer must refuse 172.20.0.0/16 — it would route docker network to that peer."""
        result = self.wg.add_peer('carol', 'Y2Fyb2xfdGVzdF93Z19wZWVyX2tleV8xMjM0NTY3OCE=', '', allowed_ips='10.0.0.0/24, 172.20.0.0/16')
        self.assertFalse(result,
            "172.20.0.0/16 in server peer AllowedIPs routes docker network traffic to that peer")

    def test_remove_peer_cleans_config(self):
        self.wg.add_peer('dave', 'ZGF2ZV90ZXN0X3dnX3BlZXJfa2V5XzEyMzQ1Njc4OSE=', '', allowed_ips='10.0.0.4/32')
        self.wg.remove_peer('ZGF2ZV90ZXN0X3dnX3BlZXJfa2V5XzEyMzQ1Njc4OSE=')
        cfg = self._config()
        self.assertNotIn('ZGF2ZV90ZXN0X3dnX3BlZXJfa2V5XzEyMzQ1Njc4OSE=', cfg)

    def test_syncconf_called_on_add(self):
        self.wg.add_peer('eve', 'ZXZlX3Rlc3Rfd2dfcGVlcl9rZXlfXzEyMzQ1Njc4OSE=', '', allowed_ips='10.0.0.5/32')
        self.mock_sync.assert_called()

    def test_syncconf_called_on_remove(self):
        self.wg.add_peer('frank', 'ZnJhbmtfdGVzdF93Z19wZWVyX2tleV8xMjM0NTY3OCE=', '', allowed_ips='10.0.0.6/32')
        self.mock_sync.reset_mock()
        self.wg.remove_peer('ZnJhbmtfdGVzdF93Z19wZWVyX2tleV8xMjM0NTY3OCE=')
        self.mock_sync.assert_called()


class TestAutoAssignIP(unittest.TestCase):
    """Auto-assigned peer IPs must be unique /32s starting at 10.0.0.2."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp()
        self.registry = PeerRegistry(data_dir=self.tmp, config_dir=self.tmp)

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

    def _next_ip(self):
        import ipaddress
        used = {p.get('ip', '').split('/')[0] for p in self.registry.list_peers()}
        for host in ipaddress.ip_network('10.0.0.0/24').hosts():
            ip = str(host)
            if ip != '10.0.0.1' and ip not in used:
                return ip
        raise ValueError('No free IPs')

    def test_first_peer_gets_10_0_0_2(self):
        ip = self._next_ip()
        self.assertEqual(ip, '10.0.0.2')

    def test_second_peer_gets_10_0_0_3(self):
        self.registry.add_peer({'peer': 'p1', 'ip': '10.0.0.2'})
        ip = self._next_ip()
        self.assertEqual(ip, '10.0.0.3')

    def test_no_duplicate_ips(self):
        assigned = []
        for i in range(5):
            ip = self._next_ip()
            self.assertNotIn(ip, assigned, f"Duplicate IP assigned: {ip}")
            assigned.append(ip)
            self.registry.add_peer({'peer': f'peer{i}', 'ip': ip})

    def test_server_ip_never_assigned(self):
        # Fill up .2 through .10
        for i in range(2, 11):
            self.registry.add_peer({'peer': f'p{i}', 'ip': f'10.0.0.{i}'})
        ip = self._next_ip()
        self.assertNotEqual(ip, '10.0.0.1', "Server IP 10.0.0.1 must never be assigned to a peer")


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