packmate/mcp/tests/integration.py

"""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())