feat(netflow): автоматизировать настройку туннеля коллектора

This commit is contained in:
Denozordec committed 2026-10-16 15:03:23 +07:00
1 parent fcca92706b
commit db81cf0f97
33 files changed
+1503 -22

No files matched your search

+201
View File
@@ -0,0 +1,201 @@
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()