mirror of
https://github.com/Av0idD4rk/packmate.git
synced 2026-10-23 16:31:01 +00:00
195 lines
12 KiB
Python
195 lines
12 KiB
Python
"""Run only in tests/compose.yml's isolated network and disposable PostgreSQL database."""
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import math
|
|
import struct
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import psycopg
|
|
from mcp import ClientSession
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
from websockets.asyncio.client import connect
|
|
|
|
API = "http://172.31.250.10:65000"
|
|
TOKEN = "mcp-integration-test-token-000000000000"
|
|
|
|
|
|
async def main():
|
|
report = {"baseline_streams": 200, "live_streams": 3000, "checks": []}
|
|
db = await psycopg.AsyncConnection.connect("host=db dbname=packmate user=packmate password=network-stress", autocommit=True)
|
|
|
|
async def scalar(sql, params=()):
|
|
cursor = await db.execute(sql, params)
|
|
return (await cursor.fetchone())[0]
|
|
|
|
async with httpx.AsyncClient(auth=("stress", "stress"), timeout=10) as rest:
|
|
for _ in range(60):
|
|
try:
|
|
if (await rest.get(API + "/api/analysis/overview")).status_code == 200:
|
|
break
|
|
except httpx.HTTPError:
|
|
pass
|
|
await asyncio.sleep(1)
|
|
else:
|
|
raise AssertionError("backend failed to become ready")
|
|
|
|
headers = {"Authorization": "Bearer " + TOKEN}
|
|
async with httpx.AsyncClient(headers=headers, timeout=45) as transport:
|
|
async with streamable_http_client("http://mcp:65002/mcp", http_client=transport) as (read, write, _):
|
|
async with ClientSession(read, write) as client:
|
|
await client.initialize()
|
|
tools = {tool.name for tool in (await client.list_tools()).tools}
|
|
assert {"create_service", "create_pattern", "delete_service", "search_traffic", "start_lookback"} <= tools
|
|
report["tool_count"] = len(tools)
|
|
|
|
async def call(name, arguments=None):
|
|
response = await client.call_tool(name, arguments or {})
|
|
assert not response.isError, (name, response.content)
|
|
return response.structuredContent
|
|
|
|
async def finish(job):
|
|
for _ in range(150):
|
|
status = await call("get_job", {"job_id": job["job_id"]})
|
|
if status["status"] != "running":
|
|
assert status["status"] == "completed", status
|
|
return status
|
|
await asyncio.sleep(0.2)
|
|
raise AssertionError("job never finished")
|
|
|
|
async def send(i):
|
|
reader, writer = await asyncio.open_connection("172.31.250.10", 18080)
|
|
payload = f"PMMCP-{i:08d} FLAG{{mcp_test}}".encode()
|
|
writer.write(struct.pack("!I", len(payload)) + payload)
|
|
await writer.drain()
|
|
length = struct.unpack("!I", await reader.readexactly(4))[0]
|
|
assert await reader.readexactly(length) == payload
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
authorization = "Basic " + base64.b64encode(b"stress:stress").decode()
|
|
async with connect("ws://172.31.250.10:65000/api/ws/websocket", additional_headers={"Authorization": authorization}) as ws:
|
|
service = await call("create_service", {"port": 18080, "settings": {"name": "MCP echo", "http": False,
|
|
"urldecode_http_requests": False, "merge_adjacent_packets": False}})
|
|
assert service["service"]["port"] == 18080
|
|
pattern = await call("create_pattern", {"settings": {"name": "MCP flag", "match": {"type": "regex", "value": r"FLAG\{[^}]+\}"}, "service_port": 18080}})
|
|
pid = pattern["pattern"]["id"]
|
|
assert pattern["pattern"]["regexEngine"] == "RE2J"
|
|
events = []
|
|
async with asyncio.timeout(10):
|
|
while not {"SAVE_SERVICE", "SAVE_PATTERN"} <= set(events):
|
|
event = json.loads(await ws.recv())
|
|
events.append(event["type"])
|
|
report["checks"].append("MCP-created service and pattern broadcast to browser websocket")
|
|
|
|
for first in range(0, 200, 20):
|
|
await asyncio.gather(*(send(i) for i in range(first, first + 20)))
|
|
for _ in range(30):
|
|
if await scalar("SELECT count(*) FROM stream") == 200:
|
|
break
|
|
await asyncio.sleep(0.5)
|
|
assert await scalar("SELECT count(*) FROM stream") == 200
|
|
baseline_ids = [row[0] for row in await (await db.execute("SELECT id FROM stream ORDER BY id")).fetchall()]
|
|
assert await scalar("SELECT count(*) FROM stream_found_patterns WHERE pattern_id = %s", (pid,)) == 200
|
|
report["checks"].append("all baseline TCP streams have live pattern highlights")
|
|
|
|
search = await finish(await call("search_traffic", {"query": {"service_ports": [18080], "direction": "incoming",
|
|
"match": {"type": "regex", "value": "PMMCP-"}}, "limit": 500}))
|
|
assert search["coverage_complete"] and search["matched_packets"] == 200, search
|
|
rows = (await call("get_job_results", {"job_id": search["job_id"], "limit": 5}))["items"]
|
|
selected = rows[0]
|
|
packet = await call("read_packet", {"packet_id": selected["packet_id"], "format": "hex"})
|
|
assert b"PMMCP-" in bytes.fromhex(packet["content"])
|
|
assert (await rest.post(API + f"/api/packet/{selected['stream_id']}", json={"pageSize": 20})).status_code == 200
|
|
|
|
ip = await finish(await call("search_traffic", {"query": {"direction": "incoming", "match": {
|
|
"target": "source_ip", "type": "regex", "value": r"^172\.31\.250\."}}, "limit": 500}))
|
|
assert ip["matched_packets"] == 200, ip
|
|
preserved = [row["stream_id"] for row in rows]
|
|
favored = await call("set_favorites", {"stream_ids": preserved})
|
|
assert len(favored["updated"]) == 5
|
|
report["checks"].append("content and IP regex, exact payload reads, favorites confirmed")
|
|
|
|
historical = await call("create_pattern", {"settings": {"name": "Historical MCP", "match": {"type": "substring", "value": "PMMCP-"}, "direction": "INPUT"}})
|
|
hpid = historical["pattern"]["id"]
|
|
history = await finish(await call("start_lookback", {"pattern_id": hpid, "query": {}, "limit": 500}))
|
|
assert history["applied_packets"] == 200, history
|
|
count = await scalar("SELECT count(*) FROM found_pattern WHERE pattern_id = %s", (hpid,))
|
|
assert count == 200, count
|
|
await finish(await call("start_lookback", {"pattern_id": hpid, "query": {}, "limit": 500}))
|
|
assert await scalar("SELECT count(*) FROM found_pattern WHERE pattern_id = %s", (hpid,)) == count
|
|
visible = (await rest.post(API + f"/api/packet/{selected['stream_id']}", json={"pageSize": 20})).json()
|
|
assert any(any(m["patternId"] == hpid for m in packet["matches"]) for packet in visible), visible
|
|
report["checks"].append("lookback stores UI-visible highlights and repeat creates no duplicates")
|
|
|
|
# Reject requests before work, without trusting an agent to observe the limit.
|
|
bad = await client.call_tool("search_traffic", {"query": {"match": {"value": r"(x)\1"}}})
|
|
assert bad.isError
|
|
forbidden = await transport.post("http://mcp:65002/mcp", headers={"Authorization": "Bearer wrong"}, json={})
|
|
assert forbidden.status_code == 401
|
|
bound = await rest.get(API + f"/api/analysis/packets/{selected['packet_id']}", params={"length": 9999999})
|
|
assert bound.status_code == 400
|
|
report["checks"].append("invalid regex, missing token and oversized reads rejected")
|
|
|
|
# Keep later searches alive while traffic is written and the old baseline is cleaned.
|
|
await db.execute("UPDATE stream SET start_timestamp = start_timestamp - 120000, end_timestamp = end_timestamp - 120000, persisted_at = persisted_at - 120000 WHERE favorite = false")
|
|
load_start = time.monotonic()
|
|
latencies = []
|
|
sent = 0
|
|
overlap = False
|
|
|
|
async def load():
|
|
nonlocal sent
|
|
for start in range(0, 3000, 20):
|
|
await asyncio.gather(*(send(10000 + i) for i in range(start, start + 20)))
|
|
sent += 20
|
|
await asyncio.sleep(max(0, load_start + (start + 20) * 15 / 3000 - time.monotonic()))
|
|
|
|
async def monitor():
|
|
nonlocal overlap
|
|
while sent < 3000:
|
|
before = time.monotonic()
|
|
response = await rest.get(API + "/api/pcap/metrics")
|
|
response.raise_for_status()
|
|
latencies.append((time.monotonic() - before) * 1000)
|
|
remaining = await scalar("SELECT count(*) FROM stream WHERE id = ANY(%s) AND favorite = false", (baseline_ids,))
|
|
if 0 < remaining < 195 and 0 < sent < 3000:
|
|
overlap = True
|
|
await asyncio.sleep(0.1)
|
|
|
|
async def investigations():
|
|
while sent < 3000:
|
|
job = await call("search_traffic", {"query": {"match": {"value": "(a+)+$"}}, "limit": 1000})
|
|
await finish(job)
|
|
await call("delete_job", {"job_id": job["job_id"]})
|
|
|
|
await asyncio.gather(load(), monitor(), investigations())
|
|
assert overlap, "cleanup did not overlap live capture and MCP"
|
|
for _ in range(60):
|
|
if await scalar("SELECT count(*) FROM stream WHERE id = ANY(%s) AND favorite = false", (baseline_ids,)) == 0:
|
|
break
|
|
await asyncio.sleep(0.5)
|
|
assert await scalar("SELECT count(*) FROM stream WHERE id = ANY(%s)", (baseline_ids,)) == 5
|
|
for _ in range(30):
|
|
if await scalar("SELECT count(*) FROM stream") == 3005:
|
|
break
|
|
await asyncio.sleep(0.5)
|
|
assert await scalar("SELECT count(*) FROM stream") == 3005
|
|
metrics = (await rest.get(API + "/api/pcap/metrics")).json()
|
|
assert all(metrics[key] == 0 for key in ("droppedPackets", "interfaceDroppedPackets", "applicationDroppedPackets")), metrics
|
|
p99 = sorted(latencies)[math.ceil(len(latencies) * 0.99) - 1]
|
|
assert p99 < 500, p99
|
|
assert await scalar("SELECT count(*) FROM packet p LEFT JOIN stream s ON s.id = p.stream_id WHERE s.id IS NULL") == 0
|
|
report.update(api_p99_ms=round(p99, 3), api_samples=len(latencies), metrics=metrics, cleanup_overlap=overlap)
|
|
report["checks"].append("3000 live TCP streams + concurrent regex jobs + cleanup: no stream loss, no drops, favorites preserved")
|
|
await db.close()
|
|
Path("/reports").mkdir(exist_ok=True)
|
|
Path("/reports/integration.json").write_text(json.dumps(report, indent=2))
|
|
print(json.dumps(report, indent=2), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|