import base64 import copy import json import os from pathlib import Path import tempfile import unittest import shutil import uuid from unittest.mock import patch import agent def key(n): return base64.b64encode(bytes([n]) * 32).decode() class ControllerTest(unittest.TestCase): def setUp(self): # Windows Python 3.13's TemporaryDirectory mode 0700 excludes the restricted sandbox token. base = Path.cwd() / "tmp" if os.name == "nt" else Path(tempfile.gettempdir()) base.mkdir(exist_ok=True) self.root = base / ("flow-agent-test-" + uuid.uuid4().hex) self.root.mkdir() def cleanup(): if self.root.resolve().parent != base.resolve(): raise RuntimeError("Unexpected test directory") shutil.rmtree(self.root) self.addCleanup(cleanup) self.addCleanup(patch.stopall) patch.object(agent, "CONFIG", self.root).start() self.interface = dict(PrivateKey=key(1), Address="10.255.254.1/24", ListenPort="51821", MTU="1420") self.third = dict(PublicKey=key(3), AllowedIPs="192.0.2.5/32", PresharedKey=key(9), PersistentKeepalive="0") self.desired = dict(collectorIp="10.255.254.1", prefix="10.255.254.0/24", flowListenPort=4739, wgListenPort=51821, hostPublicKey=key(2), hostPrivateKey=key(1), peers=[dict(serverId=1, publicKey=key(4), allowedIps=["10.255.254.2/32"], endpoint="router.test:13232")]) (self.root / "binding.json").write_text(json.dumps({k: self.desired[k] for k in ("collectorIp", "prefix", "flowListenPort", "wgListenPort")})) self.previous = agent.render(self.interface, [self.third]) (self.root / "wg-flow.conf").write_text(self.previous) self.runtime = [copy.deepcopy(self.third)] self.calls = [] def run(*args, input=None): self.calls.append(args) if args == ("wg", "pubkey"): return key(2) if input == key(1) + "\n" else key(8) if args == ("wg", "showconf", "wg-flow"): return agent.render(self.interface, self.runtime, native=True) raise AssertionError(args) patch.object(agent, "run", run).start() patch.object(agent, "live", self.live).start() patch.object(agent, "sync", self.sync).start() self.synced = [] def live(self): return dict(interfaceUp=True, publicKey=key(2), listenPort=51821, addresses=["10.255.254.1/24"], peers=[]), copy.deepcopy(self.runtime) def sync(self, text): self.synced.append(text) # Native configs don't include wg-quick's Address; simulate applying their peers. _, self.runtime = agent.parse_conf(text.replace("ListenPort = 51821", "ListenPort = 51821\nAddress = 10.255.254.1/24")) def test_apply_preserves_third_party_psk_and_is_idempotent(self): self.assertFalse(agent.check(self.desired)["configMatches"]) self.assertTrue(agent.check(self.desired, apply=True)["configMatches"]) text = (self.root / "wg-flow.conf").read_text() self.assertIn(key(3), text) self.assertIn(key(9), text) self.assertIn(key(4), text) self.assertNotIn("Address", self.synced[0]) self.assertTrue(agent.check(self.desired)["configMatches"]) agent.check(self.desired, apply=True) self.assertEqual((self.root / "wg-flow.conf").read_text(), text) if os.name != "nt": self.assertEqual((self.root / "wg-flow.conf").stat().st_mode & 0o777, 0o600) def test_runtime_only_peer_preserved(self): self.runtime.append(dict(PublicKey=key(6), AllowedIPs="198.51.100.7/32", PersistentKeepalive="0")) agent.check(self.desired, apply=True) self.assertIn(key(6), (self.root / "wg-flow.conf").read_text()) def test_runtime_peer_without_allowed_ips_preserved(self): self.runtime.append(dict(PublicKey=key(6), AllowedIPs="", PersistentKeepalive="0")) agent.check(self.desired, apply=True) self.assertTrue(agent.check(self.desired)["configMatches"]) def test_third_party_runtime_settings_are_not_reset(self): # Runtime may contain a roaming endpoint and intentionally unsaved settings. self.runtime[0].update(Endpoint="198.51.100.2:50000", PersistentKeepalive="10") agent.check(self.desired, apply=True) third = next(p for p in self.runtime if p["PublicKey"] == key(3)) self.assertEqual(third["Endpoint"], "198.51.100.2:50000") self.assertEqual(third["PersistentKeepalive"], "10") _, stored = agent.parse_conf((self.root / "wg-flow.conf").read_text()) self.assertEqual(next(p for p in stored if p["PublicKey"] == key(3))["PersistentKeepalive"], "0") def test_managed_roaming_endpoint_preserved_without_changing_desired_dns(self): agent.check(self.desired, apply=True) next(p for p in self.runtime if p["PublicKey"] == key(4))["Endpoint"] = "198.51.100.2:50000" agent.check(self.desired, apply=True) _, stored = agent.parse_conf((self.root / "wg-flow.conf").read_text()) self.assertEqual(next(p for p in stored if p["PublicKey"] == key(4))["Endpoint"], "router.test:13232") self.assertEqual(next(p for p in self.runtime if p["PublicKey"] == key(4))["Endpoint"], "198.51.100.2:50000") self.assertTrue(agent.check(self.desired)["configMatches"]) def test_down_interface_cannot_report_applied(self): status, peers = self.live() with patch.object(agent, "live", return_value=({**status, "interfaceUp": False}, peers)): with self.assertRaisesRegex(agent.AgentError, "выключен"): agent.check(self.desired, apply=True) self.assertEqual(self.synced, []) def test_overlap_rejected_before_any_write(self): self.runtime.append(dict(PublicKey=key(6), AllowedIPs="10.255.254.0/24")) with self.assertRaisesRegex(agent.AgentError, "пересекаются"): agent.check(self.desired, apply=True) self.assertEqual((self.root / "wg-flow.conf").read_text(), self.previous) self.assertEqual(self.synced, []) def test_different_interface_identity_rejected(self): self.desired["hostPrivateKey"], self.desired["hostPublicKey"] = key(7), key(8) with self.assertRaisesRegex(agent.AgentError, "существующего"): agent.check(self.desired, apply=True) self.assertEqual(self.synced, []) def test_failed_apply_restores_file_and_runtime(self): rollback = agent.render(self.interface, self.runtime, native=True) with patch.object(agent, "sync", side_effect=[agent.AgentError("secret error"), None]) as sync: with self.assertRaisesRegex(agent.AgentError, "предыдущая конфигурация восстановлена"): agent.check(self.desired, apply=True) self.assertEqual(sync.call_args_list[-1].args, (rollback,)) self.assertEqual((self.root / "wg-flow.conf").read_text(), self.previous) def test_failed_rollback_is_explicit(self): with patch.object(agent, "sync", side_effect=agent.AgentError("secret error")): with self.assertRaisesRegex(agent.AgentError, "восстановление требует ручной проверки"): agent.check(self.desired, apply=True) def test_invalid_or_shell_input_rejected(self): for change in [dict(flowListenPort=4740), dict(prefix="10.255.254.0/28"), dict(hostPrivateKey="$(id)"), dict(command="id"), dict(peers=[dict(serverId=1, publicKey=key(4), allowedIps=["0.0.0.0/0"])]), dict(peers=[dict(serverId=1, publicKey=key(4), allowedIps=["10.255.254.2/32"], endpoint="router.test:13232\nPostUp = id")])]: with self.subTest(change=list(change)): with self.assertRaises(agent.AgentError): agent.check({**self.desired, **change}, apply=True) self.assertEqual(self.synced, []) def test_hooks_and_duplicate_fields_rejected(self): for suffix in ("\nPostUp = id", "\nListenPort = 3", "\n[Interface]\nPrivateKey = secret"): with self.assertRaises(agent.AgentError): agent.parse_conf(agent.render(self.interface, []) + suffix) def test_endpoint_change_requires_apply(self): agent.check(self.desired, apply=True) self.desired["peers"][0]["endpoint"] = "new-router.test:13232" self.assertFalse(agent.check(self.desired)["configMatches"]) def test_command_errors_do_not_expose_keys(self): with patch.object(agent.subprocess, "run", return_value=type("Result", (), {"returncode": 1, "stdout": key(1), "stderr": key(9)})()): # Use the original function, not the fixture runner. with self.assertRaisesRegex(agent.AgentError, "не выполнена") as error: ORIGINAL_RUN("wg", "show") self.assertNotIn(key(1), str(error.exception)) self.assertNotIn(key(9), str(error.exception)) ORIGINAL_RUN = agent.run @unittest.skipUnless(os.environ.get("FLOW_AGENT_KERNEL_TEST") == "1", "Requires Linux kernel WireGuard and NET_ADMIN in an isolated container") class KernelTest(unittest.TestCase): def test_real_interface_apply_and_preservation(self): with tempfile.TemporaryDirectory() as folder, patch.object(agent, "CONFIG", Path(folder)): private = agent.run("wg", "genkey") public = agent.run("wg", "pubkey", input=private + "\n") other = agent.run("wg", "pubkey", input=agent.run("wg", "genkey") + "\n") desired_key = agent.run("wg", "pubkey", input=agent.run("wg", "genkey") + "\n") interface = dict(PrivateKey=private, Address="10.255.254.1/24", ListenPort="51821") third = dict(PublicKey=other, AllowedIPs="192.0.2.3/32") path = Path(folder) (path / "wg-flow.conf").write_text(agent.render(interface, [third])) snapshot = dict(collectorIp="10.255.254.1", prefix="10.255.254.0/24", flowListenPort=4739, wgListenPort=51821, hostPublicKey=public, hostPrivateKey=private, peers=[dict(serverId=1, publicKey=desired_key, allowedIps=["10.255.254.2/32"], endpoint="127.0.0.1:13232")]) (path / "binding.json").write_text(json.dumps({k: snapshot[k] for k in ("collectorIp", "prefix", "flowListenPort", "wgListenPort")})) agent.run("ip", "link", "add", "dev", agent.IFACE, "type", "wireguard") try: agent.run("ip", "address", "add", "10.255.254.1/24", "dev", agent.IFACE) agent.sync(agent.render(interface, [third], native=True)) agent.run("ip", "link", "set", "dev", agent.IFACE, "up") before = json.loads(agent.run("ip", "-j", "link", "show", "dev", agent.IFACE))[0]["ifindex"] self.assertTrue(agent.check(snapshot, apply=True)["configMatches"]) self.assertTrue(agent.check(snapshot)["configMatches"]) status, _ = agent.live() self.assertEqual({p["publicKey"] for p in status["peers"]}, {other, desired_key}) self.assertNotIn(private, json.dumps(status)) after = json.loads(agent.run("ip", "-j", "link", "show", "dev", agent.IFACE))[0]["ifindex"] self.assertEqual(before, after) finally: agent.run("ip", "link", "delete", "dev", agent.IFACE) if __name__ == "__main__": unittest.main()