dmo/MorkNetVisualizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
5f8a43e6529b5a0c38d2827bf95aff56c9094210
MorkNetVisualizer/backend/main.py
T

599 lines
23 KiB
Python
Raw Blame History

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)
Reference in New Issue View Git Blame Copy Permalink