malmert
dfad09fa21
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 12s
394 lines
14 KiB
Python
394 lines
14 KiB
Python
"""Manage optional NFStream flow enrichment workers and packet lookups."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, Optional
|
|
|
|
import src.shared_objects as shared_objects
|
|
from src.config import settings
|
|
from src.utilities.flow_identity import flow_cache_key, flow_key_from_endpoints, flow_key_from_packet, packet_observed_at_ms
|
|
|
|
logger = logging.getLogger("nfstream_manager")
|
|
|
|
|
|
def _merge_metadata(base: Optional[Dict[str, Any]], extra: Dict[str, Any]) -> Dict[str, Any]:
|
|
merged = dict(base or {})
|
|
for key, value in extra.items():
|
|
if value not in (None, "", [], {}):
|
|
merged[key] = value
|
|
return merged
|
|
|
|
|
|
def _safe_text(value: Any) -> Optional[str]:
|
|
if value is None:
|
|
return None
|
|
text = str(value).strip()
|
|
return text if text else None
|
|
|
|
|
|
def _normalized_label(value: Any) -> Optional[str]:
|
|
text = _safe_text(value)
|
|
if text is None:
|
|
return None
|
|
if text.isdigit():
|
|
return None
|
|
if text.lower() in {"unknown", "other", "unclassified", "none", "null"}:
|
|
return None
|
|
return text
|
|
|
|
|
|
def _infer_is_encrypted(payload: Dict[str, Any]) -> Optional[bool]:
|
|
application_name = str(_normalized_label(payload.get("application_name")) or "").upper()
|
|
if any(token in application_name for token in ("TLS", "HTTPS", "QUIC", "SSL")):
|
|
return True
|
|
return None
|
|
|
|
|
|
def _build_enrichment(payload: Dict[str, Any], iface: str, flow_key: str) -> Dict[str, Any]:
|
|
app_protocol = _normalized_label(payload.get("application_name"))
|
|
if app_protocol is None and any(payload.get(key) for key in ("user_agent", "content_type")):
|
|
app_protocol = "HTTP"
|
|
|
|
metadata = {
|
|
"iface": iface,
|
|
"flow_key": flow_key,
|
|
"first_seen_ms": int(payload.get("first_seen_ms") or 0) or None,
|
|
"last_seen_ms": int(payload.get("last_seen_ms") or 0) or None,
|
|
"event": payload.get("event"),
|
|
"application_name": _safe_text(payload.get("application_name")),
|
|
"application_category_name": _safe_text(payload.get("application_category_name")),
|
|
"application_confidence": _safe_text(payload.get("application_confidence")),
|
|
"bidirectional_packets": payload.get("bidirectional_packets"),
|
|
"bidirectional_bytes": payload.get("bidirectional_bytes"),
|
|
}
|
|
metadata = _merge_metadata(
|
|
metadata,
|
|
{
|
|
"requested_server_name": payload.get("requested_server_name"),
|
|
"client_fingerprint": payload.get("client_fingerprint"),
|
|
"server_fingerprint": payload.get("server_fingerprint"),
|
|
"user_agent": payload.get("user_agent"),
|
|
"content_type": payload.get("content_type"),
|
|
},
|
|
)
|
|
|
|
return {
|
|
"app_protocol": app_protocol,
|
|
"app_master_protocol": app_protocol,
|
|
"app_category": _normalized_label(payload.get("application_category_name")),
|
|
"app_confidence": _safe_text(payload.get("application_confidence")),
|
|
"app_hostname": _safe_text(payload.get("requested_server_name")),
|
|
"app_is_encrypted": _infer_is_encrypted(payload),
|
|
"dpi_metadata": {"nfstream": metadata},
|
|
}
|
|
|
|
|
|
def _has_useful_metadata(payload: Dict[str, Any]) -> bool:
|
|
return any(
|
|
payload.get(key) not in (None, "", [], {})
|
|
for key in (
|
|
"application_name",
|
|
"application_category_name",
|
|
"application_confidence",
|
|
"requested_server_name",
|
|
"client_fingerprint",
|
|
"server_fingerprint",
|
|
"user_agent",
|
|
"content_type",
|
|
)
|
|
)
|
|
|
|
|
|
class NFStreamManager:
|
|
"""Own long-lived NFStream subprocesses and a short-lived flow metadata cache."""
|
|
|
|
def __init__(self) -> None:
|
|
self._enabled = settings.nfstream_enabled
|
|
self._workers: Dict[str, Dict[str, Any]] = {}
|
|
self._cache: Dict[tuple[str, str], Dict[str, Any]] = {}
|
|
self._last_error_by_iface: Dict[str, str] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return self._enabled
|
|
|
|
def update_interfaces(self, interfaces: Iterable[str]) -> None:
|
|
"""Align worker set with currently active interfaces."""
|
|
if not self._enabled:
|
|
return
|
|
|
|
targets = {iface.strip() for iface in interfaces if iface and iface.strip()}
|
|
with self._lock:
|
|
current = set(self._workers.keys())
|
|
for iface in sorted(current - targets):
|
|
self._stop_worker_locked(iface)
|
|
for iface in sorted(targets - current):
|
|
self._start_worker_locked(iface)
|
|
self._purge_cache_locked()
|
|
|
|
def stop(self) -> None:
|
|
"""Stop all workers and clear transient state."""
|
|
with self._lock:
|
|
for iface in list(self._workers.keys()):
|
|
self._stop_worker_locked(iface)
|
|
self._cache.clear()
|
|
|
|
def lookup_packet(self, pkt: Any, iface: Optional[str]) -> Dict[str, Any]:
|
|
"""Return NFStream-derived enrichment for the packet if a recent flow is known."""
|
|
if not self._enabled or not iface:
|
|
return {}
|
|
|
|
flow_key = flow_key_from_packet(pkt)
|
|
if flow_key is None:
|
|
return {}
|
|
|
|
now_ms = packet_observed_at_ms(pkt)
|
|
with self._lock:
|
|
self._purge_cache_locked(now_ms=now_ms)
|
|
payload = self._cache.get(flow_cache_key(iface, flow_key))
|
|
if payload is None:
|
|
return {}
|
|
|
|
first_seen_ms = int(payload.get("first_seen_ms") or 0)
|
|
last_seen_ms = int(payload.get("last_seen_ms") or 0)
|
|
window_ms = settings.nfstream_lookup_window_ms
|
|
if now_ms + window_ms < first_seen_ms or now_ms - window_ms > last_seen_ms:
|
|
return {}
|
|
return _build_enrichment(payload, iface, flow_key)
|
|
|
|
def get_debug_snapshot(self) -> Dict[str, Any]:
|
|
"""Expose current worker/cache state for API debugging."""
|
|
with self._lock:
|
|
self._purge_cache_locked()
|
|
return {
|
|
"enabled": self._enabled,
|
|
"workers": {
|
|
iface: {
|
|
"running": bool(worker.get("process") and worker["process"].poll() is None),
|
|
"thread_alive": bool(worker.get("reader") and worker["reader"].is_alive()),
|
|
"last_error": self._last_error_by_iface.get(iface),
|
|
}
|
|
for iface, worker in self._workers.items()
|
|
},
|
|
"cache_entries": len(self._cache),
|
|
"last_error_by_iface": dict(self._last_error_by_iface),
|
|
}
|
|
|
|
def _start_worker_locked(self, iface: str) -> None:
|
|
helper = Path(__file__).with_name("nfstream_flow_worker.py")
|
|
python_bin = sys.executable or "python3"
|
|
|
|
env = os.environ.copy()
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
backend_root = str(helper.parents[2])
|
|
existing_pythonpath = env.get("PYTHONPATH", "")
|
|
env["PYTHONPATH"] = backend_root if not existing_pythonpath else f"{backend_root}:{existing_pythonpath}"
|
|
|
|
cmd = [
|
|
python_bin,
|
|
str(helper),
|
|
"--iface",
|
|
iface,
|
|
"--idle-timeout",
|
|
str(settings.nfstream_idle_timeout_seconds),
|
|
"--active-timeout",
|
|
str(settings.nfstream_active_timeout_seconds),
|
|
"--snapshot-length",
|
|
str(settings.nfstream_snapshot_length),
|
|
"--n-dissections",
|
|
str(settings.nfstream_n_dissections),
|
|
"--n-meters",
|
|
str(settings.nfstream_n_meters),
|
|
]
|
|
if settings.nfstream_promiscuous_mode:
|
|
cmd.append("--promiscuous-mode")
|
|
|
|
logger.info("Starting NFStream worker for %s", iface)
|
|
try:
|
|
process = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
env=env,
|
|
start_new_session=True,
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to start NFStream worker for %s", iface)
|
|
self._last_error_by_iface[iface] = "failed_to_start_process"
|
|
return
|
|
|
|
reader = threading.Thread(
|
|
target=self._read_loop,
|
|
args=(iface, process),
|
|
daemon=True,
|
|
name=f"nfstream-reader-{iface}",
|
|
)
|
|
self._workers[iface] = {
|
|
"process": process,
|
|
"reader": reader,
|
|
}
|
|
reader.start()
|
|
|
|
def _stop_worker_locked(self, iface: str) -> None:
|
|
worker = self._workers.pop(iface, None)
|
|
if worker is None:
|
|
return
|
|
|
|
process = worker.get("process")
|
|
reader = worker.get("reader")
|
|
|
|
if process is not None and process.poll() is None:
|
|
try:
|
|
os.killpg(os.getpgid(process.pid), signal.SIGTERM)
|
|
process.wait(timeout=settings.nfstream_process_stop_timeout_seconds)
|
|
except subprocess.TimeoutExpired:
|
|
try:
|
|
os.killpg(os.getpgid(process.pid), signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
except Exception:
|
|
logger.exception("Failed to stop NFStream worker for %s cleanly", iface)
|
|
|
|
if reader is not None and reader.is_alive():
|
|
reader.join(timeout=settings.nfstream_reader_join_timeout_seconds)
|
|
|
|
for key in [cache_key for cache_key in self._cache.keys() if cache_key[0] == iface]:
|
|
self._cache.pop(key, None)
|
|
|
|
def _read_loop(self, iface: str, process: subprocess.Popen[str]) -> None:
|
|
stdout = process.stdout
|
|
if stdout is None:
|
|
return
|
|
|
|
for line in stdout:
|
|
text = line.strip()
|
|
if not text:
|
|
continue
|
|
try:
|
|
payload = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
logger.info("nfstream[%s]: %s", iface, text)
|
|
continue
|
|
|
|
payload_type = payload.get("type")
|
|
if payload_type == "worker_ready":
|
|
logger.info("NFStream worker ready on %s", iface)
|
|
continue
|
|
if payload_type == "worker_error":
|
|
message = str(payload.get("message") or "unknown error")
|
|
logger.warning("NFStream worker error on %s: %s", iface, message)
|
|
with self._lock:
|
|
self._last_error_by_iface[iface] = message
|
|
continue
|
|
if payload_type != "flow_update":
|
|
logger.debug("Ignoring NFStream payload on %s: %s", iface, payload)
|
|
continue
|
|
|
|
flow_key = flow_key_from_endpoints(
|
|
payload.get("ip_version"),
|
|
payload.get("protocol"),
|
|
payload.get("src_ip"),
|
|
payload.get("src_port"),
|
|
payload.get("dst_ip"),
|
|
payload.get("dst_port"),
|
|
)
|
|
if flow_key is None:
|
|
continue
|
|
|
|
with self._lock:
|
|
previous = self._cache.get(flow_cache_key(iface, flow_key))
|
|
self._cache[flow_cache_key(iface, flow_key)] = dict(payload)
|
|
self._purge_cache_locked()
|
|
|
|
if self._should_backfill(previous, payload):
|
|
self._schedule_backfill(iface, flow_key, payload)
|
|
|
|
rc = process.poll()
|
|
if rc not in (0, None):
|
|
logger.warning("NFStream worker for %s exited with code %s", iface, rc)
|
|
|
|
def _purge_cache_locked(self, now_ms: Optional[int] = None) -> None:
|
|
if now_ms is None:
|
|
now_ms = int(time.time() * 1000)
|
|
expiry_ms = int(settings.nfstream_cache_ttl_seconds * 1000)
|
|
|
|
stale_keys = []
|
|
for key, payload in self._cache.items():
|
|
last_seen_ms = int(payload.get("last_seen_ms") or 0)
|
|
if last_seen_ms and now_ms - last_seen_ms > expiry_ms:
|
|
stale_keys.append(key)
|
|
|
|
for key in stale_keys:
|
|
self._cache.pop(key, None)
|
|
|
|
def _should_backfill(self, previous: Optional[Dict[str, Any]], current: Dict[str, Any]) -> bool:
|
|
if not _has_useful_metadata(current):
|
|
return False
|
|
if previous is None:
|
|
return True
|
|
|
|
interesting_keys = (
|
|
"application_name",
|
|
"application_category_name",
|
|
"application_confidence",
|
|
"requested_server_name",
|
|
"client_fingerprint",
|
|
"server_fingerprint",
|
|
"user_agent",
|
|
"content_type",
|
|
)
|
|
return any(previous.get(key) != current.get(key) for key in interesting_keys)
|
|
|
|
def _schedule_backfill(self, iface: str, flow_key: str, payload: Dict[str, Any]) -> None:
|
|
web_db = getattr(shared_objects, "db", None)
|
|
web_loop = getattr(shared_objects, "web_loop", None)
|
|
if web_db is None or web_loop is None:
|
|
return
|
|
|
|
enrichment = _build_enrichment(payload, iface, flow_key)
|
|
|
|
async def _backfill() -> None:
|
|
updated = await web_db.backfill_flow_metadata(
|
|
iface=iface,
|
|
src_ip=str(payload.get("src_ip")),
|
|
dst_ip=str(payload.get("dst_ip")),
|
|
src_port=int(payload.get("src_port") or 0),
|
|
dst_port=int(payload.get("dst_port") or 0),
|
|
protocol=int(payload.get("protocol") or 0),
|
|
first_seen_ms=int(payload.get("first_seen_ms") or 0),
|
|
last_seen_ms=int(payload.get("last_seen_ms") or 0),
|
|
enrichment=enrichment,
|
|
window_ms=settings.nfstream_lookup_window_ms,
|
|
)
|
|
if updated:
|
|
logger.debug("Backfilled NFStream metadata for %s packets on %s flow=%s", updated, iface, flow_key)
|
|
|
|
try:
|
|
future = asyncio.run_coroutine_threadsafe(_backfill(), web_loop)
|
|
future.add_done_callback(self._log_backfill_result)
|
|
except Exception:
|
|
logger.exception("Failed to schedule NFStream metadata backfill for %s", flow_key)
|
|
|
|
@staticmethod
|
|
def _log_backfill_result(future: Any) -> None:
|
|
try:
|
|
future.result()
|
|
except Exception:
|
|
logger.exception("NFStream metadata backfill task failed")
|
|
|
|
|
|
nfstream_manager = NFStreamManager()
|