import asyncio import base64 import ipaddress import time import uuid from collections import Counter, OrderedDict from dataclasses import dataclass, field from .models import TrafficQuery, compile_regex def address_matches(value, rule): if rule is None: return True if not value: return False if rule.type == "regex": return bool(compile_regex(rule.value).search(value)) try: address = ipaddress.ip_address(value) return address in ipaddress.ip_network(rule.value) if rule.type == "cidr" else address == ipaddress.ip_address(rule.value) except ValueError: return False def inspect_packet(row, query, context_bytes): client = row.get("source_ip") if row["incoming"] else row.get("destination_ip") if not all((address_matches(row.get("source_ip"), query.source_ip), address_matches(row.get("destination_ip"), query.destination_ip), address_matches(client, query.client_ip))): return None, False data = base64.b64decode(row.get("content") or "") match = query.match start = end = 0 target = None units = "bytes" if match: if match.target == "content": if len(data) < row["size_bytes"]: # Do not match a truncated prefix: $ and similar expressions could give false positives. return None, True if match.type == "hex": needle = bytes.fromhex(match.value) start = data.find(needle) end = start + len(needle) else: text = data.decode("utf-8", errors="replace") if match.type == "regex": found = next((m for m in compile_regex(match.value).finditer(text) if m.start() != m.end()), None) if not found: return None, False a, b = found.span() else: a = text.find(match.value) if a < 0: return None, False b = a + len(match.value) # Character offsets are explicit; they are not offsets into the original binary bytes. units = "unicode_codepoints" start, end = a, b snippet = text[max(0, a - context_bytes):b + context_bytes][:2 * context_bytes + 256] if start < 0: return None, False target = "content" else: fields = ["source_ip", "destination_ip"] if match.target == "either_ip" else [match.target] for name in fields: found = compile_regex(match.value).search(row.get(name) or "") if found: start, end = found.span() target = name break if target is None: return None, False units = "unicode_codepoints" result = {k: v for k, v in row.items() if k != "content"} result["client_ip"] = client if target: result["match"] = {"target": target, "start": start, "end_exclusive": end, "offset_units": units} if match and match.target == "content" and match.type != "hex": result["snippet"] = snippet elif data: result["snippet_hex"] = data[max(0, start - context_bytes):start + context_bytes].hex() return result, False @dataclass class Job: id: str kind: str query: TrafficQuery status: str = "running" reason: str | None = None created: float = field(default_factory=time.monotonic) finished: float | None = None task: asyncio.Task | None = None results: list = field(default_factory=list) cursor: dict | None = None scanned_packets: int = 0 scanned_bytes: int = 0 matched_packets: int = 0 skipped_oversize: int = 0 unavailable: int = 0 applied_packets: int = 0 exhausted: bool = False groups: Counter = field(default_factory=Counter) groups_truncated: bool = False class Jobs: def __init__(self, backend, *, max_active=2, max_jobs=32, ttl=900, max_seconds=30, max_bytes=32 * 1024 * 1024): self.backend = backend self.max_active, self.max_jobs, self.ttl = max_active, max_jobs, ttl self.max_seconds, self.max_bytes = max_seconds, max_bytes self.jobs = OrderedDict() def reap(self): for key, job in list(self.jobs.items()): if job.finished is not None and time.monotonic() - job.finished > self.ttl: del self.jobs[key] def get(self, key): self.reap() if key not in self.jobs: raise ValueError("NOT_FOUND: job expired, was lost on MCP restart, or does not exist") return self.jobs[key] def describe(self, job): return {"job_id": job.id, "kind": job.kind, "status": job.status, "reason": job.reason, "query": job.query.model_dump(), "scanned_packets": job.scanned_packets, "scanned_bytes": job.scanned_bytes, "matched_packets": job.matched_packets, "skipped_oversize_packets": job.skipped_oversize, "unavailable_packets": job.unavailable, "applied_packets": job.applied_packets, "scan_exhausted": job.exhausted, "coverage_complete": job.exhausted and not job.skipped_oversize and not job.unavailable, "next_scan_cursor": job.cursor, "result_count": len(job.results), "groups_truncated": job.groups_truncated, "consistency": "live_read_committed; retention and concurrent commits may change visibility", "result_ttl_seconds": self.ttl} def start(self, query, *, kind="search", limit=100, context_bytes=160, group_by="service_port", pattern_id=None): self.reap() if sum(j.status == "running" for j in self.jobs.values()) >= self.max_active: raise ValueError("BUSY: analysis job slots are full; wait or cancel a job") if len(self.jobs) >= self.max_jobs: raise ValueError("BUSY: retained job limit reached; delete completed jobs or wait for TTL") now = int(time.time() * 1000) query = query.model_copy(deep=True) if query.from_ms is None: query.to_ms = now query.from_ms = max(0, now - query.last_minutes * 60000) query.snapshot_ms = query.snapshot_ms or now job = Job(uuid.uuid4().hex, kind, query) job.cursor = query.cursor.model_dump() if query.cursor else None self.jobs[job.id] = job job.task = asyncio.create_task(self._run(job, limit, context_bytes, group_by, pattern_id)) return self.describe(job) async def _run(self, job, limit, context_bytes, group_by, pattern_id): try: async with asyncio.timeout(self.max_seconds): while True: q = job.query confirmed_cursor = job.cursor needs_content = q.match is not None and q.match.target == "content" body = {key: getattr(q, key) for key in ("from_ms", "to_ms", "snapshot_ms", "service_ports", "protocol", "direction", "favorites_only", "pattern_id")} body.update(cursor=job.cursor, limit=16, content_bytes=65536 if needs_content else 0) page = await self.backend.request("POST", "/api/analysis/scan", json=body) apply_ids = [] for row in page["items"]: charged = min(row["size_bytes"], body["content_bytes"]) + 256 if job.scanned_bytes + charged > self.max_bytes: job.reason = "byte_budget" break result, skipped = inspect_packet(row, q, context_bytes) job.scanned_packets += 1 job.scanned_bytes += charged job.skipped_oversize += int(skipped) job.cursor = {"stream_timestamp": row["stream_timestamp"], "stream_id": row["stream_id"], "packet_id": row["packet_id"]} if result is not None: job.matched_packets += 1 if job.kind == "aggregate": key = str(result.get(group_by) if group_by != "minute" else row["timestamp"] // 60000 * 60000) if key in job.groups or len(job.groups) < 128: job.groups[key] += 1 else: job.groups_truncated = True else: if len(job.results) < limit: job.results.append(result) if pattern_id is not None: apply_ids.append(row["packet_id"]) if job.matched_packets >= limit: job.reason = "result_limit" break if apply_ids: try: applied = await self.backend.request("POST", f"/api/analysis/lookback/{pattern_id}", json={"packet_ids": apply_ids}) except (Exception, asyncio.CancelledError): # A cancelled/failed HTTP write may or may not have committed. # Resume before this page; idempotent lookback makes replay safe. job.cursor = confirmed_cursor job.unavailable += len(apply_ids) raise job.applied_packets += len(applied["processed"]) job.unavailable += len(applied["unavailable"]) if job.reason: break if page["exhausted"]: job.exhausted = True break await asyncio.sleep(0.025) job.status = "completed" except TimeoutError: job.status, job.reason = "partial", "time_budget" except asyncio.CancelledError: job.status, job.reason = "cancelled", "cancelled_by_client; an in-flight write may already have committed" except Exception as error: job.status, job.reason = "failed", str(error)[:256] finally: if job.kind == "aggregate": job.results = [{"group": key, "packets": count} for key, count in job.groups.most_common()] job.finished = time.monotonic() async def cancel(self, key): job = self.get(key) if job.status == "running": job.task.cancel() await asyncio.gather(job.task, return_exceptions=True) if job.status == "running": # cancelled before the coroutine entered its try block job.status, job.reason, job.finished = "cancelled", "cancelled_by_client", time.monotonic() return self.describe(job) async def close(self): for job in list(self.jobs.values()): if job.status == "running": await self.cancel(job.id) await self.backend.close()