202 lines
11 KiB
Python
202 lines
11 KiB
Python
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()
|