mirror of
https://github.com/Av0idD4rk/packmate.git
synced 2026-10-23 15:11:01 +00:00
189 lines
9.2 KiB
Python
189 lines
9.2 KiB
Python
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"}}})
|