From fbcb4a92ac8c3e129a442438636bbb0fc188a6fc Mon Sep 17 00:00:00 2001 From: kwonta3wook <152182535+Kowntaewook@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:34:34 +0900 Subject: [PATCH] feat: add Windows MCPTrace Volatility 3 plugin 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. --- .github/workflows/tests.yml | 29 +++ .gitignore | 14 ++ README.md | 201 +++++++++++++++++++- docs/validation.md | 45 +++++ examples/policy.json | 23 +++ mcptrace/__init__.py | 3 + mcptrace/core.py | 359 ++++++++++++++++++++++++++++++++++++ mcptrace/memory.py | 53 ++++++ mcptrace/policy.py | 164 ++++++++++++++++ plugins/windows/__init__.py | 5 + plugins/windows/mcptrace.py | 284 ++++++++++++++++++++++++++++ pyproject.toml | 26 +++ tests/test_core.py | 190 +++++++++++++++++++ tests/test_plugin.py | 216 ++++++++++++++++++++++ tests/test_policy.py | 112 +++++++++++ 15 files changed, 1723 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/tests.yml create mode 100644 .gitignore create mode 100644 docs/validation.md create mode 100644 examples/policy.json create mode 100644 mcptrace/__init__.py create mode 100644 mcptrace/core.py create mode 100644 mcptrace/memory.py create mode 100644 mcptrace/policy.py create mode 100644 plugins/windows/__init__.py create mode 100644 plugins/windows/mcptrace.py create mode 100644 pyproject.toml create mode 100644 tests/test_core.py create mode 100644 tests/test_plugin.py create mode 100644 tests/test_policy.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..c1ab46f --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,29 @@ +name: Tests + +on: + push: + branches: [main] + pull_request: + +permissions: + contents: read + +jobs: + test: + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + python: ['3.10', '3.12'] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: ${{ matrix.python }} + - run: python -m pip install -e ".[dev]" + - run: python -m ruff check . + - run: python -m ruff format --check . + - run: python -m pytest -q diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8ccc544 --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +.venv/ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +*.egg-info/ +build/ +dist/ +*.raw +*.vmem +*.dmp +*.mem +evidence/ +output/ diff --git a/README.md b/README.md index b04644d..c8ede4b 100644 --- a/README.md +++ b/README.md @@ -1 +1,200 @@ -# MCPTrace \ No newline at end of file +# MCPTrace + +Windows 프로세스 메모리에서 MCP/JSON-RPC 메시지 후보를 추출하는 **Volatility 3 플러그인**. +수집된 메모리 이미지를 분석하고, 메시지 원문·주소·도구 인자·정책 검사 결과를 남깁니다. + +**상태: 초기 구현.** 프레임워크 패키지 버전은 `volatility3==2.28.2`로 고정했습니다. +Windows x64를 대상으로 하며 첫 실증 대상은 +Windows 11 x64입니다. 현재 검증은 합성 데이터, 실제 Volatility 메모리 레이어·출력기, +CLI 플러그인 로딩까지입니다. 실제 Windows RAM 이미지에서의 복원율과 호환성은 아직 검증하지 않았습니다. + +## 기능 + +- `windows.mcptrace.MCPTrace`: `pslist` 프로세스의 VAD를 읽어 UTF-8·UTF-16LE JSON 후보 추출. +- 문자열 안의 중괄호, 이스케이프, 중첩 객체, 청크 경계 처리. 누락된 페이지와 VAD 경계는 연결하지 않음. +- MCP 메서드, `tools/call` 도구명과 인자, 응답의 `result`/`error` 보존. +- PID·VAD 시작 주소·프로세스 가상 주소·인코딩·원문 SHA-256 표시. +- 동일 PID와 **타입까지 같은 요청 ID**를 기준으로 요청·응답 후보 제시. ID 충돌은 모호함으로 표시. +- 선택적으로 사용자 정책에 맞는 호출을 표시하고, 원문 바이트를 포함한 JSONL 내보내기. +- 프로세스마다 읽은 바이트, 누락된 바이트, 검색 제한, 오류를 요약 행으로 표시. + +메모리에서 JSON이 발견됐다는 사실만으로 실제 전송, 도구 실행, 공격 성공을 입증할 수는 없습니다. +로그·캐시·예제 문자열도 같은 형태로 남을 수 있습니다. `complete`는 아래 파서가 완전한 JSON-RPC +형태로 읽었다는 뜻이며, MCP 전체 스키마 검증이나 실제 통신 여부 검증을 의미하지 않습니다. + +## 설치 + +Python 3.10 이상과 이 저장소의 체크아웃이 필요합니다. 저장소 루트에서 실행합니다. + +Windows PowerShell: + +```powershell +py -3.12 -m venv .venv +.\.venv\Scripts\python.exe -m pip install -e ".[dev]" +.\.venv\Scripts\vol.exe -p .\plugins windows.mcptrace.MCPTrace -h +``` + +Linux 분석 호스트에서도 Windows 이미지를 대상으로 실행할 수 있습니다. + +```bash +python3 -m venv .venv +.venv/bin/python -m pip install -e '.[dev]' +.venv/bin/vol -p ./plugins windows.mcptrace.MCPTrace -h +``` + +`mcptrace/`는 설치할 Python 패키지이고 `plugins/`는 `-p`로 지정할 Volatility 플러그인 경로입니다. +두 경로를 함께 유지해야 합니다. 개발 도구가 불필요하면 `pip install -e .`로 설치합니다. +아래 예시는 가상환경의 `vol`이 PATH에 있다는 전제입니다. PowerShell에서는 위의 `vol.exe` 전체 경로를 +사용해도 됩니다. 플러그인 폴더를 옮긴 경우 `-p`에 절대 경로를 지정합니다. + +## 실행 + +먼저 Volatility가 이미지와 Windows 심볼을 정상적으로 해석하는지 확인합니다. + +```powershell +vol -f C:\Evidence\memory.raw windows.info.Info +vol -f C:\Evidence\memory.raw windows.pslist.PsList +``` + +MCP 클라이언트와 서버는 서로 다른 프로세스에 있을 수 있으므로 해당 PID들을 함께 확인합니다. +아래 PID는 예시입니다. PID를 생략하면 `pslist`가 열거한 모든 프로세스를 시도합니다. + +```powershell +vol -p .\plugins -f C:\Evidence\memory.raw windows.mcptrace.MCPTrace --pid 4242 5252 +``` + +정책 적용 및 JSONL 저장: + +```powershell +New-Item -ItemType Directory -Force .\output | Out-Null +vol -p .\plugins -f C:\Evidence\memory.raw -o .\output windows.mcptrace.MCPTrace --pid 4242 --policy .\examples\policy.json --dump +``` + +전체 TreeGrid를 JSON으로 저장할 때는 전역 렌더러 옵션을 플러그인 이름 앞에 둡니다. + +```powershell +vol -p .\plugins -f C:\Evidence\memory.raw -r json windows.mcptrace.MCPTrace --pid 4242 > findings.json +``` + +Volatility는 심볼 확보를 위해 네트워크를 사용할 수 있습니다. 사전에 필요한 심볼·캐시를 준비한 환경에서는 +전역 옵션 `--offline`을 사용할 수 있습니다. MCPTrace 자체는 외부 API나 LLM을 호출하지 않습니다. + +## 옵션 + +| 옵션 | 기본값 | 의미 | +| --- | --- | --- | +| `--pid 4242 5252` | 전체 | 지정한 PID만 분석 | +| `--include-partial` | 끔 | 파싱되지 않은 JSON-RPC 앵커 주변 조각을 별도 표시 | +| `--include-rpc` | 끔 | MCP 관련성이 확인되지 않은 일반 JSON-RPC 후보도 표시 | +| `--max-object-bytes` | 65536 | 완전한 JSON 하나의 최대 바이트 크기, 256–4194304 | +| `--max-records` | 10000 | 프로세스당 후보 수 제한, 필터 적용 **전** 기준 | +| `--max-bytes` | 0 | 프로세스당 읽기를 시도할 VAD 범위 바이트 예산, 0은 제한 없음 | +| `--policy` | 없음 | 로컬 정책 JSON 파일 | +| `--suspicious-only` | 끔 | `policy-match`인 후보와 요약만 표시, 비어 있지 않은 정책 필요 | +| `--dump` | 끔 | 표시되는 후보를 `-o` 디렉터리에 JSONL로 내보내기 | + +`--include-rpc`가 없으면 알려진 MCP 메서드, MCP 형태의 응답, 해당 요청과 같은 PID·ID의 응답을 +표시합니다. 메서드 목록은 `mcptrace/core.py`의 `MCP_METHODS`에서 관리합니다. +MCP 메서드명과 응답 형태는 다른 프로토콜과 겹칠 수 있으므로 이 필터는 관련성 휴리스틱입니다. + +## 결과 해석 + +| 필드 / 값 | 의미 | +| --- | --- | +| `VirtualAddress` | 해당 PID의 프로세스 가상 주소. 이미지 파일 오프셋이 아님 | +| `complete` | 중복 키·잘못된 인코딩 등을 거부하는 파서가 읽은 JSON-RPC 형태의 객체 | +| `fragment` | JSON-RPC 앵커부터 최대 512바이트의 원문 조각. ID·인자·응답을 추정하지 않음 | +| `candidate-only` | 같은 PID·타입별 ID의 상대 메시지가 있음. 실제 세션 연결은 미확정 | +| `ambiguous-id` | 같은 PID·ID에 서로 다른 요청 또는 응답 텍스트가 남음 | +| `unpaired` | 비교 대상 요청/응답이 발견되지 않음 | +| `not-applicable` | ID가 없거나 파싱하지 못해 연결을 시도하지 않음 | +| `CandidateAddresses` | 상대 후보 주소 목록·전체 주소 수·목록 생략 여부. 목록은 최대 256개 | +| `policy-match` | 사용자가 정의한 조건과 인자가 일치함. 악성 판정이 아님 | +| `not-evaluated` | 정책 대상 도구이지만 인자가 없거나 지원하지 않는 형식이라 검사 불가 | +| `scan-summary` | 검사 범위, 제한, 오류. `Raw` 열에 JSON으로 표시 | + +요청 ID `1`과 `"1"`은 구분하며 서로 다른 PID의 메시지는 연결하지 않습니다. +같은 JSON 텍스트가 여러 주소·인코딩으로 남은 경우 각각의 증거를 유지합니다. +서로 다른 공백이나 멤버 순서도 보수적으로 다른 텍스트로 취급합니다. +메모리 주소의 크기로 시간 순서를 추정하지 않으며 세션과 통신 방향은 복원하지 않습니다. +`RequestID`의 JSON `null`은 ID 없음과 명시적 null을 모두 나타낼 수 있으므로 원문을 확인합니다. + +`--dump` 출력은 프로세스별 `mcptrace.pid-4242.jsonl` 파일입니다. 이름이 이미 존재하면 Volatility가 +접미사를 추가하며 요약의 `export`에 실제 파일명을 기록합니다. +첫 줄은 메타데이터, 중간 줄은 아티팩트, 마지막 줄은 검사 요약입니다. +아티팩트의 `raw_base64`는 **추출된 범위의 원문 바이트**, `raw_sha256`은 그 바이트의 해시입니다. +전체 메모리 이미지의 해시가 아닙니다. `message`는 해석된 JSON이고 숫자는 Python의 정수·실수 표현을 +사용하므로 원문 검증에는 `raw_base64`를 사용합니다. 터미널의 `Raw`는 제어 문자를 이스케이프해 표시합니다. + +## 정책 + +`examples/policy.json`은 예시 도구 `read_file`, `fetch`의 허용 경로·호스트를 검사합니다. +사용 중인 MCP 서버의 **정확한 도구 이름과 인자 구조**에 맞게 수정해야 합니다. + +```json +{ + "version": 1, + "rules": [{ + "id": "outside-case-directory", + "description": "사건 분석 디렉터리 밖의 파일 접근 요청", + "severity": "medium", + "tool": "read_file", + "argument": "/path", + "operator": "windows_path_outside", + "value": ["C:\\Cases"] + }] +} +``` + +`argument`는 `params.arguments`에 대한 JSON Pointer이며 `~0`, `~1`, 배열 인덱스를 지원합니다. +완전하게 파싱된 `tools/call`의 문자열 인자에만 적용합니다. + +| 연산자 | 비교 | +| --- | --- | +| `equals` | 문자열 전체 일치, 대소문자 구분 | +| `contains` | 문자열 포함, 대소문자 구분 | +| `url_host_not_in` | HTTP(S) URL의 파싱된 호스트가 허용 목록 밖인지 비교 | +| `windows_path_outside` | 절대 Windows 경로를 정규화한 뒤 허용 디렉터리 경계 밖인지 비교 | + +URL 검사는 사용자 정보·경로의 문자열을 호스트로 오인하지 않으며 와일드카드나 암묵적 서브도메인 허용은 +없습니다. 포트는 검사 대상이 아니며 IPv6 리터럴은 현재 지원하지 않습니다. DNS·리다이렉트를 조회하지 않습니다. +경로 검사는 문자열 기준이며 정션, 심볼릭 링크, 8.3 별칭의 실제 대상은 해석하지 않습니다. +환경 변수, 상대 경로, ADS, 와일드카드 등 지원하지 않는 입력은 `not-evaluated`로 남깁니다. +규칙에서 경고가 발생하지 않았다고 안전하거나 실행되지 않았다는 뜻은 아닙니다. + +## 제한과 검증 + +- 평문 UTF-8·UTF-16LE JSON을 대상으로 합니다. 압축·암호화·바이너리 객체·V8/Python 힙 객체 재조립은 미지원입니다. +- 리터럴 `"jsonrpc": "2.0"` 앵커를 사용합니다. 키/버전 문자가 `\u` 이스케이프로 표현되면 놓칠 수 있습니다. +- 깊이 128, 앵커당 직전 객체 시작 후보 64개, 인코딩별 버퍼당 앵커 20000개의 제한이 있습니다. + 후보·앵커 제한에 도달하면 요약 카운터에 표시합니다. 한 메시지의 크기 제한은 바이트 단위입니다. +- 빠진 바이트를 0으로 채우거나 다른 VAD와 이어 붙이지 않습니다. 프로세스 읽기 오류가 나면 이미 추출한 + 결과를 유지하고 해당 프로세스의 나머지 검사를 중단합니다. +- `pslist`에서 보이지 않는 프로세스, 해제되어 VAD 밖에 남은 흔적, 현재 레이어가 읽지 못하는 페이지는 범위 밖입니다. +- `bytes_read`는 읽기에 성공한 바이트 수입니다. 후보 제한으로 중단하면 읽은 버퍼 전체가 처리됐다는 뜻은 아닙니다. + 후보 제한은 일반 JSON-RPC에도 적용되므로 잡음이 많으면 관련 메시지에 도달하기 전에 중단할 수 있습니다. +- 검색 결과 없음은 흔적 부재의 증명이 아닙니다. 모든 프로세스에 요약 행을 남기며 탐지율·복원율 수치를 주장하지 않습니다. + +테스트 실행: + +```powershell +python -m pytest -q +python -m ruff check . +python -m ruff format --check . +``` + +검증 구성과 실제 Windows 이미지에서 확인해야 할 항목은 [docs/validation.md](docs/validation.md)에 있습니다. + +## 구조 + +```text +mcptrace/core.py JSON 추출·증거 모델·ID 후보 연결 +mcptrace/memory.py 누락 페이지를 보존하는 레이어 읽기 +mcptrace/policy.py 사용자 정의 도구 인자 검사 +plugins/windows/mcptrace.py Volatility Windows 어댑터·TreeGrid·JSONL +tests/ 합성 증거와 실제 Volatility API 테스트 +``` + +참고 인터페이스: [Volatility 플러그인 문서](https://volatility3.readthedocs.io/en/stable/simple-plugin.html), +[MCP 2025-11-25 기본 프로토콜](https://modelcontextprotocol.io/specification/2025-11-25/basic), +[MCP 도구 호출](https://modelcontextprotocol.io/specification/2025-11-25/server/tools). diff --git a/docs/validation.md b/docs/validation.md new file mode 100644 index 0000000..b4e46fd --- /dev/null +++ b/docs/validation.md @@ -0,0 +1,45 @@ +# 검증 범위 + +이 문서는 실제 Windows 호환성 검증을 수행했다는 보고서가 아닙니다. +현재 테스트는 알려진 바이트를 메모리 레이어에 넣어 추출 결과와 비교합니다. + +## 자동 검증 + +로컬 검증 환경: Linux, Python 3.12, Volatility 3 패키지 2.28.2. +CI는 Ubuntu/Windows 분석 호스트와 Python 3.10/3.12 조합을 구성합니다. +CI에서 Windows 호스트를 사용하더라도 테스트 데이터는 합성 버퍼입니다. + +| 영역 | 검증 내용 | +| --- | --- | +| 원문 보존 | UTF-8/UTF-16LE, 한글·이모지, 실제 주소·원문·SHA-256 일치 | +| 파싱 | 문자열 내부 중괄호·이스케이프·중첩, 뒤에 배치된 jsonrpc 멤버, 모든 청크 분할 위치 | +| 손상 | 중복 키, 잘못된 인코딩, 비정상 JSON-RPC ID, 비유한 숫자 거부; 조각에서 필드 추정 금지 | +| 메모리 경계 | 실제 BufferDataLayer 위의 누락 페이지 시뮬레이션, 인접 VAD 비결합 | +| 후보 연결 | PID와 숫자/문자열 ID 분리, 복사본 보존, ID 충돌, 부동소수점 반올림으로 인한 오병합 방지 | +| 정책 | 디렉터리 경계, URL 사용자 정보·가짜 접미사, JSON Pointer, 미지원 입력·잘못된 규칙 | +| 어댑터 | 실제 TreeGrid, 프로세스별 제한·요약, 뒤쪽 VAD 읽기 오류 이후 기존 증거 보존 | +| 출력 | 실제 Volatility CLI 파일 핸들러, JSONL 원문 왕복·해시, 기존 파일 비덮어쓰기 | +| 로딩 | 새 CLI 프로세스에서 -p 경로로 플러그인 검색 및 도움말 생성 | + +`tests/test_plugin.py`의 EPROCESS/VAD 열거 객체와 커널 모듈은 테스트용 대역입니다. +Volatility의 Context, BufferDataLayer, TreeGrid, CLI 파일 핸들러는 실제 구현을 사용합니다. +따라서 CLI 검색·어댑터 출력까지는 검증하지만 커널 심볼 해석, 실제 페이지 테이블 변환, OS 빌드 호환성은 +이 테스트만으로 입증할 수 없습니다. + +## 실제 Windows 11 x64 이미지 확인 절차 + +1. 격리된 테스트 VM에서 OS 빌드, MCP 클라이언트·서버 버전, 전송 방식, 프로세스 PID를 기록합니다. +2. 정확한 요청·응답 JSON과 ID를 별도 기준 자료로 저장합니다. 정상 호출, 정책에 해당하는 호출, + 같은 ID의 반복 사용, 문자열 ID, 긴 응답을 각각 포함합니다. +3. 메모리를 수집하고 수집 도구·버전·시각·이미지 SHA-256, 호출부터 수집까지의 경과 시간을 기록합니다. +4. 같은 이미지에서 `windows.info.Info`와 `windows.pslist.PsList`가 정상 동작하는지 확인합니다. +5. 해당 PID에 대해 MCPTrace를 실행합니다. 처음에는 `--suspicious-only` 없이 전체 후보와 요약을 저장합니다. +6. 내보낸 `raw_base64`를 디코딩하고 SHA-256을 확인합니다. 해당 프로세스 레이어의 가상 주소와 길이로 + 다시 읽은 원문과 일치하는지 확인합니다. 가상 주소를 이미지 파일 오프셋으로 사용하지 않습니다. +7. 기준 자료와 비교해 누락·잘못된 추출·복사본·연결 모호함을 구분합니다. 예제/로그에 의도적으로 남긴 + JSON과 일반 JSON-RPC 트래픽도 넣어 실제 통신 증거와 혼동되는 범위를 확인합니다. +8. 기본 제한과 조정한 제한을 비교하고, 도달한 제한·누락 페이지·오류가 요약에 남는지 확인합니다. + +재현 자료에는 명령 전체, MCPTrace 커밋, Python/Volatility 버전, OS 빌드, 정책 파일 및 그 해시, +입력 이미지 해시, 표준 오류 로그, JSONL을 포함합니다. 실제 증거·메모리 덤프는 저장소에 커밋하지 않습니다. +실제 이미지 결과가 확보되기 전에는 특정 클라이언트 지원, 정밀도·재현율, 세션 복원 성공률을 표시하지 않습니다. diff --git a/examples/policy.json b/examples/policy.json new file mode 100644 index 0000000..4e1ea41 --- /dev/null +++ b/examples/policy.json @@ -0,0 +1,23 @@ +{ + "version": 1, + "rules": [ + { + "id": "FILE-OUTSIDE-CASE", + "description": "read_file 요청의 경로가 분석자가 지정한 사건 폴더 밖에 있음", + "severity": "medium", + "tool": "read_file", + "argument": "/path", + "operator": "windows_path_outside", + "value": ["C:\\Cases"] + }, + { + "id": "URL-OUTSIDE-ALLOWLIST", + "description": "fetch 요청의 URL 호스트가 분석자의 허용목록에 없음", + "severity": "medium", + "tool": "fetch", + "argument": "/url", + "operator": "url_host_not_in", + "value": ["api.example.test"] + } + ] +} diff --git a/mcptrace/__init__.py b/mcptrace/__init__.py new file mode 100644 index 0000000..501c426 --- /dev/null +++ b/mcptrace/__init__.py @@ -0,0 +1,3 @@ +"""MCPTrace parsing core; the Volatility adapter is in plugins/windows.""" + +__version__ = "0.1.0" diff --git a/mcptrace/core.py b/mcptrace/core.py new file mode 100644 index 0000000..251ea04 --- /dev/null +++ b/mcptrace/core.py @@ -0,0 +1,359 @@ +"""Bounded carving of literal JSON-RPC artifacts, never a proof of transmission.""" + +from __future__ import annotations + +import base64 +import bisect +import hashlib +import json +import math +import re +from collections import defaultdict +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field +from typing import Any + +MCP_METHODS = frozenset( + "initialize tools/list tools/call resources/list resources/read resources/subscribe " + "resources/unsubscribe resources/templates/list prompts/list prompts/get " + "completion/complete logging/setLevel sampling/createMessage roots/list elicitation/create " + "tasks/get tasks/result tasks/list tasks/cancel notifications/initialized " + "notifications/cancelled notifications/progress notifications/message " + "notifications/tools/list_changed notifications/resources/list_changed " + "notifications/resources/updated notifications/prompts/list_changed " + "notifications/roots/list_changed notifications/elicitation/complete " + "notifications/tasks/status".split() +) + + +def _patterns(encoding: str) -> tuple[re.Pattern, re.Pattern]: + def literal(text: str) -> bytes: + return re.escape(text.encode(encoding)) + + space = rb"[ \t\r\n]*" if encoding == "utf-8" else rb"(?:[ \t\r\n]\x00)*" + anchor = literal('"jsonrpc"') + space + literal(":") + space + literal('"2.0"') + start = literal("{") + space + literal('"') + return re.compile(anchor), re.compile(start) + + +PATTERNS = {encoding: _patterns(encoding) for encoding in ("utf-8", "utf-16-le")} + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict: + result: dict = {} + for key, value in pairs: + if key in result: + raise ValueError("Duplicate JSON member") + result[key] = value + return result + + +def _reject_constant(value: str) -> None: + raise ValueError(f"Non-JSON number: {value}") + + +def strict_loads(text: str) -> Any: + def finite_float(value: str) -> float: + number = float(value) + if not math.isfinite(number): + raise ValueError("JSON number is outside the supported finite float range") + return number + + return json.loads( + text, + object_pairs_hook=_unique_object, + parse_constant=_reject_constant, + parse_float=finite_float, + ) + + +def _valid_id(value: Any) -> bool: + # bool is an int subclass, but not a JSON-RPC request identifier. + return type(value) in (str, int) + + +def classify(obj: Any) -> tuple[str, bool] | None: + if not isinstance(obj, dict) or obj.get("jsonrpc") != "2.0": + return None + if "method" in obj: + if not isinstance(obj["method"], str) or "result" in obj or "error" in obj: + return None + if "params" in obj and not isinstance(obj["params"], (dict, list)): + return None + if "id" in obj and not _valid_id(obj["id"]): + return None + kind = "request" if "id" in obj else "notification" + return kind, obj["method"] in MCP_METHODS + if ("result" in obj) == ("error" in obj): + return None + if "id" not in obj or not (_valid_id(obj["id"]) or obj["id"] is None): + return None + if "error" in obj: + error = obj["error"] + if not isinstance(error, dict): + return None + if type(error.get("code")) is not int or not isinstance(error.get("message"), str): + return None + result = obj.get("result") + mcp_shape = isinstance(result, dict) and any( + isinstance(result.get(key), list) for key in ("tools", "content", "resources", "prompts") + ) + return "response", mcp_shape + + +@dataclass +class Artifact: + address: int + encoding: str + raw: bytes + state: str + kind: str + obj: dict | None = None + mcp: bool = False + pid: int = 0 + vad_start: int = 0 + relation: str = "not-applicable" + candidate_addresses: list[int] = field(default_factory=list) + candidate_count: int = 0 + findings: list[dict] = field(default_factory=list) + + @property + def method(self) -> str: + return (self.obj or {}).get("method", "") + + @property + def request_id(self) -> str | int | None: + return (self.obj or {}).get("id") + + @property + def tool(self) -> str: + params = (self.obj or {}).get("params") + if self.method == "tools/call" and isinstance(params, dict): + name = params.get("name") + return name if isinstance(name, str) else "" + return "" + + @property + def sha256(self) -> str: + return hashlib.sha256(self.raw).hexdigest() + + def to_dict(self) -> dict: + return { + "pid": self.pid, + "vad_start": self.vad_start, + "virtual_address": self.address, + "encoding": self.encoding, + "state": self.state, + "kind": self.kind, + "mcp_candidate": self.mcp, + "request_id": self.request_id, + "method": self.method, + "tool": self.tool, + "message": self.obj, + "raw_base64": base64.b64encode(self.raw).decode("ascii"), + "raw_sha256": self.sha256, + "byte_length": len(self.raw), + "relation": self.relation, + "candidate_virtual_addresses": self.candidate_addresses, + "candidate_address_count": self.candidate_count, + "candidate_addresses_truncated": self.candidate_count > len(self.candidate_addresses), + "findings": self.findings, + } + + +@dataclass +class ScanStats: + bytes_read: int = 0 + unreadable_bytes: int = 0 + anchor_limit_hits: int = 0 + candidate_limit_hits: int = 0 + + +def _object_end(data: bytes, start: int, step: int, limit: int) -> int | None: + """Respect quoted braces, escapes, arrays and a bounded nesting depth.""" + stack: list[int] = [] + quoted = escaped = False + for pos in range(start, min(len(data), start + limit) - step + 1, step): + char = data[pos] if step == 1 else int.from_bytes(data[pos : pos + 2], "little") + if char == 0: + return None + if quoted: + if escaped: + escaped = False + elif char == 92: + escaped = True + elif char == 34: + quoted = False + continue + if char == 34: + quoted = True + elif char in (123, 91): + stack.append(char) + if len(stack) > 128: + return None + elif char in (125, 93): + if not stack or stack.pop() != {125: 123, 93: 91}[char]: + return None + if not stack: + return pos + step + return None + + +def _carve_buffer( + data: bytes, base: int, maximum: int, partial: bool, stats: ScanStats +) -> Iterator[Artifact]: + for encoding, (anchors, starts_pattern) in PATTERNS.items(): + step = 1 if encoding == "utf-8" else 2 + starts = [m.start() for m in starts_pattern.finditer(data)] + parsed: dict[int, Artifact | None] = {} + complete: dict[int, Artifact] = {} + unmatched: list[int] = [] + for count, match in enumerate(anchors.finditer(data)): + if count >= 20000: + stats.anchor_limit_hits += 1 + break + anchor = match.start() + low = bisect.bisect_left(starts, max(0, anchor - maximum + 1)) + high = bisect.bisect_right(starts, anchor) + if high - low > 64: + stats.candidate_limit_hits += 1 + low = high - 64 + found = False + for index in range(high - 1, low - 1, -1): + start = starts[index] + if (anchor - start) % step: + continue + if start not in parsed: + parsed[start] = None + end = _object_end(data, start, step, maximum) + if end is not None and end > anchor: + raw = data[start:end] + try: + obj = strict_loads(raw.decode(encoding, errors="strict")) + classification = classify(obj) + except (ValueError, UnicodeError, RecursionError): + classification = None + if classification: + kind, mcp = classification + parsed[start] = Artifact( + base + start, encoding, raw, "complete", kind, obj, mcp + ) + record = parsed[start] + if record and record.address + len(record.raw) > base + anchor: + complete[start] = record + found = True + break + if not found: + unmatched.append(anchor) + yield from complete.values() + if partial: + ranges = sorted((start, start + len(rec.raw)) for start, rec in complete.items()) + range_starts = [start for start, _ in ranges] + for anchor in unmatched: + index = bisect.bisect_right(range_starts, anchor) - 1 + if index >= 0 and anchor < ranges[index][1]: + continue + end = min(len(data), anchor + min(maximum, 512)) + end -= (end - anchor) % step + for pos in range(anchor, end, step): + if data[pos : pos + step] == b"\0" * step: + end = pos + break + raw = data[anchor:end] + if raw: + # Raw fragment only: no guessed ID, arguments, or repaired JSON. + yield Artifact(base + anchor, encoding, raw, "fragment", "fragment") + + +def scan_chunks( + chunks: Iterable[tuple[int, bytes]], + *, + max_object_bytes: int = 65536, + include_partial: bool = False, + stats: ScanStats | None = None, +) -> Iterator[Artifact]: + """Carve contiguous virtual bytes, flushing at gaps and at every VAD boundary. + + Keep two windows so a late jsonrpc member and a chunk-spanning object survive. + No bytes are invented for unreadable pages. Addresses are process virtual addresses. + """ + if not 256 <= max_object_bytes <= 4 * 1024 * 1024: + raise ValueError("max_object_bytes must be between 256 and 4194304") + stats = stats if stats is not None else ScanStats() + pending = b"" + base = 0 + seen: set[tuple[int, str]] = set() + + def emit(final: bool) -> Iterator[Artifact]: + safe_end = base + len(pending) if final else base + len(pending) - max_object_bytes + for record in _carve_buffer(pending, base, max_object_bytes, include_partial, stats): + key = record.address, record.encoding + if record.address < safe_end and key not in seen: + seen.add(key) + yield record + + for address, data in chunks: + if not data: + continue + if pending and address != base + len(pending): + yield from emit(True) + pending = b"" + if not pending: + base = address + pending += data + yield from emit(False) + keep = 2 * max_object_bytes + if len(pending) > keep: + base += len(pending) - keep + pending = pending[-keep:] + if pending: + yield from emit(True) + + +def correlate(records: list[Artifact]) -> None: + """ID matches are candidates, never inferred sessions or verified executions.""" + groups: dict[tuple, list[Artifact]] = defaultdict(list) + for record in records: + record.candidate_addresses = [] + record.candidate_count = 0 + if record.state != "complete" or record.request_id is None: + record.relation = "not-applicable" + continue + groups[(record.pid, type(record.request_id), record.request_id)].append(record) + for group in groups.values(): + requests = [r for r in group if r.kind == "request"] + responses = [r for r in group if r.kind == "response"] + + def identities(items: list[Artifact]) -> set[str]: + # Do not collapse distinct numeric lexemes through float rounding. + # Different JSON formatting is conservatively considered ambiguous. + return {r.raw.decode(r.encoding) for r in items} + + ambiguous = len(identities(requests)) > 1 or len(identities(responses)) > 1 + request_addresses = sorted({r.address for r in requests}) + response_addresses = sorted({r.address for r in responses}) + for record in group: + opposite = response_addresses if record.kind == "request" else request_addresses + record.candidate_count = len(opposite) + # Bound repeated address lists; the individual artifacts are retained. + record.candidate_addresses = opposite[:256] + record.relation = ( + "ambiguous-id" if ambiguous else "candidate-only" if opposite else "unpaired" + ) + + +def relevant_records(records: list[Artifact], include_rpc: bool = False) -> list[Artifact]: + """Retain generic responses only when MCP-shaped or sharing a process/typed ID.""" + ids = { + (r.pid, type(r.request_id), r.request_id) + for r in records + if r.mcp and r.kind == "request" and r.request_id is not None + } + return [ + r + for r in records + if include_rpc + or r.mcp + or r.state == "fragment" + or (r.kind == "response" and (r.pid, type(r.request_id), r.request_id) in ids) + ] diff --git a/mcptrace/memory.py b/mcptrace/memory.py new file mode 100644 index 0000000..8400023 --- /dev/null +++ b/mcptrace/memory.py @@ -0,0 +1,53 @@ +"""Read a Volatility layer without filling missing evidence with zero bytes.""" + +from collections.abc import Iterator + +from mcptrace.core import ScanStats + + +def readable_chunks( + layer, + start: int, + size: int, + stats: ScanStats, + invalid_address: type[Exception], + chunk_bytes: int = 1024 * 1024, +) -> Iterator[tuple[int, bytes]]: + if size < 0 or start < 0 or chunk_bytes < 4096: + raise ValueError("Invalid virtual range or read chunk size") + end = start + size + offset = start + while offset < end: + amount = min(chunk_bytes, end - offset) + try: + data = layer.read(offset, amount, pad=False) + if len(data) != amount: + raise invalid_address(layer.name, offset, "Short read") + except invalid_address: + # One absent page must not discard other resident pages in this chunk. + page = offset + run_start = page + run = bytearray() + while page < offset + amount: + page_amount = min(4096 - page % 4096, offset + amount - page) + try: + part = layer.read(page, page_amount, pad=False) + if len(part) != page_amount: + raise invalid_address(layer.name, page, "Short read") + except invalid_address: + if run: + yield run_start, bytes(run) + run.clear() + stats.unreadable_bytes += page_amount + else: + if not run: + run_start = page + run.extend(part) + stats.bytes_read += len(part) + page += page_amount + if run: + yield run_start, bytes(run) + else: + stats.bytes_read += len(data) + yield offset, data + offset += amount diff --git a/mcptrace/policy.py b/mcptrace/policy.py new file mode 100644 index 0000000..be48500 --- /dev/null +++ b/mcptrace/policy.py @@ -0,0 +1,164 @@ +"""Explicit analyst policy checks on complete tools/call arguments only.""" + +from __future__ import annotations + +import ntpath +import re +from urllib.parse import urlsplit + +from mcptrace.core import Artifact + +OPERATORS = {"equals", "contains", "url_host_not_in", "windows_path_outside"} + + +def validate_policy(document: object) -> list[dict]: + if ( + not isinstance(document, dict) + or type(document.get("version")) is not int + or document["version"] != 1 + or set(document) != {"version", "rules"} + ): + raise ValueError("Policy must be an object with version 1") + rules = document.get("rules") + if not isinstance(rules, list) or len(rules) > 100: + raise ValueError("Policy rules must be an array of at most 100 entries") + seen = set() + for rule in rules: + if not isinstance(rule, dict): + raise ValueError("Each rule must be an object") + if set(rule) - {"id", "description", "severity", "tool", "argument", "operator", "value"}: + raise ValueError("Unknown rule field") + for key in ("id", "description", "tool", "argument", "operator", "severity"): + if not isinstance(rule.get(key), str) or not rule[key]: + raise ValueError(f"Rule requires a nonempty {key}") + if rule["id"] in seen: + raise ValueError("Duplicate rule ID") + seen.add(rule["id"]) + if rule["severity"] not in {"low", "medium", "high"}: + raise ValueError("Severity must be low, medium, or high") + if ( + not rule["argument"].startswith("/") + or re.search(r"~(?![01])", rule["argument"]) + or rule["operator"] not in OPERATORS + ): + raise ValueError("Invalid argument JSON pointer or operator") + value = rule.get("value") + if rule["operator"] in {"equals", "contains"}: + if not isinstance(value, str) or not value: + raise ValueError("String operators require a nonempty string value") + else: + if ( + not isinstance(value, list) + or not value + or any(not isinstance(v, str) or not v for v in value) + ): + raise ValueError("Allowlist operators require a nonempty string array") + if rule["operator"] == "windows_path_outside" and any( + _absolute_windows(v) is None for v in value + ): + raise ValueError("Allowed Windows paths must be absolute, without variables") + if rule["operator"] == "url_host_not_in" and any(_host(v) is None for v in value): + raise ValueError("Allowed URL hosts must be host names, without ports or wildcards") + return rules + + +def _pointer(obj: object, pointer: str) -> object: + for token in pointer.split("/")[1:]: + token = token.replace("~1", "/").replace("~0", "~") + if isinstance(obj, dict): + obj = obj[token] + elif isinstance(obj, list) and re.fullmatch(r"0|[1-9][0-9]*", token): + obj = obj[int(token)] + else: + raise KeyError(pointer) + return obj + + +def _host(value: str) -> str | None: + if not value or any(c in value for c in "/:@* \\?#") or any(ord(c) < 33 for c in value): + return None + try: + return value.rstrip(".").encode("idna").decode("ascii").lower() or None + except UnicodeError: + return None + + +def _absolute_windows(value: str) -> str | None: + value = value.replace("/", "\\") + if value.startswith("\\\\?\\UNC\\"): + value = "\\\\" + value[8:] + elif value.startswith("\\\\?\\"): + value = value[4:] + if any(ord(c) < 32 for c in value) or "%" in value or "$" in value: + return None + if value.startswith("\\\\.\\"): + return None + drive, tail = ntpath.splitdrive(value) + if not drive or not tail.startswith("\\"): + return None + if not re.fullmatch(r"[a-zA-Z]:", drive) and not re.fullmatch(r"\\\\[^\\]+\\[^\\]+", drive): + return None + # Alternate streams, wildcards, and Win32-trimmed components need filesystem + # semantics we cannot recover here. Do not silently classify these as allowed. + if any(c in tail for c in ':*?"<>|'): + return None + if any(part not in {"", ".", ".."} and part.endswith((".", " ")) for part in tail.split("\\")): + return None + return ntpath.normcase(ntpath.normpath(value)) + + +def _check(value: str, rule: dict) -> bool | None: + operator, expected = rule["operator"], rule["value"] + if operator == "equals": + return value == expected + if operator == "contains": + return expected in value + if operator == "url_host_not_in": + try: + if any(ord(c) < 33 for c in value): + return None + parsed = urlsplit(value) + host = _host(parsed.hostname or "") + if parsed.scheme.lower() not in {"http", "https"} or host is None: + return None + return host not in {_host(v) for v in expected} + except ValueError: + return None + path = _absolute_windows(value) + if path is None: + return None + for allowed in expected: + root = _absolute_windows(allowed) + try: + if root is not None and ntpath.commonpath([path, root]) == root: + return False + except ValueError: # Different drives. + pass + return True + + +def evaluate(record: Artifact, rules: list[dict]) -> list[dict]: + if record.state != "complete" or record.method != "tools/call": + return [] + params = (record.obj or {}).get("params") + arguments = params.get("arguments") if isinstance(params, dict) else None + findings = [] + for rule in rules: + if rule["tool"] != record.tool: + continue + try: + value = _pointer(arguments, rule["argument"]) + matched = _check(value, rule) if isinstance(value, str) else None + except (KeyError, IndexError): + matched = None + if matched is not False: + findings.append( + { + "rule_id": rule["id"], + "severity": rule["severity"], + "status": "policy-match" if matched else "not-evaluated", + "argument": rule["argument"], + "description": rule["description"], + } + ) + return findings diff --git a/plugins/windows/__init__.py b/plugins/windows/__init__.py new file mode 100644 index 0000000..1ca1489 --- /dev/null +++ b/plugins/windows/__init__.py @@ -0,0 +1,5 @@ +"""Extend the Windows plugin namespace without hiding built-in plugins.""" + +from pkgutil import extend_path + +__path__ = extend_path(__path__, __name__) diff --git a/plugins/windows/mcptrace.py b/plugins/windows/mcptrace.py new file mode 100644 index 0000000..40a331a --- /dev/null +++ b/plugins/windows/mcptrace.py @@ -0,0 +1,284 @@ +"""MCP artifacts in Windows process virtual memory (Volatility 3).""" + +from __future__ import annotations + +import json +import logging +from dataclasses import asdict +from pathlib import Path + +from volatility3.framework import exceptions, interfaces, renderers +from volatility3.framework.configuration import requirements +from volatility3.framework.objects import utility +from volatility3.framework.renderers import format_hints +from volatility3.plugins.windows import pslist + +from mcptrace import __version__ +from mcptrace.core import ScanStats, correlate, relevant_records, scan_chunks, strict_loads +from mcptrace.memory import readable_chunks +from mcptrace.policy import evaluate, validate_policy + +LOG = logging.getLogger(__name__) + + +def _json(value) -> str: + # Escape terminal control characters and preserve string versus numeric IDs. + return json.dumps(value, ensure_ascii=True, separators=(",", ":")) + + +class MCPTrace(interfaces.plugins.PluginInterface): + """Extract MCP/JSON-RPC candidates and policy findings from Windows VADs.""" + + _required_framework_version = (2, 28, 0) + _version = (0, 1, 0) + + @classmethod + def get_requirements(cls): + return [ + requirements.ModuleRequirement( + name="kernel", description="Windows x64 kernel", architectures=["Intel64"] + ), + requirements.VersionRequirement( + name="pslist", component=pslist.PsList, version=(3, 0, 0) + ), + requirements.ListRequirement( + name="pid", + description="Only scan these process IDs", + element_type=int, + optional=True, + ), + requirements.IntRequirement( + name="max-object-bytes", + description="Maximum complete JSON size in bytes", + default=65536, + optional=True, + ), + requirements.IntRequirement( + name="max-records", + description="Maximum candidate artifacts per process", + default=10000, + optional=True, + ), + requirements.IntRequirement( + name="max-bytes", + description="VAD bytes to attempt per process; 0 is unlimited", + default=0, + optional=True, + ), + requirements.BooleanRequirement( + name="include-partial", + description="Include raw, unparsed JSON-RPC fragments", + default=False, + optional=True, + ), + requirements.BooleanRequirement( + name="include-rpc", + description="Also show unrelated generic JSON-RPC artifacts", + default=False, + optional=True, + ), + requirements.StringRequirement( + name="policy", + description="Local JSON policy path (optional; no built-in verdicts)", + optional=True, + ), + requirements.BooleanRequirement( + name="suspicious-only", + description="Show policy matches only, plus scan summaries", + default=False, + optional=True, + ), + requirements.BooleanRequirement( + name="dump", + description="Export JSONL with exact raw bytes and SHA-256 using -o", + default=False, + optional=True, + ), + ] + + def _scan_process(self, proc, stats, records, limits): + layer = self.context.layers[proc.add_process_layer()] + pid = int(proc.UniqueProcessId) + seen = set() + attempted = 0 + maximum = self.config["max-object-bytes"] + budget = self.config["max-bytes"] + for vad in proc.get_vad_root().traverse(): + start, size = int(vad.get_start()), int(vad.get_size()) + if start < 0 or size <= 0: + continue + if budget and attempted >= budget: + limits.append("max-bytes") + break + amount = min(size, budget - attempted) if budget else size + chunks = readable_chunks( + layer, start, amount, stats, exceptions.InvalidAddressException + ) + for record in scan_chunks( + chunks, + max_object_bytes=maximum, + include_partial=self.config["include-partial"], + stats=stats, + ): + key = record.address, record.encoding + if key in seen: + continue + seen.add(key) + record.pid, record.vad_start = pid, start + records.append(record) + if len(records) >= self.config["max-records"]: + limits.append("max-records") + return + attempted += amount + if amount < size: + limits.append("max-bytes") + break + + @staticmethod + def _row(pid, name, record): + return ( + pid, + name, + format_hints.Hex(record.vad_start), + format_hints.Hex(record.address), + record.encoding, + record.state, + record.kind, + _json(record.request_id), + _json(record.method), + _json(record.tool), + record.relation, + _json( + { + "addresses": [hex(a) for a in record.candidate_addresses], + "total": record.candidate_count, + "truncated": record.candidate_count > len(record.candidate_addresses), + } + ), + _json(record.findings), + _json(record.raw.decode(record.encoding, errors="backslashreplace")), + record.sha256, + ) + + @staticmethod + def _summary_row(pid, name, summary): + return ( + pid, + name, + format_hints.Hex(0), + format_hints.Hex(0), + "", + "scan-summary", + "scan-summary", + "", + "", + "", + "", + "", + "", + _json(summary), + "", + ) + + def _export(self, pid, process, records, summary): + with self.open(f"mcptrace.pid-{pid}.jsonl") as output: + header = { + "type": "metadata", + "mcptrace_version": __version__, + "pid": pid, + "process": process, + "address_space": "process-virtual", + "correlation": "same PID and typed ID only; session/direction not recovered", + } + output.write((_json(header) + "\n").encode()) + for record in records: + output.write((_json({"type": "artifact", **record.to_dict()}) + "\n").encode()) + output.write((_json({"type": "scan-summary", **summary}) + "\n").encode()) + # The CLI may rename an existing destination when the handler closes. + return output.preferred_filename + + def _generator(self, rules): + filters = pslist.PsList.create_pid_filter(self.config.get("pid")) + processes = pslist.PsList.list_processes( + context=self.context, kernel_module_name=self.config["kernel"], filter_func=filters + ) + for proc in processes: + pid = int(proc.UniqueProcessId) + try: + name = utility.array_to_string(proc.ImageFileName) + except exceptions.InvalidAddressException: + name = "" + display_name = _json(name)[1:-1] + stats = ScanStats() + error = None + records, limits = [], [] + try: + self._scan_process(proc, stats, records, limits) + except (exceptions.InvalidAddressException, exceptions.SymbolError) as exc: + # Keep evidence recovered before a later VAD or symbol failure. + limits.append("process-read-error") + error = str(exc) + LOG.warning("PID %d: %s", pid, exc) + correlate(records) + for record in records: + record.findings = evaluate(record, rules) + relevant = relevant_records(records, self.config["include-rpc"]) + visible = relevant + if self.config["suspicious-only"]: + visible = [ + r for r in relevant if any(f["status"] == "policy-match" for f in r.findings) + ] + summary = { + **asdict(stats), + "candidate_records": len(records), + "relevant_records": len(relevant), + "displayed_records": len(visible), + "limits_hit": limits, + "error": error, + "scope": "readable bytes exposed by pslist process VAD layers", + "absence_is_not_proof": True, + } + if limits or stats.anchor_limit_hits or stats.candidate_limit_hits: + LOG.warning("PID %d scan is limited: %s", pid, _json(summary)) + if self.config["dump"]: + summary["export"] = self._export(pid, name, visible, summary) + for record in sorted(visible, key=lambda r: (r.address, r.encoding)): + yield 0, self._row(pid, display_name, record) + yield 0, self._summary_row(pid, display_name, summary) + + def run(self): + maximum = self.config["max-object-bytes"] + if not 256 <= maximum <= 4 * 1024 * 1024: + raise ValueError("--max-object-bytes must be between 256 and 4194304") + if not 1 <= self.config["max-records"] <= 100000: + raise ValueError("--max-records must be between 1 and 100000") + if self.config["max-bytes"] < 0: + raise ValueError("--max-bytes cannot be negative") + rules = [] + if self.config.get("policy"): + policy_path = Path(self.config["policy"]) + if policy_path.stat().st_size > 1024 * 1024: + raise ValueError("Policy file exceeds 1 MiB") + rules = validate_policy(strict_loads(policy_path.read_text(encoding="utf-8"))) + if self.config["suspicious-only"] and not rules: + raise ValueError("--suspicious-only requires a nonempty --policy") + return renderers.TreeGrid( + [ + ("PID", int), + ("Process", str), + ("VADStart", format_hints.Hex), + ("VirtualAddress", format_hints.Hex), + ("Encoding", str), + ("State", str), + ("Kind", str), + ("RequestID", str), + ("Method", str), + ("Tool", str), + ("Relation", str), + ("CandidateAddresses", str), + ("Findings", str), + ("Raw", str), + ("SHA256", str), + ], + self._generator(rules), + ) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a255df8 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,26 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "mcptrace" +version = "0.1.0" +description = "Evidence-preserving MCP artifact analysis for Volatility 3 on Windows" +requires-python = ">=3.10" +dependencies = ["volatility3==2.28.2"] + +[project.optional-dependencies] +dev = ["pytest>=8,<10", "ruff>=0.9,<1"] + +[tool.setuptools.packages.find] +include = ["mcptrace*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.ruff] +target-version = "py310" +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "I", "UP", "B"] diff --git a/tests/test_core.py b/tests/test_core.py new file mode 100644 index 0000000..0359585 --- /dev/null +++ b/tests/test_core.py @@ -0,0 +1,190 @@ +import base64 +import hashlib +import json + +import pytest + +from mcptrace.core import ScanStats, classify, correlate, relevant_records, scan_chunks + + +def request(rid=1, text="hello", tool="echo", arguments=None): + return { + "jsonrpc": "2.0", + "id": rid, + "method": "tools/call", + "params": { + "name": tool, + "arguments": arguments if arguments is not None else {"text": text}, + }, + } + + +def encode(obj, encoding="utf-8"): + return json.dumps(obj, ensure_ascii=False, separators=(",", ":")).encode(encoding) + + +def carve(obj, **kwargs): + return list(scan_chunks([(0x1000, encode(obj))], **kwargs))[0] + + +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16-le"]) +@pytest.mark.parametrize( + "text", ["normal", "unbalanced }", "unbalanced {", 'quote \\" and }', "한글😀"] +) +def test_intact_message_and_exact_provenance(encoding, text): + obj = request(text=text) + raw = encode(obj, encoding) + records = list(scan_chunks([(0x1001, b"xyz" + raw + b"\0" * 8)])) + assert len(records) == 1 + record = records[0] + assert record.address == 0x1004 + assert record.obj == obj and record.encoding == encoding + assert record.raw == raw and record.sha256 == hashlib.sha256(raw).hexdigest() + assert base64.b64decode(record.to_dict()["raw_base64"]) == raw + + +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16-le"]) +def test_last_jsonrpc_member_and_every_chunk_boundary(encoding): + obj = { + "params": {"name": "echo", "arguments": {"a": {"b": "} { 한글"}}}, + "method": "tools/call", + "id": 7, + "jsonrpc": "2.0", + } + raw = encode(obj, encoding) + for split in range(1, len(raw)): + records = list(scan_chunks([(123, raw[:split]), (123 + split, raw[split:])])) + assert len(records) == 1 and records[0].obj == obj + + +def test_large_stream_overlap_keeps_one_record_per_address(): + raw = encode(request()) + data = b"\0" * 500 + raw + b"\0" * 2000 + raw + b"\0" * 600 + records = list( + scan_chunks( + ((i, data[i : i + 128]) for i in range(0, len(data), 128)), max_object_bytes=256 + ) + ) + assert [r.address for r in records] == [500, 500 + len(raw) + 2000] + + +def test_gap_is_not_filled_or_joined(): + raw = encode(request(text="x" * 100)) + chunks = [(0x1000, raw[:80]), (0x3000, raw[80:])] + records = list(scan_chunks(chunks, include_partial=True)) + assert records and all(r.state == "fragment" for r in records) + assert all(r.obj is None and r.request_id is None and not r.tool for r in records) + + +@pytest.mark.parametrize( + "raw", + [ + b'{"jsonrpc":"2.0","id":1,"method":"tools/list","id":2}', + b'{"jsonrpc":"2.0","id":true,"method":"tools/list"}', + b'{"jsonrpc":"2.0","id":1.5,"method":"tools/list"}', + b'{"jsonrpc":"2.0","id":1,"result":{},"error":{"code":1,"message":"x"}}', + b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"x":NaN}}', + b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"x":1e999}}', + b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"x":"\xff"}}', + ], +) +def test_malformed_data_not_silently_repaired(raw): + assert list(scan_chunks([(0, raw)])) == [] + fragments = list(scan_chunks([(0, raw)], include_partial=True)) + assert fragments and all(r.state == "fragment" and r.obj is None for r in fragments) + + +def test_plain_text_mentions_and_non_rpc_json_do_not_count(): + assert list(scan_chunks([(0, b'tools/call jsonrpc {"method":"tools/call"}')])) == [] + + +def test_notifications_and_errors(): + assert carve({"jsonrpc": "2.0", "method": "notifications/initialized"}).kind == "notification" + response = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": "parse"}} + assert carve(response).kind == "response" + assert classify({"jsonrpc": "2.0", "id": 1, "error": {"code": True, "message": "x"}}) is None + + +def test_pid_and_id_type_boundaries(): + req = carve(request(1)) + other_pid = carve({"jsonrpc": "2.0", "id": 1, "result": {}}) + other_pid.pid = 2 + string_id = carve({"jsonrpc": "2.0", "id": "1", "result": {}}) + records = [req, other_pid, string_id] + correlate(records) + assert all(r.relation == "unpaired" for r in records) + assert relevant_records(records) == [req] + + +def test_pair_is_only_a_candidate_and_copies_keep_addresses(): + req, duplicate = carve(request()), carve(request()) + duplicate.address += 500 + response = carve({"jsonrpc": "2.0", "id": 1, "result": {}}) + response.address += 1000 + records = [req, duplicate, response] + correlate(records) + assert all(r.relation == "candidate-only" for r in records) + assert response.candidate_addresses == [req.address, duplicate.address] + assert len(relevant_records(records)) == 3 + + +def test_id_collision_is_ambiguous_without_inventing_session(): + records = [ + carve(request(text="session A")), + carve(request(text="session B")), + carve({"jsonrpc": "2.0", "id": 1, "result": {}}), + ] + correlate(records) + assert all(r.relation == "ambiguous-id" for r in records) + assert all("session" not in r.to_dict() for r in records) + + +def test_same_json_in_two_encodings_preserves_both_locations(): + obj = request() + first = encode(obj) + records = list(scan_chunks([(0, first + b"\0" * 8 + encode(obj, "utf-16-le"))])) + assert len(records) == 2 + assert {r.encoding for r in records} == {"utf-8", "utf-16-le"} + assert records[0].address != records[1].address + + +def test_oversized_messages_are_not_complete(): + records = list( + scan_chunks( + [(0, encode(request(text="x" * 400)))], max_object_bytes=256, include_partial=True + ) + ) + assert records and all(r.state == "fragment" for r in records) + + +def test_backtrack_limit_is_observable(): + obj = { + "params": {"x": [{"v": i} for i in range(70)]}, + "method": "tools/list", + "id": 1, + "jsonrpc": "2.0", + } + stats = ScanStats() + list(scan_chunks([(0, encode(obj))], stats=stats)) + assert stats.candidate_limit_hits > 0 + + +def test_candidate_addresses_are_bounded_without_dropping_artifacts(): + req = carve(request()) + responses = [carve({"jsonrpc": "2.0", "id": 1, "result": {}}) for _ in range(300)] + for i, response in enumerate(responses): + response.address += i * 100 + records = [req, *responses] + correlate(records) + assert len(relevant_records(records)) == 301 + assert len(req.candidate_addresses) == 256 and req.candidate_count == 300 + assert req.to_dict()["candidate_addresses_truncated"] is True + + +def test_distinct_numeric_lexemes_are_not_collapsed_by_float_rounding(): + template = b'{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"value":%s}}' + data = template % b"1.00000000000000001" + b"\0" + template % b"1.00000000000000002" + records = list(scan_chunks([(0, data)])) + correlate(records) + assert len(records) == 2 + assert all(r.relation == "ambiguous-id" for r in records) diff --git a/tests/test_plugin.py b/tests/test_plugin.py new file mode 100644 index 0000000..b25668b --- /dev/null +++ b/tests/test_plugin.py @@ -0,0 +1,216 @@ +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" diff --git a/tests/test_policy.py b/tests/test_policy.py new file mode 100644 index 0000000..12740b1 --- /dev/null +++ b/tests/test_policy.py @@ -0,0 +1,112 @@ +import copy + +import pytest +from test_core import carve, request + +from mcptrace.policy import evaluate, validate_policy + + +def rule(operator, value, argument="/path"): + return { + "id": "R1", + "description": "Analyst-defined policy", + "severity": "medium", + "tool": "read_file", + "argument": argument, + "operator": operator, + "value": value, + } + + +def run_rule(value, operator="windows_path_outside", expected=None): + spec = rule(operator, expected if expected is not None else [r"C:\Cases"]) + rules = validate_policy({"version": 1, "rules": [spec]}) + return evaluate(carve(request(tool="read_file", arguments={"path": value})), rules) + + +@pytest.mark.parametrize( + "value", + [r"C:\Cases\evidence.txt", r"c:\CASES\sub\a.txt", r"C:/Cases/a.txt", r"\\?\C:\Cases\a.txt"], +) +def test_allowed_path_boundaries(value): + assert run_rule(value) == [] + + +@pytest.mark.parametrize("value", [r"C:\Cases2\a.txt", r"C:\Cases\..\secret.txt", r"D:\Cases\a"]) +def test_outside_paths_match(value): + assert run_rule(value)[0]["status"] == "policy-match" + + +@pytest.mark.parametrize( + "value", + [ + r"C:relative.txt", + r"\rooted.txt", + r"%USERPROFILE%\a", + 3, + None, + r"C:\Cases\a.txt:secret", + r"C:\Cases\a. ", + r"C:\Cases\*.txt", + ], +) +def test_unresolved_inputs_are_not_marked_safe(value): + assert run_rule(value)[0]["status"] == "not-evaluated" + + +@pytest.mark.parametrize( + "url,matched", + [ + ("https://API.EXAMPLE.TEST./x", False), + ("https://api.example.test.evil.test/x", True), + ("https://api.example.test@evil.test/x", True), + ("https://evil.test/?allowed=api.example.test", True), + ("https://api.example.test:443/path", False), + ], +) +def test_url_host_uses_parsed_authority(url, matched): + findings = run_rule(url, "url_host_not_in", ["api.example.test"]) + assert bool(findings) == matched + + +def test_policy_requires_exact_tool_and_complete_request(): + rules = validate_policy({"version": 1, "rules": [rule("contains", "secret")]}) + assert evaluate(carve(request(tool="another_tool", arguments={"path": "secret"})), rules) == [] + record = carve(request(tool="read_file", arguments={"path": "secret"})) + assert evaluate(record, rules)[0]["status"] == "policy-match" + record.state = "fragment" + assert evaluate(record, rules) == [] + + +def test_json_pointer_and_missing_arguments(): + spec = rule("equals", "secret", "/nested/a~1b/0") + rules = validate_policy({"version": 1, "rules": [spec]}) + rec = carve(request(tool="read_file", arguments={"nested": {"a/b": ["secret"]}})) + assert evaluate(rec, rules)[0]["status"] == "policy-match" + rec.obj["params"]["arguments"] = {} + assert evaluate(rec, rules)[0]["status"] == "not-evaluated" + + +@pytest.mark.parametrize( + "change", + [ + {"operator": "execute"}, + {"severity": "critical"}, + {"extra": "ignored?"}, + {"argument": "path"}, + {"value": []}, + {"value": ["relative"]}, + {"argument": "/bad~2pointer"}, + ], +) +def test_bad_policy_is_rejected(change): + spec = rule("windows_path_outside", [r"C:\Cases"]) + spec.update(change) + with pytest.raises(ValueError): + validate_policy({"version": 1, "rules": [spec]}) + + +def test_duplicate_rule_ids_rejected(): + spec = rule("contains", "x") + with pytest.raises(ValueError): + validate_policy({"version": 1, "rules": [spec, copy.deepcopy(spec)]})