mirror of
https://github.com/Av0idD4rk/packmate.git
synced 2026-10-23 14:31:01 +00:00
239 lines
11 KiB
Python
239 lines
11 KiB
Python
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()
|