import asyncio import json import logging import os import time from contextlib import asynccontextmanager from collections import defaultdict import geoip2.database import geoip2.errors import httpx from dotenv import load_dotenv from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware load_dotenv() logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger(__name__) # ── Config ──────────────────────────────────────────────────────────────────── UNIFI_URL = os.getenv("UNIFI_URL", "https://192.168.1.1") UNIFI_API_KEY = os.getenv("UNIFI_API_KEY", "") UNIFI_SITE = os.getenv("UNIFI_SITE", "default") UNIFI_VERIFY_SSL = os.getenv("UNIFI_VERIFY_SSL", "false").lower() == "true" GEOIP_DB_PATH = os.getenv("GEOIP_DB_PATH", "/data/GeoLite2-City.mmdb") POLL_INTERVAL = int(os.getenv("POLL_INTERVAL", "10")) HOME_LAT = float(os.getenv("HOME_LAT", "59.9139")) HOME_LON = float(os.getenv("HOME_LON", "10.7522")) HOME_NAME = os.getenv("HOME_NAME", "Home") NETFLOW_PORT = int(os.getenv("NETFLOW_PORT", "2055")) CROWDSEC_URL = os.getenv("CROWDSEC_URL", "") CROWDSEC_API_KEY = os.getenv("CROWDSEC_API_KEY", "") # ── State ───────────────────────────────────────────────────────────────────── ws_clients: dict[int, WebSocket] = {} geo_reader: geoip2.database.Reader | None = None device_registry: dict[str, dict] = {} recent_flows: list[dict] = [] MAX_RECENT = 500 ip_to_mac: dict[str, str] = {} netflow_stats = {"packets": 0, "flows": 0, "geolocated": 0, "templates": 0} # ── GeoIP ───────────────────────────────────────────────────────────────────── def load_geoip(): global geo_reader if not os.path.exists(GEOIP_DB_PATH): log.warning(f"GeoIP database not found at {GEOIP_DB_PATH}") return try: geo_reader = geoip2.database.Reader(GEOIP_DB_PATH) log.info("GeoIP database loaded.") except Exception as e: log.error(f"Failed to load GeoIP: {e}") def geolocate(ip: str) -> dict | None: if not geo_reader: return None try: r = geo_reader.city(ip) if r.location.latitude is None: return None return {"lat": r.location.latitude, "lon": r.location.longitude, "city": r.city.name or "", "country": r.country.name or "", "country_code": r.country.iso_code or ""} except (geoip2.errors.AddressNotFoundError, Exception): return None def is_private_ip(ip: str) -> bool: import ipaddress try: addr = ipaddress.ip_address(ip) return addr.is_private or addr.is_loopback or addr.is_link_local or addr.is_multicast except ValueError: return True def int_to_ip(n: int) -> str: return f"{(n >> 24) & 0xff}.{(n >> 16) & 0xff}.{(n >> 8) & 0xff}.{n & 0xff}" # ── NetFlow v5 parser ───────────────────────────────────────────────────────── import struct NF5_HEADER = struct.Struct("!HHIIIhHBB") # 24 bytes NF5_RECORD = struct.Struct("!IIIHHIIIIHHBBBBHHBBH") # 48 bytes def parse_netflow_v5(data: bytes) -> list[dict]: if len(data) < 24: return [] hdr = NF5_HEADER.unpack_from(data, 0) count = hdr[1] flows = [] for i in range(count): off = 24 + i * 48 if off + 48 > len(data): break r = NF5_RECORD.unpack_from(data, off) src_ip = int_to_ip(r[0]) dst_ip = int_to_ip(r[1]) packets = r[6] octets = r[7] proto = r[14] src_port = r[10] dst_port = r[11] flows.append({"src": src_ip, "dst": dst_ip, "bytes": octets, "packets": packets, "proto": proto, "src_port": src_port, "dst_port": dst_port}) return flows # ── NetFlow v9 parser ───────────────────────────────────────────────────────── # Store templates per (src_addr, source_id, template_id) nf9_templates: dict[tuple, dict] = {} NF9_FIELD_TYPES = { 1: ("IN_BYTES", 4), 2: ("IN_PKTS", 4), 4: ("PROTOCOL", 1), 7: ("L4_SRC_PORT", 2), 8: ("IPV4_SRC_ADDR", 4), 11: ("L4_DST_PORT", 2), 12: ("IPV4_DST_ADDR", 4), 16: ("SRC_AS", 4), 17: ("DST_AS", 4), 21: ("LAST_SWITCHED", 4), 22: ("FIRST_SWITCHED", 4), 23: ("OUT_BYTES", 4), 24: ("OUT_PKTS", 4), 32: ("ICMP_TYPE", 2), 60: ("IP_PROTOCOL_VERSION", 1), } def parse_nf9_template(data: bytes, offset: int, length: int, src: str, source_id: int): end = offset + length pos = offset while pos + 4 <= end: tmpl_id = struct.unpack_from("!H", data, pos)[0] field_count = struct.unpack_from("!H", data, pos + 2)[0] pos += 4 fields = [] for _ in range(field_count): if pos + 4 > end: break ftype = struct.unpack_from("!H", data, pos)[0] flen = struct.unpack_from("!H", data, pos + 2)[0] fields.append((ftype, flen)) pos += 4 key = (src, source_id, tmpl_id) nf9_templates[key] = fields netflow_stats["templates"] += 1 log.debug(f"Stored NF9 template {tmpl_id} from {src} ({len(fields)} fields)") def parse_nf9_data(data: bytes, offset: int, length: int, src: str, source_id: int, flowset_id: int) -> list[dict]: key = (src, source_id, flowset_id) template = nf9_templates.get(key) if not template: return [] record_len = sum(f[1] for f in template) if record_len == 0: return [] flows = [] pos = offset end = offset + length while pos + record_len <= end: record = {} rpos = pos for ftype, flen in template: raw = data[rpos:rpos + flen] rpos += flen name = NF9_FIELD_TYPES.get(ftype, (None, None))[0] if name == "IPV4_SRC_ADDR" and flen == 4: record["src"] = int_to_ip(struct.unpack("!I", raw)[0]) elif name == "IPV4_DST_ADDR" and flen == 4: record["dst"] = int_to_ip(struct.unpack("!I", raw)[0]) elif name == "IN_BYTES" and flen <= 4: record["bytes"] = int.from_bytes(raw, "big") elif name == "IN_PKTS" and flen <= 4: record["packets"] = int.from_bytes(raw, "big") elif name == "PROTOCOL" and flen == 1: record["proto"] = raw[0] elif name == "L4_SRC_PORT" and flen == 2: record["src_port"] = struct.unpack("!H", raw)[0] elif name == "L4_DST_PORT" and flen == 2: record["dst_port"] = struct.unpack("!H", raw)[0] if "src" in record and "dst" in record: flows.append(record) pos += record_len return flows def parse_netflow_v9(data: bytes, src_addr: str) -> list[dict]: if len(data) < 20: return [] count = struct.unpack_from("!H", data, 2)[0] source_id = struct.unpack_from("!I", data, 16)[0] pos = 20 all_flows = [] for _ in range(count): if pos + 4 > len(data): break flowset_id = struct.unpack_from("!H", data, pos)[0] length = struct.unpack_from("!H", data, pos + 2)[0] if length < 4: break payload_off = pos + 4 payload_len = length - 4 if flowset_id == 0: # Template flowset parse_nf9_template(data, payload_off, payload_len, src_addr, source_id) elif flowset_id == 1: pass # Options template, skip elif flowset_id >= 256: # Data flowset flows = parse_nf9_data(data, payload_off, payload_len, src_addr, source_id, flowset_id) all_flows.extend(flows) pos += length # Align to 4-byte boundary if pos % 4: pos += 4 - (pos % 4) return all_flows # ── Flow processing ─────────────────────────────────────────────────────────── # Deduplicate noisy flows: track recent (src,dst) pairs recent_pairs: dict[tuple, float] = {} DEDUP_WINDOW = 30 # seconds — don't re-emit same src→dst within this window def process_raw_flow(raw: dict) -> dict | None: src = raw.get("src", "") dst = raw.get("dst", "") if not src or not dst: return None src_local = is_private_ip(src) dst_local = is_private_ip(dst) # Skip purely local flows if src_local and dst_local: return None now_ts = time.time() now_ms = int(now_ts * 1000) if src_local and not dst_local: # Outbound pair_key = ("out", src, dst) if now_ts - recent_pairs.get(pair_key, 0) < DEDUP_WINDOW: return None recent_pairs[pair_key] = now_ts geo = geolocate(dst) mac = ip_to_mac.get(src, "") name = device_registry.get(mac, {}).get("name") or src return {"id": f"nf-{src}-{dst}-{now_ms}", "type": "outbound", "device_mac": mac, "device_name": name, "device_ip": src, "bytes": raw.get("bytes", 0), "packets": raw.get("packets", 0), "proto": raw.get("proto", 0), "src_port": raw.get("src_port", 0), "dst_port": raw.get("dst_port", 0), "remote_ip": dst, "src_lat": HOME_LAT, "src_lon": HOME_LON, "timestamp": now_ms, "remote_geo": geo} elif not src_local and dst_local: # Inbound pair_key = ("in", src, dst) if now_ts - recent_pairs.get(pair_key, 0) < DEDUP_WINDOW: return None recent_pairs[pair_key] = now_ts geo = geolocate(src) mac = ip_to_mac.get(dst, "") name = device_registry.get(mac, {}).get("name") or dst return {"id": f"nf-{src}-{dst}-{now_ms}", "type": "inbound", "device_mac": mac, "device_name": name, "device_ip": dst, "bytes": raw.get("bytes", 0), "packets": raw.get("packets", 0), "proto": raw.get("proto", 0), "src_port": raw.get("src_port", 0), "dst_port": raw.get("dst_port", 0), "remote_ip": src, "src_lat": HOME_LAT, "src_lon": HOME_LON, "timestamp": now_ms, "remote_geo": geo} return None # ── NetFlow UDP server ──────────────────────────────────────────────────────── class NetFlowProtocol(asyncio.DatagramProtocol): def __init__(self): self._loop = asyncio.get_event_loop() def datagram_received(self, data: bytes, addr): netflow_stats["packets"] += 1 src_addr = addr[0] try: if len(data) < 2: return version = struct.unpack_from("!H", data, 0)[0] if version == 5: raw_flows = parse_netflow_v5(data) elif version == 9: raw_flows = parse_netflow_v9(data, src_addr) else: log.debug(f"Unknown NetFlow version {version}") return good_flows = [] for raw in raw_flows: flow = process_raw_flow(raw) if flow: good_flows.append(flow) netflow_stats["flows"] += 1 if flow.get("remote_geo"): netflow_stats["geolocated"] += 1 if good_flows: self._loop.create_task(emit_flows(good_flows)) except Exception as e: log.error(f"NetFlow parse error: {e}", exc_info=True) async def emit_flows(flows: list[dict]): recent_flows.extend(flows) if len(recent_flows) > MAX_RECENT: del recent_flows[:-MAX_RECENT] await broadcast({"type": "flows", "flows": flows}) async def start_netflow_server(): loop = asyncio.get_event_loop() try: await loop.create_datagram_endpoint( NetFlowProtocol, local_addr=("0.0.0.0", NETFLOW_PORT) ) log.info(f"NetFlow UDP collector listening on port {NETFLOW_PORT}") except Exception as e: log.error(f"Failed to start NetFlow listener: {e}") # ── UniFi API ───────────────────────────────────────────────────────────────── def make_unifi_client() -> httpx.AsyncClient: return httpx.AsyncClient( base_url=UNIFI_URL, headers={"X-API-KEY": UNIFI_API_KEY, "Accept": "application/json"}, verify=UNIFI_VERIFY_SSL, timeout=10.0, ) async def fetch_clients(http: httpx.AsyncClient) -> list[dict]: try: r = await http.get(f"/proxy/network/api/s/{UNIFI_SITE}/stat/sta") r.raise_for_status() return r.json().get("data", []) except Exception as e: log.error(f"Error fetching clients: {e}") return [] async def fetch_known_clients(http: httpx.AsyncClient) -> list[dict]: try: r = await http.get(f"/proxy/network/api/s/{UNIFI_SITE}/rest/user") r.raise_for_status() return r.json().get("data", []) except Exception as e: log.error(f"Error fetching known clients: {e}") return [] # ── Broadcast ───────────────────────────────────────────────────────────────── async def broadcast(msg: dict): dead = [] for cid, ws in ws_clients.items(): try: await ws.send_text(json.dumps(msg)) except Exception: dead.append(cid) for cid in dead: ws_clients.pop(cid, None) # ── Poll loop ───────────────────────────────────────────────────────────────── async def poll_loop(): log.info("Starting UniFi poll loop...") async with make_unifi_client() as http: known = await fetch_known_clients(http) for k in known: mac = k.get("mac", "") if mac: device_registry[mac] = { "mac": mac, "name": k.get("name") or k.get("hostname") or mac, "oui": k.get("oui", ""), "ip": k.get("last_ip", ""), "tx_bytes": 0, "rx_bytes": 0, } log.info(f"Loaded {len(device_registry)} known devices.") while True: try: active = await fetch_clients(http) for sta in active: mac = sta.get("mac", "") if not mac: continue ip = sta.get("ip", "") if ip: ip_to_mac[ip] = mac tx = sta.get("wired-tx_bytes") or sta.get("tx_bytes") or 0 rx = sta.get("wired-rx_bytes") or sta.get("rx_bytes") or 0 device_registry[mac] = { "mac": mac, "name": (sta.get("name") or sta.get("hostname") or device_registry.get(mac, {}).get("name") or mac), "ip": ip, "oui": sta.get("oui", ""), "is_wired": sta.get("is_wired", False), "uptime": sta.get("uptime"), "signal": sta.get("signal"), "tx_bytes": tx, "rx_bytes": rx, "last_seen": sta.get("last_seen"), } await broadcast({"type": "devices", "devices": list(device_registry.values())}) except Exception as e: log.error(f"Poll loop error: {e}", exc_info=True) await asyncio.sleep(POLL_INTERVAL) # ── Cleanup task ───────────────────────────────────────────────────────────── async def cleanup_loop(): """Periodically clean up the dedup table so it doesn't grow forever.""" while True: await asyncio.sleep(120) cutoff = time.time() - DEDUP_WINDOW * 2 stale = [k for k, t in recent_pairs.items() if t < cutoff] for k in stale: recent_pairs.pop(k, None) # ── FastAPI ─────────────────────────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): load_geoip() await start_netflow_server() poll_task = asyncio.create_task(poll_loop()) cleanup_task = asyncio.create_task(cleanup_loop()) yield poll_task.cancel() cleanup_task.cancel() app = FastAPI(lifespan=lifespan) app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]) @app.get("/api/devices") async def get_devices(): return {"devices": list(device_registry.values())} @app.get("/api/status") async def get_status(): return { "geoip_loaded": geo_reader is not None, "device_count": len(device_registry), "home": {"lat": HOME_LAT, "lon": HOME_LON, "name": HOME_NAME}, "netflow_port": NETFLOW_PORT, "netflow_stats": netflow_stats, "templates_learned": len(nf9_templates), } @app.get("/api/crowdsec/enabled") async def get_crowdsec_enabled(): """Returns whether CrowdSec is configured.""" return {"enabled": bool(CROWDSEC_URL and CROWDSEC_API_KEY)} @app.get("/api/crowdsec") async def get_crowdsec(origin: str = "all"): """Proxy CrowdSec decisions with geo enrichment. origin=all|local|capi""" if not CROWDSEC_URL or not CROWDSEC_API_KEY: return [] try: async with httpx.AsyncClient(timeout=8.0) as http: # Fetch all decisions — filter locally since CrowdSec API origin param is unreliable r = await http.get( f"{CROWDSEC_URL}/v1/decisions", headers={"X-Api-Key": CROWDSEC_API_KEY}, params={"limit": 2000} ) if r.status_code == 204: return [] r.raise_for_status() decisions = r.json() if not decisions: return [] # Filter by origin locally — handles any casing CrowdSec uses if origin == "local": decisions = [d for d in decisions if d.get("origin", "").lower() not in ("capi", "lists")] elif origin == "capi": decisions = [d for d in decisions if d.get("origin", "").upper() == "CAPI"] results = [] for d in decisions: ip = d.get("value", "") ip_clean = ip.split("/")[0] geo = geolocate(ip_clean) if ip_clean else None results.append({ "ip": ip, "geo": geo, "reason": d.get("reason", ""), "scenario": d.get("scenario", d.get("reason", "")), "origin": d.get("origin", ""), "type": d.get("type", "ban"), "expires_in": d.get("duration", ""), "scope": d.get("scope", ""), }) return results except Exception as e: log.error(f"CrowdSec fetch error: {e}") return [] @app.get("/api/crowdsec/counts") async def get_crowdsec_counts(): """Return ban counts per origin without geo enrichment (fast).""" if not CROWDSEC_URL or not CROWDSEC_API_KEY: return {"local": 0, "capi": 0, "total": 0} try: async with httpx.AsyncClient(timeout=8.0) as http: r = await http.get( f"{CROWDSEC_URL}/v1/decisions", headers={"X-Api-Key": CROWDSEC_API_KEY}, params={"limit": 5000} ) if r.status_code == 204: return {"local": 0, "capi": 0, "total": 0} r.raise_for_status() decisions = r.json() or [] capi = sum(1 for d in decisions if d.get("origin","").upper() == "CAPI") local = sum(1 for d in decisions if d.get("origin","").lower() not in ("capi","lists")) return {"local": local, "capi": capi, "total": len(decisions)} except Exception as e: log.error(f"CrowdSec counts error: {e}") return {"local": 0, "capi": 0, "total": 0} @app.get("/api/rdns/{ip}") async def get_rdns(ip: str): """Reverse DNS lookup.""" import socket try: hostname = await asyncio.get_event_loop().run_in_executor( None, lambda: socket.gethostbyaddr(ip)[0] ) return {"hostname": hostname} except Exception: return {"hostname": None} # Server-side ASN cache to avoid hammering ip-api.com _asn_cache: dict[str, str] = {} @app.get("/api/asn/{ip}") async def get_asn(ip: str): """Proxy ASN/org lookup with caching and fallback.""" if ip in _asn_cache: return {"org": _asn_cache[ip]} try: async with httpx.AsyncClient(timeout=5.0) as http: # Primary: ip-api.com (45 req/min free) r = await http.get( f"http://ip-api.com/json/{ip}", params={"fields": "org,isp,status"} ) if r.status_code == 200: d = r.json() if d.get("status") == "success": org = d.get("org") or d.get("isp") or "Unknown" _asn_cache[ip] = org return {"org": org} # Fallback: ipinfo.io r2 = await http.get(f"https://ipinfo.io/{ip}/json") if r2.status_code == 200: d2 = r2.json() org = d2.get("org") or "Unknown" _asn_cache[ip] = org return {"org": org} except Exception as e: log.debug(f"ASN lookup error for {ip}: {e}") return {"org": "Unknown"} @app.websocket("/ws") async def websocket_endpoint(ws: WebSocket): await ws.accept() cid = id(ws) ws_clients[cid] = ws log.info(f"WS client connected. Total: {len(ws_clients)}") try: await ws.send_text(json.dumps({ "type": "init", "home": {"lat": HOME_LAT, "lon": HOME_LON, "name": HOME_NAME}, "geoip_loaded": geo_reader is not None, "devices": list(device_registry.values()), "recent_flows": recent_flows[-50:], })) while True: await ws.receive_text() except WebSocketDisconnect: log.info(f"WS client disconnected.") finally: ws_clients.pop(cid, None)