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

+11
View File
@@ -0,0 +1,11 @@
FROM alpine:3.22 AS base
RUN apk add --no-cache python3 wireguard-tools iproute2
COPY agent.py /app/agent.py
ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1
HEALTHCHECK --interval=15s --timeout=5s --retries=3 CMD ["python3", "/app/agent.py", "--healthcheck"]
ENTRYPOINT ["python3", "/app/agent.py"]
FROM base AS test
COPY test_agent.py /app/test_agent.py
FROM base AS runtime
+367
View File
@@ -0,0 +1,367 @@
"""Local, single-interface WireGuard controller. Never executes wg-quick/hooks."""
import base64
import binascii
import http.client
import http.server
import ipaddress
import json
import os
from pathlib import Path
import re
import socket
import socketserver
import subprocess
import sys
import tempfile
import threading
IFACE = "wg-flow"
CONFIG = Path(os.environ.get("FLOW_AGENT_CONFIG_DIR", "/etc/wireguard"))
SOCKET = os.environ.get("FLOW_AGENT_SOCKET", "/run/mm-netflow-host/agent.sock")
LOCK = threading.RLock()
MAX_BODY = 262144
class AgentError(Exception):
pass
def run(*args, input=None):
try:
result = subprocess.run(args, input=input, text=True, capture_output=True, timeout=10)
except (OSError, subprocess.TimeoutExpired):
raise AgentError("Команда WireGuard/IP недоступна или превысила время ожидания") from None
if result.returncode:
# stderr/stdout can contain keys and must never reach an HTTP response/log.
raise AgentError("Операция WireGuard/IP не выполнена")
return result.stdout.strip()
def key(value):
try:
if not isinstance(value, str) or len(value) != 44 or len(base64.b64decode(value, validate=True)) != 32:
raise ValueError()
except (ValueError, binascii.Error):
raise AgentError("Некорректный ключ WireGuard") from None
return value
def port(value):
if type(value) is not int or not 1 <= value <= 65535:
raise AgentError("Некорректный UDP-порт")
return value
def endpoint(value):
if not isinstance(value, str) or not re.fullmatch(r"(?:\[[0-9a-fA-F:]+\]|[a-zA-Z0-9._-]+):[0-9]{1,5}", value):
raise AgentError("Некорректный endpoint")
port(int(value.rsplit(":", 1)[1]))
return value
def parse_conf(text):
interface, peers, section = {}, [], None
allowed = {"interface": {"PrivateKey", "Address", "ListenPort", "MTU"},
"peer": {"PublicKey", "PresharedKey", "AllowedIPs", "Endpoint", "PersistentKeepalive"}}
for raw in text.splitlines():
line = raw.split("#", 1)[0].strip()
if not line:
continue
if line == "[Interface]" and section is None:
section, current = "interface", interface
elif line == "[Peer]" and section is not None:
section, current = "peer", {}
peers.append(current)
else:
name, sep, value = line.partition("=")
name, value = name.strip(), value.strip()
if not sep or section is None or name not in allowed[section] or name in current or not value:
raise AgentError("Конфигурация wg-flow содержит неподдерживаемые или повторные параметры; требуется ручная проверка")
current[name] = value
if not interface or not {"PrivateKey", "Address", "ListenPort"} <= interface.keys():
raise AgentError("Неполная конфигурация wg-flow")
key(interface["PrivateKey"])
try:
ipaddress.IPv4Interface(interface["Address"])
port(int(interface["ListenPort"]))
if "MTU" in interface and not 576 <= int(interface["MTU"]) <= 9000:
raise ValueError()
for peer in peers:
key(peer.get("PublicKey"))
if "PresharedKey" in peer:
key(peer["PresharedKey"])
if "Endpoint" in peer:
endpoint(peer["Endpoint"])
if peer.get("AllowedIPs"):
networks(peer["AllowedIPs"].split(","))
if "PersistentKeepalive" in peer and not 0 <= int(peer["PersistentKeepalive"]) <= 65535:
raise ValueError()
except (ValueError, TypeError):
raise AgentError("Некорректные параметры wg-flow") from None
if len({p["PublicKey"] for p in peers}) != len(peers):
raise AgentError("Повторный публичный ключ пира")
return interface, peers
def networks(values):
if not isinstance(values, list) or not values:
raise AgentError("AllowedIPs должны содержать адреса")
try:
return [ipaddress.ip_network(v.strip(), strict=False) for v in values]
except (ValueError, AttributeError):
raise AgentError("Некорректные AllowedIPs") from None
def validate(raw):
fields = {"collectorIp", "prefix", "flowListenPort", "wgListenPort", "hostPublicKey", "hostPrivateKey", "peers"}
if not isinstance(raw, dict) or set(raw) != fields:
raise AgentError("Некорректный снимок NetFlow")
try:
collector = ipaddress.IPv4Address(raw["collectorIp"])
prefix = ipaddress.IPv4Network(raw["prefix"], strict=False)
if prefix.prefixlen != 24 or collector not in prefix or collector in (prefix.network_address, prefix.broadcast_address):
raise ValueError()
except (ValueError, TypeError):
raise AgentError("Требуется IPv4 overlay /24 и адрес коллектора внутри него") from None
port(raw["flowListenPort"])
port(raw["wgListenPort"])
key(raw["hostPublicKey"])
key(raw["hostPrivateKey"])
if run("wg", "pubkey", input=raw["hostPrivateKey"] + "\n") != raw["hostPublicKey"]:
raise AgentError("Публичный и приватный ключи MM не соответствуют друг другу")
binding = json.loads((CONFIG / "binding.json").read_text())
for name in ("collectorIp", "prefix", "flowListenPort", "wgListenPort"):
if binding.get(name) != raw[name]:
raise AgentError("Адрес, префикс или порт отличаются от установки агента; повторите установщик")
if not isinstance(raw["peers"], list) or len(raw["peers"]) > 253:
raise AgentError("Некорректный список пиров")
peers, ids, keys = [], set(), set()
for peer in raw["peers"]:
if not isinstance(peer, dict) or not {"serverId", "publicKey", "allowedIps"} <= peer.keys() or set(peer) - {"serverId", "publicKey", "allowedIps", "endpoint"}:
raise AgentError("Некорректный пир")
peer_key = key(peer["publicKey"])
if type(peer["serverId"]) is not int or peer["serverId"] <= 0 or peer["serverId"] in ids or peer_key in keys or peer_key == raw["hostPublicKey"]:
raise AgentError("Повторный или некорректный пир")
ids.add(peer["serverId"])
keys.add(peer_key)
nets = networks(peer["allowedIps"])
if len(nets) != 1 or nets[0].version != 4 or nets[0].prefixlen != 32 or nets[0].network_address not in prefix or nets[0].network_address in (collector, prefix.network_address, prefix.broadcast_address):
raise AgentError("Пиру разрешён только собственный IPv4 /32 внутри overlay")
data = {"PublicKey": peer_key, "AllowedIPs": str(nets[0])}
if peer.get("endpoint"):
data.update(Endpoint=endpoint(peer["endpoint"]), PersistentKeepalive="25")
peers.append(data)
return peers
def merge_peers(existing, desired):
merged = {peer["PublicKey"]: dict(peer) for peer in existing}
for peer in desired:
public = peer["PublicKey"]
# Replacing a PSK would silently break a previously working tunnel.
if merged.get(public, {}).get("PresharedKey", "A" * 43 + "=") != "A" * 43 + "=":
raise AgentError("Управляемый пир использует preshared key; требуется ручная проверка")
merged[public] = dict(peer)
rows = list(merged.values())
wanted = {p["PublicKey"] for p in desired}
for i, left in enumerate(rows):
for right in rows[i + 1:]:
if left["PublicKey"] not in wanted and right["PublicKey"] not in wanted:
continue
if not left.get("AllowedIPs") or not right.get("AllowedIPs"):
continue
if any(a.version == b.version and a.overlaps(b) for a in networks(left["AllowedIPs"].split(",")) for b in networks(right["AllowedIPs"].split(","))):
raise AgentError("AllowedIPs управляемого пира пересекаются с другим пиром")
return rows
def render(interface, peers, native=False):
rows = ["[Interface]"]
rows += [f"{k} = {v}" for k, v in interface.items() if not native or k not in ("Address", "MTU")]
for peer in peers:
rows += ["", "[Peer]"] + [f"{k} = {v}" for k, v in peer.items() if k != "AllowedIPs" or v]
return "\n".join(rows) + "\n"
def atomic_write(path, text):
fd, name = tempfile.mkstemp(prefix=".flow-", dir=path.parent)
try:
with os.fdopen(fd, "w") as stream:
stream.write(text)
stream.flush()
os.fsync(stream.fileno())
os.replace(name, path)
if os.name == "posix":
directory = os.open(path.parent, os.O_RDONLY)
try:
os.fsync(directory)
finally:
os.close(directory)
finally:
if os.path.exists(name):
os.unlink(name)
def sync(text):
# wg requires a file. mkstemp creates it with mode 0600.
fd, name = tempfile.mkstemp(prefix="flow-sync-")
try:
with os.fdopen(fd, "w") as stream:
stream.write(text)
run("wg", "syncconf", IFACE, name)
finally:
os.unlink(name)
def live():
dump = run("wg", "show", IFACE, "dump").splitlines()
if not dump or len(dump[0].split("\t")) != 4:
raise AgentError("Интерфейс wg-flow недоступен")
interface = dump[0].split("\t")
peers, raw_peers = [], []
for line in dump[1:]:
cols = line.split("\t")
if len(cols) != 8:
raise AgentError("Некорректный ответ WireGuard")
public, psk, ep, ips, handshake, rx, tx, keepalive = cols
peers.append({"publicKey": public, "allowedIps": ips.split(",") if ips != "(none)" else [],
"lastHandshakeAt": int(handshake) or None, "receivedBytes": int(rx), "sentBytes": int(tx)})
item = {"PublicKey": public, "AllowedIPs": ips if ips != "(none)" else ""}
if psk != "(none)":
item["PresharedKey"] = psk
if ep != "(none)":
item["Endpoint"] = ep
item["PersistentKeepalive"] = keepalive
raw_peers.append(item)
addresses = json.loads(run("ip", "-j", "-4", "address", "show", "dev", IFACE))
cidrs = [f"{a['local']}/{a['prefixlen']}" for row in addresses for a in row.get("addr_info", []) if a.get("family") == "inet"]
return {"interfaceUp": any("UP" in row.get("flags", []) for row in addresses), "publicKey": interface[1], "listenPort": int(interface[2]), "addresses": cidrs, "peers": peers}, raw_peers
def check(raw, apply=False):
desired = validate(raw)
path = CONFIG / "wg-flow.conf"
previous = path.read_text()
interface, stored = parse_conf(previous)
expected_address = f"{raw['collectorIp']}/{raw['prefix'].split('/')[1]}"
if run("wg", "pubkey", input=interface["PrivateKey"] + "\n") != raw["hostPublicKey"] or interface["Address"] != expected_address or int(interface["ListenPort"]) != raw["wgListenPort"]:
raise AgentError("Ключ, адрес или порт существующего wg-flow отличаются от MM; требуется ручная проверка")
status, runtime = live()
if not status["interfaceUp"]:
raise AgentError("Интерфейс wg-flow выключен; проверьте wg-quick на хосте")
if status["publicKey"] != raw["hostPublicKey"] or expected_address not in status["addresses"] or status["listenPort"] != raw["wgListenPort"]:
raise AgentError("Работающий wg-flow отличается от конфигурации; требуется ручная проверка")
# Preserve both persistent and runtime-only third-party peers. Detect conflicts before writing.
combined = {p["PublicKey"]: p for p in runtime}
combined.update({p["PublicKey"]: p for p in stored})
merged = merge_peers(list(combined.values()), desired)
runtime_merged = merge_peers(runtime, desired)
# Do not reset a third-party peer's runtime settings to an older persistent file.
# Preserve roaming endpoints of managed peers unless the persisted endpoint was edited.
stored_by_key = {p["PublicKey"]: p for p in stored}
runtime_by_key = {p["PublicKey"]: p for p in runtime}
desired_by_key = {p["PublicKey"]: p for p in desired}
for peer in runtime_merged:
public = peer["PublicKey"]
if public in desired_by_key and stored_by_key.get(public, {}).get("Endpoint") == desired_by_key[public].get("Endpoint"):
if runtime_by_key.get(public, {}).get("Endpoint"):
peer["Endpoint"] = runtime_by_key[public]["Endpoint"]
def matches(actual, persistent=False):
by_key = {p["PublicKey"]: p for p in actual}
return all(p["PublicKey"] in by_key and all(by_key[p["PublicKey"]].get(k) == v for k, v in p.items() if persistent or k != "Endpoint") for p in desired)
matched = matches(stored, persistent=True) and matches(runtime)
if apply:
rollback = run("wg", "showconf", IFACE)
try:
atomic_write(path, render(interface, merged))
# Keep the running interface header verbatim (including a possible fwmark).
header = rollback.split("[Peer]", 1)[0].rstrip()
peer_config = render({}, runtime_merged, native=True).removeprefix("[Interface]\n")
sync(header + "\n" + peer_config)
status, after = live()
if not matches(after):
raise AgentError("Применение пиров не подтверждено")
except Exception:
failures = []
for action in (lambda: atomic_write(path, previous), lambda: sync(rollback)):
try:
action()
except Exception:
failures.append(True)
raise AgentError("Применение не удалось; " + ("восстановление требует ручной проверки" if failures else "предыдущая конфигурация восстановлена")) from None
matched = True
return {**status, "configMatches": matched}
class Handler(http.server.BaseHTTPRequestHandler):
def log_message(self, *_):
pass
def handle_request(self):
try:
with LOCK:
if self.command == "GET" and self.path == "/status":
result, _ = live()
result["configMatches"] = None
elif self.command == "POST" and self.path in ("/check", "/apply"):
size = int(self.headers.get("Content-Length", "0"))
if not 0 < size <= MAX_BODY or self.headers.get("Transfer-Encoding"):
raise AgentError("Некорректный размер запроса")
raw = json.loads(self.rfile.read(size))
result = check(raw, apply=self.path == "/apply")
else:
self.respond(404, {"error": "Операция не поддерживается"})
return
self.respond(200, result)
except AgentError as error:
self.respond(409, {"error": str(error)})
except Exception:
self.respond(503, {"error": "Агент NetFlow не готов; проверьте конфигурацию и интерфейс хоста"})
def setup(self):
super().setup()
self.connection.settimeout(15)
def respond(self, code, body):
data = json.dumps(body, ensure_ascii=False).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json; charset=utf-8")
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
do_GET = handle_request
do_POST = handle_request
class Server(socketserver.ThreadingMixIn, getattr(socketserver, "UnixStreamServer", socketserver.TCPServer)):
daemon_threads = True
class UnixConnection(http.client.HTTPConnection):
def connect(self):
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
self.sock.settimeout(3)
self.sock.connect(SOCKET)
def main():
if not hasattr(socketserver, "UnixStreamServer"):
raise SystemExit("NetFlow agent requires Linux Unix sockets")
if "--healthcheck" in sys.argv:
client = UnixConnection("localhost")
client.request("GET", "/status")
response = client.getresponse()
sys.exit(0 if response.status == 200 and json.loads(response.read()).get("interfaceUp") else 1)
os.umask(0o077)
Path(SOCKET).parent.mkdir(parents=True, exist_ok=True)
if os.path.lexists(SOCKET):
Path(SOCKET).unlink()
with Server(SOCKET, Handler) as server:
os.chmod(SOCKET, 0o660)
print("NetFlow agent: Unix API ready", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()
+154
View File
@@ -0,0 +1,154 @@
#!/usr/bin/env bash
# Usage: sudo bash install.sh /absolute/path/wg-flow.conf /opt/cdn-mm 4739
set -euo pipefail
umask 077
[[ ${EUID:-$(id -u)} -eq 0 ]] || { echo "Запустите от root" >&2; exit 1; }
SOURCE=$(realpath "${1:?Укажите скачанный wg-flow.conf}")
COMPOSE_DIR=$(realpath "${2:-/opt/cdn-mm}")
FLOW_PORT=${3:-4739}
[[ "$FLOW_PORT" =~ ^[0-9]+$ ]] && (( FLOW_PORT >= 1 && FLOW_PORT <= 65535 )) || exit 1
command -v docker >/dev/null
docker compose version >/dev/null
systemctl show docker.service --property=LoadState --value | grep -qx loaded
if ! command -v wg >/dev/null || ! command -v python3 >/dev/null || ! command -v modprobe >/dev/null; then
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y wireguard-tools python3 iproute2 kmod
fi
modprobe wireguard
[[ -f "$COMPOSE_DIR/docker-compose.yml" ]] || { echo "Не найден docker-compose.yml" >&2; exit 1; }
AGENT_IMAGE=${MM_NETFLOW_AGENT_IMAGE:-git.shx.one/denozord/mikrotikmanager-netflow-agent:latest}
[[ "$AGENT_IMAGE" =~ ^[A-Za-z0-9][A-Za-z0-9._:/@-]+$ ]] || { echo "Некорректное имя образа" >&2; exit 1; }
# The image must be available before changing the host configuration.
docker pull "$AGENT_IMAGE"
TASK_TEMP=$(mktemp -d)
trap 'rm -rf -- "$TASK_TEMP"' EXIT
python3 - "$SOURCE" "$COMPOSE_DIR" "$FLOW_PORT" "$TASK_TEMP" <<'PY'
import base64, ipaddress, json, os, pathlib, re, subprocess, sys
source, directory, port, temp = pathlib.Path(sys.argv[1]), pathlib.Path(sys.argv[2]), int(sys.argv[3]), pathlib.Path(sys.argv[4])
def read(text):
section, fields, peers = None, {}, []
allowed = {'Interface': {'PrivateKey','Address','ListenPort','MTU'}, 'Peer': {'PublicKey','AllowedIPs','Endpoint','PersistentKeepalive','PresharedKey'}}
for raw in text.splitlines():
line = raw.split('#',1)[0].strip()
if not line: continue
if line == '[Interface]' and section is None: section, current = 'Interface', fields
elif line == '[Peer]' and section is not None:
section, current = 'Peer', {}; peers.append(current)
else:
key, sep, value = line.partition('='); key, value = key.strip(), value.strip()
if not sep or section not in allowed or key not in allowed[section] or key in current or not value:
raise SystemExit('Конфигурация содержит неподдерживаемые параметры; требуется ручная проверка')
current[key] = value
try:
addr = ipaddress.IPv4Interface(fields['Address'])
if addr.network.prefixlen != 24 or addr.ip in (addr.network.network_address, addr.network.broadcast_address): raise ValueError()
if len(base64.b64decode(fields['PrivateKey'], validate=True)) != 32 or not 1 <= int(fields['ListenPort']) <= 65535: raise ValueError()
if 'MTU' in fields and not 576 <= int(fields['MTU']) <= 9000: raise ValueError()
for peer in peers:
if len(base64.b64decode(peer['PublicKey'], validate=True)) != 32: raise ValueError()
for network in peer['AllowedIPs'].split(','): ipaddress.ip_network(network.strip(), strict=False)
if 'Endpoint' in peer and not re.fullmatch(r'(?:\[[0-9a-fA-F:]+\]|[a-zA-Z0-9._-]+):[0-9]{1,5}', peer['Endpoint']): raise ValueError()
if 'PersistentKeepalive' in peer and not 0 <= int(peer['PersistentKeepalive']) <= 65535: raise ValueError()
except (ValueError, KeyError): raise SystemExit('Некорректная конфигурация WireGuard; ключи не изменены')
return fields, addr
expected, addr = read(source.read_text())
active = pathlib.Path('/etc/wireguard/wg-flow.conf')
candidate = active.read_text() if active.exists() else source.read_text()
actual, actual_addr = read(candidate)
for key in ('PrivateKey','Address','ListenPort'):
if actual[key] != expected[key]: raise SystemExit('Ключ, адрес или порт существующего туннеля отличаются от MM; установка остановлена')
public = subprocess.run(['wg','pubkey'], input=actual['PrivateKey']+'\n',text=True,capture_output=True,check=True).stdout.strip()
if subprocess.run(['ip','link','show','dev','wg-flow'],capture_output=True).returncode == 0:
dump = subprocess.run(['wg','show','wg-flow','dump'],capture_output=True,text=True,check=True).stdout.splitlines()
head = dump[0].split('\t')
addresses = json.loads(subprocess.run(['ip','-j','-4','address','show','dev','wg-flow'],capture_output=True,text=True,check=True).stdout)
if head[1] != public or int(head[2]) != int(actual['ListenPort']) or not any(a.get('local') == str(addr.ip) and a.get('prefixlen') == addr.network.prefixlen for row in addresses for a in row.get('addr_info',[])):
raise SystemExit('Работающий wg-flow отличается от MM; установка остановлена')
binding = dict(collectorIp=str(addr.ip), prefix=str(addr.network), flowListenPort=port, wgListenPort=int(actual['ListenPort']))
(temp/'binding.json').write_text(json.dumps(binding))
(temp/'wg-flow.conf').write_text(candidate)
(temp/'collector').write_text(str(addr.ip))
PY
COLLECTOR_IP=$(cat "$TASK_TEMP/collector")
# Keep a dedicated override; never replace docker-compose.override.yml.
cat > "$TASK_TEMP/compose.yml" <<EOF
services:
netflow-agent:
image: "$AGENT_IMAGE"
container_name: mmapp-netflow-agent
restart: unless-stopped
network_mode: host
cap_drop: [ALL]
cap_add: [NET_ADMIN]
security_opt: ["no-new-privileges:true"]
read_only: true
tmpfs: ["/tmp:rw,noexec,nosuid,size=8m"]
volumes:
- ./data/netflow-wireguard:/etc/wireguard
- ./data/netflow-agent-socket:/run/mm-netflow-host
logging:
driver: json-file
options: {max-size: "5m", max-file: "2"}
backend:
environment:
FLOW_AGENT_SOCKET: /run/mm-netflow-host/agent.sock
FLOW_LISTEN_HOST: "0.0.0.0"
volumes:
- ./data/netflow-agent-socket:/run/mm-netflow-host:ro
ports:
- "$COLLECTOR_IP:$FLOW_PORT:$FLOW_PORT/udp"
depends_on:
netflow-agent:
condition: service_healthy
EOF
COMPOSE_ARGS=(-f "$COMPOSE_DIR/docker-compose.yml")
if [[ -f "$COMPOSE_DIR/docker-compose.override.yml" ]]; then COMPOSE_ARGS+=(-f "$COMPOSE_DIR/docker-compose.override.yml"); fi
# Render first, and reject an existing public/incorrect IPFIX binding before touching WireGuard.
docker compose --project-directory "$COMPOSE_DIR" "${COMPOSE_ARGS[@]}" -f "$TASK_TEMP/compose.yml" config --format json > "$TASK_TEMP/merged.json"
python3 - "$TASK_TEMP/merged.json" "$COLLECTOR_IP" "$FLOW_PORT" <<'PY'
import json,sys
config=json.load(open(sys.argv[1])); ip=sys.argv[2]; port=int(sys.argv[3])
for bind in config['services']['backend'].get('ports',[]):
if bind.get('protocol','tcp') == 'udp' and (bind['target'] in (4739,port) or str(bind.get('published')) == str(port)):
if bind.get('host_ip') != ip or bind['target'] != port or str(bind.get('published')) != str(port):
raise SystemExit('В Compose уже есть другой IPFIX bind. Согласуйте его с адресом и портом MM и повторите установку.')
PY
if docker container inspect mmapp-netflow-agent >/dev/null 2>&1; then docker stop mmapp-netflow-agent; fi
install -d -m 700 /etc/wireguard "$COMPOSE_DIR/data/netflow-wireguard" "$COMPOSE_DIR/data/netflow-agent-socket"
if [[ -e /etc/wireguard/wg-flow.conf ]]; then
cp -L /etc/wireguard/wg-flow.conf "$COMPOSE_DIR/data/netflow-wireguard/wg-flow.conf.backup.$(date +%s)"
fi
install -m 600 "$TASK_TEMP/wg-flow.conf" "$COMPOSE_DIR/data/netflow-wireguard/wg-flow.conf"
install -m 600 "$TASK_TEMP/binding.json" "$COMPOSE_DIR/data/netflow-wireguard/binding.json"
ln -sfn "$COMPOSE_DIR/data/netflow-wireguard/wg-flow.conf" /etc/wireguard/wg-flow.conf
install -m 600 "$TASK_TEMP/compose.yml" "$COMPOSE_DIR/docker-compose.netflow-agent.yml"
install -d /etc/systemd/system/docker.service.d
cat > /etc/systemd/system/docker.service.d/mm-netflow.conf <<'EOF'
[Unit]
Wants=wg-quick@wg-flow.service
After=wg-quick@wg-flow.service
EOF
systemctl daemon-reload
systemctl enable wg-quick@wg-flow.service
if ! ip link show dev wg-flow >/dev/null 2>&1; then
systemctl start wg-quick@wg-flow.service
else
ip link set dev wg-flow up
fi
cd "$COMPOSE_DIR"
docker compose "${COMPOSE_ARGS[@]}" -f docker-compose.netflow-agent.yml up -d netflow-agent
AGENT_READY=false
for attempt in $(seq 1 60); do
if [[ "$(docker inspect --format '{{.State.Health.Status}}' mmapp-netflow-agent)" == healthy ]]; then
AGENT_READY=true
break
fi
sleep 1
done
[[ "$AGENT_READY" == true ]] || { echo "Агент не готов. Проверьте docker logs mmapp-netflow-agent; backend не пересоздан." >&2; exit 1; }
docker compose "${COMPOSE_ARGS[@]}" -f docker-compose.netflow-agent.yml up -d --no-deps --force-recreate backend
echo "Агент установлен. Используйте те же -f файлы при дальнейших операциях Compose:"
printf 'docker compose '; printf '%q ' "${COMPOSE_ARGS[@]}" -f "$COMPOSE_DIR/docker-compose.netflow-agent.yml"; echo 'up -d'
echo "В MM обновите статус хоста и нажмите «Применить настройки хоста»."
+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()