mirror of
https://github.com/Kowntaewook/MCPTrace.git
synced 2026-11-02 04:45:09 +01:00
kwonta3wook
fbcb4a92ac
Carve MCP/JSON-RPC candidates with original bytes and provenance, conservative ID correlation, opt-in argument policies, JSONL export, and scan coverage. Add synthetic Volatility integration tests, CI, and Korean documentation.
217 lines
8.0 KiB
Python
217 lines
8.0 KiB
Python
import base64
|
|
import hashlib
|
|
import importlib
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import sysconfig
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from test_core import encode, request
|
|
from volatility3.cli import CommandLine
|
|
from volatility3.framework import contexts, exceptions
|
|
from volatility3.framework.layers.physical import BufferDataLayer
|
|
|
|
from mcptrace.core import ScanStats, scan_chunks
|
|
from mcptrace.memory import readable_chunks
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def plugin_module():
|
|
import volatility3.plugins.windows
|
|
|
|
plugin_path = str(ROOT / "plugins" / "windows")
|
|
if plugin_path not in volatility3.plugins.windows.__path__:
|
|
volatility3.plugins.windows.__path__.insert(0, plugin_path)
|
|
return importlib.import_module("volatility3.plugins.windows.mcptrace")
|
|
|
|
|
|
class GappedLayer(BufferDataLayer):
|
|
def read(self, offset, length, pad=False):
|
|
if offset < 8192 and offset + length > 4096:
|
|
raise exceptions.InvalidAddressException(self.name, max(offset, 4096))
|
|
return super().read(offset, length, pad=pad)
|
|
|
|
|
|
def test_real_volatility_layer_does_not_bridge_missing_page():
|
|
context = contexts.Context()
|
|
intact = encode(request(2))
|
|
split = encode(request(1))
|
|
data = bytearray(b"\0" * 12288)
|
|
data[4000:4096] = split[:96]
|
|
data[8192 : 8192 + len(split[96:])] = split[96:]
|
|
data[9000 : 9000 + len(intact)] = intact
|
|
layer = GappedLayer(context, "memory", "memory", bytes(data))
|
|
stats = ScanStats()
|
|
records = list(
|
|
scan_chunks(
|
|
readable_chunks(layer, 0, len(data), stats, exceptions.InvalidAddressException),
|
|
include_partial=True,
|
|
)
|
|
)
|
|
assert [r.request_id for r in records if r.state == "complete"] == [2]
|
|
assert stats.bytes_read == 8192 and stats.unreadable_bytes == 4096
|
|
|
|
|
|
@pytest.fixture
|
|
def adapter(monkeypatch):
|
|
module = plugin_module()
|
|
context = contexts.Context()
|
|
req = encode(request(tool="read_file", arguments={"path": r"C:\Secret\key.txt"}))
|
|
resp = encode({"jsonrpc": "2.0", "id": 1, "result": {"content": []}}, "utf-16-le")
|
|
memory = b"\0" * 256 + req + b"\0" * 100 + resp + b"\0" * 256
|
|
context.layers.add_layer(BufferDataLayer(context, "mem", "mem", memory))
|
|
vad = SimpleNamespace(get_start=lambda: 0, get_size=lambda: len(memory))
|
|
proc = SimpleNamespace(
|
|
UniqueProcessId=42,
|
|
ImageFileName="agent.exe",
|
|
add_process_layer=lambda: "mem",
|
|
get_vad_root=lambda: SimpleNamespace(traverse=lambda: iter([vad])),
|
|
)
|
|
monkeypatch.setattr(module.pslist.PsList, "list_processes", lambda **kw: iter([proc]))
|
|
monkeypatch.setattr(module.utility, "array_to_string", lambda value: value)
|
|
requirements = module.MCPTrace.get_requirements()
|
|
# Synthetic EPROCESS/VADs, real Volatility memory layer and TreeGrid. This does
|
|
# not replace a real Windows kernel/symbol integration test.
|
|
monkeypatch.setattr(module.MCPTrace, "get_requirements", classmethod(lambda cls: []))
|
|
context.config["plugins.MCPTrace.kernel"] = "synthetic-kernel"
|
|
for requirement in requirements:
|
|
if requirement.name not in {"kernel", "pslist"}:
|
|
context.config[f"plugins.MCPTrace.{requirement.name}"] = requirement.default
|
|
plugin = module.MCPTrace(context, "plugins.MCPTrace")
|
|
return plugin, proc
|
|
|
|
|
|
def rows_of(plugin):
|
|
rows = []
|
|
grid = plugin.run()
|
|
grid.populate(lambda node, acc: acc.append(tuple(node.values)) or acc, rows)
|
|
return rows
|
|
|
|
|
|
def test_adapter_treegrid_with_real_framework_and_policy(adapter):
|
|
plugin, _ = adapter
|
|
plugin.config["policy"] = str(ROOT / "examples" / "policy.json")
|
|
rows = rows_of(plugin)
|
|
artifacts = [row for row in rows if row[5] != "scan-summary"]
|
|
assert len(artifacts) == 2
|
|
assert {row[4] for row in artifacts} == {"utf-8", "utf-16-le"}
|
|
assert all(row[0] == 42 and row[10] == "candidate-only" for row in artifacts)
|
|
assert json.loads(artifacts[0][12])[0]["status"] == "policy-match"
|
|
assert json.loads(rows[-1][13])["candidate_records"] == 2
|
|
|
|
|
|
def test_limit_is_in_summary(adapter):
|
|
plugin, _ = adapter
|
|
plugin.config["max-records"] = 1
|
|
rows = rows_of(plugin)
|
|
assert json.loads(rows[-1][13])["limits_hit"] == ["max-records"]
|
|
|
|
|
|
def test_suspicious_only_retains_summary(adapter):
|
|
plugin, _ = adapter
|
|
plugin.config["policy"] = str(ROOT / "examples" / "policy.json")
|
|
plugin.config["suspicious-only"] = True
|
|
rows = rows_of(plugin)
|
|
assert len(rows) == 2 and rows[0][6] == "request" and rows[1][5] == "scan-summary"
|
|
|
|
|
|
def test_suspicious_only_without_policy_fails(adapter):
|
|
plugin, _ = adapter
|
|
plugin.config["suspicious-only"] = True
|
|
with pytest.raises(ValueError, match="nonempty"):
|
|
plugin.run()
|
|
|
|
|
|
def test_cli_plugin_discovery_and_help():
|
|
executable = Path(sysconfig.get_path("scripts")) / (
|
|
"vol.exe" if sys.platform == "win32" else "vol"
|
|
)
|
|
result = subprocess.run(
|
|
[str(executable), "-p", str(ROOT / "plugins"), "windows.mcptrace.MCPTrace", "-h"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
)
|
|
assert result.returncode == 0, result.stderr + result.stdout
|
|
assert "--include-partial" in result.stdout and "--max-object-bytes" in result.stdout
|
|
|
|
|
|
def test_later_vad_failure_keeps_prior_evidence(adapter):
|
|
plugin, proc = adapter
|
|
first = next(proc.get_vad_root().traverse())
|
|
|
|
def broken_traversal():
|
|
yield first
|
|
raise exceptions.InvalidAddressException("synthetic-kernel", 0xBAD)
|
|
|
|
proc.get_vad_root = lambda: SimpleNamespace(traverse=broken_traversal)
|
|
rows = rows_of(plugin)
|
|
assert len([row for row in rows if row[5] == "complete"]) == 2
|
|
summary = json.loads(rows[-1][13])
|
|
assert summary["limits_hit"] == ["process-read-error"]
|
|
assert summary["error"] is not None
|
|
|
|
|
|
def test_byte_budget_keeps_summary_without_claiming_no_evidence(adapter):
|
|
plugin, _ = adapter
|
|
plugin.config["max-bytes"] = 128
|
|
rows = rows_of(plugin)
|
|
assert len(rows) == 1
|
|
summary = json.loads(rows[0][13])
|
|
assert summary["bytes_read"] == 128
|
|
assert summary["limits_hit"] == ["max-bytes"] and summary["absence_is_not_proof"]
|
|
|
|
|
|
def test_export_uses_real_cli_handler_and_roundtrips_raw_evidence(adapter, tmp_path):
|
|
plugin, _ = adapter
|
|
cli = CommandLine()
|
|
cli.output_dir = str(tmp_path)
|
|
plugin.set_open_method(cli.file_handler_class_factory())
|
|
plugin.config["dump"] = True
|
|
rows_of(plugin)
|
|
exported = [
|
|
json.loads(line) for line in (tmp_path / "mcptrace.pid-42.jsonl").read_text().splitlines()
|
|
]
|
|
assert exported[0]["type"] == "metadata"
|
|
assert exported[-1]["type"] == "scan-summary"
|
|
artifacts = [item for item in exported if item["type"] == "artifact"]
|
|
assert len(artifacts) == 2
|
|
for item in artifacts:
|
|
raw = base64.b64decode(item["raw_base64"])
|
|
assert hashlib.sha256(raw).hexdigest() == item["raw_sha256"]
|
|
assert json.loads(raw.decode(item["encoding"])) == item["message"]
|
|
|
|
|
|
def test_export_collision_reports_actual_filename(adapter, tmp_path):
|
|
plugin, _ = adapter
|
|
original = tmp_path / "mcptrace.pid-42.jsonl"
|
|
original.write_text("existing evidence", encoding="utf-8")
|
|
cli = CommandLine()
|
|
cli.output_dir = str(tmp_path)
|
|
plugin.set_open_method(cli.file_handler_class_factory())
|
|
plugin.config["dump"] = True
|
|
rows = rows_of(plugin)
|
|
filename = json.loads(rows[-1][13])["export"]
|
|
assert filename != original.name
|
|
assert (tmp_path / filename).is_file()
|
|
assert original.read_text() == "existing evidence"
|
|
|
|
|
|
def test_adjacent_vads_are_never_joined(adapter):
|
|
plugin, proc = adapter
|
|
data = encode(request(text="x" * 200))
|
|
context = plugin.context
|
|
context.layers.add_layer(BufferDataLayer(context, "split", "split", data))
|
|
proc.add_process_layer = lambda: "split"
|
|
vads = [
|
|
SimpleNamespace(get_start=lambda: 0, get_size=lambda: 80),
|
|
SimpleNamespace(get_start=lambda: 80, get_size=lambda: len(data) - 80),
|
|
]
|
|
proc.get_vad_root = lambda: SimpleNamespace(traverse=lambda: iter(vads))
|
|
rows = rows_of(plugin)
|
|
assert len(rows) == 1 and rows[0][5] == "scan-summary"
|