最終更新 1 month ago

nuno revised this gist 1 month ago. Go to revision

1 file changed, 188 insertions

gistfile1.txt(file created)

@@ -0,0 +1,188 @@
1 + #!/usr/bin/env python3
2 +
3 + import json
4 + import logging
5 + import os
6 + import signal
7 + import sys
8 + import time
9 + from datetime import datetime, timezone
10 + from pathlib import Path
11 +
12 + from scapy.all import IP, UDP, sniff
13 +
14 + # ============================================
15 + # Config
16 + # ============================================
17 +
18 + WATCHLIST_FILE = Path("/etc/wireguard/.wgctl/daemon/watchlist.json")
19 + EVENTS_LOG = Path("/etc/wireguard/.wgctl/daemon/events.log")
20 + WG_INTERFACE = os.environ.get("WG_INTERFACE", "eth0")
21 + WG_PORT = int(os.environ.get("WG_PORT", "51820"))
22 + LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO")
23 +
24 + # ============================================
25 + # Logging
26 + # ============================================
27 +
28 + logging.basicConfig(
29 + level=getattr(logging, LOG_LEVEL),
30 + format="%(asctime)s [%(levelname)s] %(message)s",
31 + handlers=[logging.StreamHandler(sys.stdout)]
32 + )
33 + log = logging.getLogger("wgctl-monitor")
34 +
35 + # ============================================
36 + # Watchlist
37 + # ============================================
38 +
39 + _watchlist: dict[str, str] = {}
40 + _watchlist_mtime: float = 0.0
41 +
42 + def load_watchlist() -> dict[str, str]:
43 + global _watchlist, _watchlist_mtime
44 +
45 + try:
46 + mtime = WATCHLIST_FILE.stat().st_mtime
47 + if mtime == _watchlist_mtime:
48 + return _watchlist
49 +
50 + with WATCHLIST_FILE.open() as f:
51 + _watchlist = json.load(f)
52 + _watchlist_mtime = mtime
53 + log.debug(f"Watchlist reloaded: {len(_watchlist)} entries")
54 +
55 + except Exception as e:
56 + log.error(f"Failed to load watchlist: {e}")
57 +
58 + return _watchlist
59 +
60 + def is_watched(ip: str) -> str | None:
61 + watchlist = load_watchlist()
62 + return watchlist.get(ip)
63 +
64 + # ============================================
65 + # Endpoint Resolution
66 + # ============================================
67 +
68 + def get_endpoint(public_key: str) -> str | None:
69 + try:
70 + import subprocess
71 + result = subprocess.run(
72 + ["wg", "show", WG_INTERFACE, "endpoints"],
73 + capture_output=True, text=True
74 + )
75 + for line in result.stdout.splitlines():
76 + parts = line.split()
77 + if len(parts) == 2 and parts[0] == public_key:
78 + # Return just the IP without port
79 + return parts[1].rsplit(":", 1)[0]
80 + except Exception as e:
81 + log.debug(f"Failed to get endpoint: {e}")
82 + return None
83 +
84 + def get_client_public_key(client_name: str) -> str | None:
85 + key_file = Path(f"/etc/wireguard/clients/{client_name}_public.key")
86 + try:
87 + return key_file.read_text().strip()
88 + except Exception:
89 + return None
90 +
91 + # ============================================
92 + # Event Logging
93 + # ============================================
94 +
95 + def log_event(ip: str, client: str, event: str, endpoint: str | None = None):
96 + entry = {
97 + "timestamp": datetime.now(timezone.utc).isoformat(),
98 + "ip": ip,
99 + "client": client,
100 + "event": event,
101 + }
102 +
103 + # Update endpoint cache when we see a packet
104 + cache_file = os.path.join(os.path.dirname(WATCHLIST_FILE), 'endpoint_cache.json')
105 + try:
106 + with open(cache_file) as f:
107 + cache = json.load(f)
108 + except:
109 + cache = {}
110 + cache[client] = ip
111 + with open(cache_file, 'w') as f:
112 + json.dump(cache, f, indent=2)
113 +
114 + if endpoint:
115 + entry["endpoint"] = endpoint
116 +
117 + try:
118 + with EVENTS_LOG.open("a") as f:
119 + f.write(json.dumps(entry) + "\n")
120 + log.debug(f"Event logged: {entry}")
121 + except Exception as e:
122 + log.error(f"Failed to write event: {e}")
123 +
124 + # ============================================
125 + # Packet Handler
126 + # ============================================
127 +
128 + def handle_packet(pkt):
129 + if not (IP in pkt and UDP in pkt):
130 + return
131 +
132 + # Only care about packets targeting WireGuard port
133 + if pkt[UDP].dport != WG_PORT:
134 + return
135 +
136 + src_ip = pkt[IP].src
137 + client = is_watched(src_ip)
138 +
139 + if not client:
140 + return
141 +
142 + # Resolve real endpoint IP
143 + public_key = get_client_public_key(client)
144 + endpoint = None
145 + if public_key:
146 + endpoint = get_endpoint(public_key)
147 +
148 + # If no endpoint from wg show, use packet source IP
149 + if not endpoint:
150 + endpoint = src_ip
151 +
152 + log_event(src_ip, client, "attempt", endpoint)
153 + log.info(f"Blocked attempt: {client} ({src_ip}) from endpoint {endpoint}")
154 +
155 + # ============================================
156 + # Signal Handling
157 + # ============================================
158 +
159 + def handle_signal(signum, frame):
160 + log.info("Shutting down wgctl-monitor")
161 + sys.exit(0)
162 +
163 + signal.signal(signal.SIGTERM, handle_signal)
164 + signal.signal(signal.SIGINT, handle_signal)
165 +
166 + # ============================================
167 + # Main
168 + # ============================================
169 +
170 + def main():
171 + log.info(f"wgctl-monitor starting on interface {WG_INTERFACE} port {WG_PORT}")
172 +
173 + if not WATCHLIST_FILE.exists():
174 + log.error(f"Watchlist not found: {WATCHLIST_FILE}")
175 + sys.exit(1)
176 +
177 + load_watchlist()
178 + log.info("Watchlist loaded, starting packet capture...")
179 +
180 + sniff(
181 + iface=WG_INTERFACE,
182 + filter=f"udp port {WG_PORT}",
183 + prn=handle_packet,
184 + store=0
185 + )
186 +
187 + if __name__ == "__main__":
188 + main()
Newer Older