import asyncio import base64 import unittest from unittest.mock import patch import httpx from pydantic import ValidationError from starlette.applications import Starlette from starlette.responses import JSONResponse from starlette.routing import Route from starlette.testclient import TestClient from packmate_mcp.jobs import Jobs, inspect_packet from packmate_mcp.models import AddressFilter, Match, TrafficQuery, PatternSettings from packmate_mcp.server import TokenGate, make_server, application def packet(id=1, data=b"prefix FLAG{abc} suffix", size=None): return {"packet_id": id, "stream_id": 10, "stream_timestamp": 1, "timestamp": 2, "service_port": 8080, "source_ip": "10.60.1.1", "destination_ip": "10.50.1.1", "incoming": True, "content": base64.b64encode(data).decode(), "size_bytes": len(data) if size is None else size} class FakeBackend: def __init__(self, rows=None, block=False): self.rows = rows or [] self.block = block self.calls = [] async def request(self, method, path, **kwargs): self.calls.append((method, path, kwargs)) if self.block: await asyncio.sleep(100) if path == "/api/analysis/scan": return {"items": self.rows, "exhausted": True} if "/lookback/" in path: return {"processed": kwargs["json"]["packet_ids"], "unavailable": []} return [] async def close(self): pass class MatchingTests(unittest.TestCase): def test_stateless_requests_share_application_lifetime(self): backend = FakeBackend() backend.closed = False async def close(): backend.closed = True backend.close = close headers = {"Authorization": "Bearer " + "a" * 32, "Accept": "application/json, text/event-stream"} with patch("packmate_mcp.server.Backend", return_value=backend), patch.dict("os.environ", { "PACKMATE_MCP_ENABLED": "true", "PACKMATE_MCP_TOKEN": "a" * 32, "PACKMATE_MCP_ALLOWED_HOSTS": "localhost:*" }): with TestClient(application(), base_url="http://localhost:65002", headers=headers) as client: for i in range(2): response = client.post("/mcp", json={"jsonrpc": "2.0", "id": i + 1, "method": "tools/call", "params": {"name": "list_services", "arguments": {}}}) self.assertEqual(response.status_code, 200) self.assertEqual(response.json()["result"]["structuredContent"]["items"], []) self.assertFalse(backend.closed) self.assertTrue(backend.closed) def test_regex_ips_cidr_and_binary_offsets(self): query = TrafficQuery(client_ip=AddressFilter(type="cidr", value="10.60.0.0/16"), match=Match(value=r"FLAG\{[^}]+\}")) result, skipped = inspect_packet(packet(), query, 20) self.assertFalse(skipped) self.assertEqual(result["match"]["start"], 7) self.assertNotIn("content", result) query.match = Match(target="source_ip", value=r"^10\.60\.") self.assertEqual(inspect_packet(packet(), query, 20)[0]["match"]["target"], "source_ip") query.match = Match(type="hex", value="ff0041") result, _ = inspect_packet(packet(data=b"x\xff\x00Az"), query, 4) self.assertEqual(result["match"]["start"], 1) self.assertEqual(result["match"]["offset_units"], "bytes") def test_truncation_does_not_invent_end_anchor_match(self): result, skipped = inspect_packet(packet(data=b"FLAG", size=100000), TrafficQuery(match=Match(value="FLAG$")), 10) self.assertIsNone(result) self.assertTrue(skipped) def test_zero_length_does_not_hide_later_nonempty_match(self): result, _ = inspect_packet(packet(data=b" FLAG"), TrafficQuery(match=Match(value="^|FLAG")), 10) self.assertEqual(result["match"]["end_exclusive"], 5) def test_invalid_regex_and_unsafe_resume_rejected(self): with self.assertRaises(ValidationError): Match(value=r"(x)\1") with self.assertRaises(ValidationError): TrafficQuery(cursor={"stream_timestamp": 1, "stream_id": 1, "packet_id": 1}) with self.assertRaises(ValidationError): PatternSettings(name="too long", match=Match(value="a" * 256)) class JobTests(unittest.IsolatedAsyncioTestCase): async def test_limited_results_and_resume_cursor(self): jobs = Jobs(FakeBackend([packet(1), packet(2)])) started = jobs.start(TrafficQuery(), limit=1) job = jobs.get(started["job_id"]) await job.task self.assertEqual(job.reason, "result_limit") self.assertEqual(job.cursor["packet_id"], 1) self.assertFalse(jobs.describe(job)["coverage_complete"]) async def test_busy_cancel_and_terminal_ttl(self): jobs = Jobs(FakeBackend(block=True), max_active=1, ttl=-1) first = jobs.start(TrafficQuery()) with self.assertRaisesRegex(ValueError, "BUSY"): jobs.start(TrafficQuery()) self.assertEqual((await jobs.cancel(first["job_id"]))["status"], "cancelled") with self.assertRaisesRegex(ValueError, "NOT_FOUND"): jobs.get(first["job_id"]) async def test_deadline_and_byte_budget(self): jobs = Jobs(FakeBackend(block=True), max_seconds=0.01) first = jobs.start(TrafficQuery()) job = jobs.get(first["job_id"]) await job.task self.assertEqual(job.reason, "time_budget") jobs = Jobs(FakeBackend([packet()]), max_bytes=1) job = jobs.get(jobs.start(TrafficQuery())["job_id"]) await job.task self.assertEqual(job.reason, "byte_budget") self.assertEqual(job.scanned_packets, 0) async def test_lookback_persists_through_backend(self): backend = FakeBackend([packet()]) jobs = Jobs(backend) job = jobs.get(jobs.start(TrafficQuery(match=Match(type="substring", value="FLAG")), kind="lookback", pattern_id=12)["job_id"]) await job.task self.assertEqual(job.applied_packets, 1) self.assertTrue(jobs.describe(job)["coverage_complete"]) self.assertTrue(any(path == "/api/analysis/lookback/12" for _, path, _ in backend.calls)) async def test_interrupted_lookback_resumes_before_unconfirmed_write(self): entered_write = asyncio.Event() backend = FakeBackend([packet()]) original_request = backend.request async def interrupted(method, path, **kwargs): if "/lookback/" in path: entered_write.set() await asyncio.sleep(100) return await original_request(method, path, **kwargs) backend.request = interrupted jobs = Jobs(backend) started = jobs.start(TrafficQuery(match=Match(type="substring", value="FLAG")), kind="lookback", pattern_id=12) await asyncio.wait_for(entered_write.wait(), 1) cancelled = await jobs.cancel(started["job_id"]) self.assertEqual(cancelled["matched_packets"], 1) self.assertEqual(cancelled["applied_packets"], 0) self.assertEqual(cancelled["unavailable_packets"], 1) self.assertIsNone(cancelled["next_scan_cursor"]) async def test_all_tools_use_one_catalog_and_strict_schemas(self): backend = FakeBackend() server = make_server(backend, Jobs(backend), "http://localhost:65000", ["localhost:*"]) tools = {tool.name: tool for tool in await server.list_tools()} self.assertIn("create_service", tools) self.assertIn("create_pattern", tools) self.assertIn("delete_service", tools) self.assertIn("start_lookback", tools) self.assertTrue(all(t.outputSchema is not None for t in tools.values())) self.assertEqual(tools["create_service"].inputSchema["properties"]["port"]["maximum"], 65535) async def test_authentication_and_disabled_mode(self): async def ok(request): return JSONResponse({"ok": True}) app = Starlette(routes=[Route("/mcp", ok)]) for enabled in (False, True): gate = TokenGate(app, "a" * 32, enabled) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=gate), base_url="http://localhost") as client: self.assertEqual((await client.get("/healthz")).status_code, 200) self.assertEqual((await client.get("/mcp")).status_code, 401 if enabled else 503) response = await client.get("/mcp", headers={"Authorization": "Bearer " + "a" * 32}) self.assertEqual(response.status_code, 200 if enabled else 503) async def test_validation_checks_capture_engine_as_well_as_python_re2(self): backend = FakeBackend() async def reject(method, path, **kwargs): self.assertEqual(path, "/api/analysis/regex/validate") self.assertEqual(kwargs["json"], {"value": r"\C"}) return {"valid": False, "error": "unsupported escape"} backend.request = reject server = make_server(backend, Jobs(backend), "http://localhost", ["localhost:*"]) with self.assertRaisesRegex(Exception, "RE2J capture engine"): await server.call_tool("validate_pattern", {"settings": {"name": "byte escape", "match": {"value": r"\C"}}})