Open_Ultra_SAST/tests/test_engine_trace.py

"""Offline controls for the trace experiment; no Docker, Joern or network."""

from __future__ import annotations

import importlib.util
import json
import subprocess
import sys
from pathlib import Path

import pytest

from openultrasast.cpg.backend import BEGIN, END, CpgResult, JoernBackend
from openultrasast.model.specs import taint_specs
from openultrasast.model.taint import request_params
from openultrasast.model.trace import parse_trace

SCRIPT = Path(__file__).resolve().parents[1] / "benchmarks/learn/engine_trace.py"
_spec = importlib.util.spec_from_file_location("engine_trace", SCRIPT)
assert _spec and _spec.loader
trace = importlib.util.module_from_spec(_spec)
with pytest.MonkeyPatch.context() as patch:
    patch.syspath_prepend(str(SCRIPT.parent))
    _spec.loader.exec_module(trace)
worker = sys.modules["engine_trace_worker"]


def row(**extra):
    return {
        "sink": "os.system(cmd)",
        "source": "request.args['cmd']",
        "sinkFile": "app.py",
        "sinkMethod": "run",
        "sinkLine": 4,
        "sourceKind": "framework",
        "length": 3,
        "paths": 2,
        "sanitized": False,
        "bounded": False,
        "bound": "",
        "trace": ["app.py:run:2 [CALL] request.args['cmd']", "app.py:run:3 [IDENTIFIER] cmd", "app.py:run:4 [CALL] os.system(cmd)"],
        **extra,
    }


def unit(id="a", pin="pin", **extra):
    return {
        "unit": id,
        "repo": "owner/repo",
        "pin": pin,
        "family": "injection",
        "source": "vibe-py",
        "excerpt": True,
        "file": "app.py",
        "function": "run",
        "language": "python",
        "pin_role": "vulnerable",
        "pair": "pair",
        "label": 1,
        "fold": "outer-0",
        "group": "owner/repo",
        "supported": True,
        "reason": "",
        **extra,
    }


def test_trace_roles_and_path_level_sanitizer():
    result = parse_trace(row(sanitized=True, bound="cmd", bounded=True))
    assert [s["roles"] for s in result["steps"]] == [["source"], ["step", "guard"], ["sink"]]
    assert result["steps"][0] == {
        "file": "app.py",
        "method": "run",
        "line": 2,
        "node_label": "CALL",
        "code": "request.args['cmd']",
        "roles": ["source"],
        "raw": "app.py:run:2 [CALL] request.args['cmd']",
    }
    assert result["sanitizer"] == {"on_path": True, "step": None}
    assert result["guard"] == {"text": "cmd", "bounded": True}
    off_path = parse_trace(row(bound="len < 20", bounded=True))
    assert all("guard" not in s["roles"] for s in off_path["steps"])
    assert off_path["guard"]["text"] == "len < 20"


def test_trace_unknown_location_and_method_colons():
    parsed = parse_trace(row(trace=["?:-1 [UNKNOWN] x", "a.php:A::run:4 [CALL] sink(x)"]))
    assert parsed["steps"][0]["line"] == -1
    assert parsed["steps"][0]["method"] == "?"
    assert parsed["steps"][1]["method"] == "A::run"
    assert parse_trace(row(trace=["a.py:run:2 [CALL] x"]))["steps"][0]["roles"] == ["source", "sink"]
    with pytest.raises(ValueError, match="invalid trace"):
        parse_trace(row(trace=["unreadable"]))


def test_rendered_default_scan_params_do_not_request_trace(tmp_path):
    spec = taint_specs(language="python")["injection"]
    rendered = []

    def runner(command, **kwargs):
        request_file = next(arg.split("=", 1)[1] for arg in command if arg.startswith("requestsFile="))
        rendered.append(json.loads(Path(request_file).read_text()))
        return subprocess.CompletedProcess(command, 0, BEGIN + '\n{"r": []}\n' + END, "")

    backend = JoernBackend(runner=runner, session_transport=False)
    cpg = tmp_path / "cpg.bin"
    cpg.write_bytes(b"graph")
    backend._batch_once(cpg, "taint", {"r": request_params(spec)})
    backend._batch_once(cpg, "taint", {"r": request_params(spec, trace=True)})
    assert "trace" not in rendered[0]["r"]
    assert rendered[1]["r"]["trace"] == "true"
    script = backend.queries_dir / "taint.sc"
    assert "traceS = trace" in script.read_text()  # single-query forwarding too


def test_plan_groups_skips_done_and_retains_other_families(tmp_path):
    units = [unit(), unit("b", label=0), unit("c", "other")]
    plan = trace.make_plan(units, tmp_path)
    assert sorted(len(p["units"]) for p in plan) == [1, 2]
    done = next(p for p in plan if p["pin"] == "pin")
    trace.write_json(tmp_path / trace.pin_name(done), {**done, "done": True})
    assert len(trace.make_plan(units, tmp_path)) == 1
    assert trace.make_plan(units, tmp_path, only_family="path") == []
    assert len(trace.make_plan(units, tmp_path, limit=1)) == 1
    # An interrupted pin is rerun, even if a checkpoint file exists.
    trace.write_json(tmp_path / trace.pin_name(done), {**done, "done": False})
    assert len(trace.make_plan(units, tmp_path)) == 2


def test_join_requires_exact_examples_and_ts_opt_in():
    u = unit()
    example = {**u, "id": "a", "candidate": "app.ts::run", "source": "pairs", "language": "typescript"}
    with pytest.raises(ValueError, match="missing 1/1"):
        trace.join_units([u], [])
    assert not trace.join_units([u], [example])[0]["supported"]
    assert trace.join_units([u], [example], include_typescript=True)[0]["supported"]


class FakeBackend:
    seconds = 20
    census_files = 1
    timed_out = False
    fatal = ""
    last_failure = ""
    seen = []

    def __init__(self, deadline, question_deadline, checkpoint):
        self.stages = [{"command": "joern-parse", "seconds": self.seconds, "exit": 0}]
        self.asked = []
        self.answers = {}

    def build(self, root, language):
        graph = root / "graph.bin"
        graph.write_bytes(b"CPG bytes")

        def batch(query, requests):
            self.seen.append(requests)
            self.asked.extend(requests)
            self.stages.append({"command": "joern", "seconds": self.seconds, "exit": 0})
            return {
                **{rid: [row()] if rid == "a" else [] for rid in requests},
                "__census__": [{"files": str(self.census_files), "file_names": ["app.py"]}],
            }

        return CpgResult(graph, lambda q, p: None, run_batch_once=batch)


def run_fake(tmp_path, monkeypatch, backend=FakeBackend, content="def run():\n    pass\n"):
    monkeypatch.setattr(worker, "MeasuredBackend", backend)
    (tmp_path / "app.py").write_text(content)
    pin = {"repo": "owner/repo", "pin": "pin", "units": [unit(), unit("b")]}
    out = tmp_path / "result.json"
    pin["questions"] = trace.prepare_questions(pin, tmp_path, 120) if content else {}
    worker.worker(pin, tmp_path, out, 900, 120)
    return json.loads(out.read_text())


def test_worker_one_batch_and_preserves_witnesses(tmp_path, monkeypatch):
    FakeBackend.seen = []
    result = run_fake(tmp_path, monkeypatch)
    assert [u["status"] for u in result["units"]] == ["path", "asked-nothing"]
    assert len(FakeBackend.seen) == 1
    assert all(p["trace"] == "true" and p["function"] == "run" and p["file"] == "app.py" for p in FakeBackend.seen[0].values())
    assert result["units"][0]["witness_rows"] == [row()]
    assert result["units"][0]["traces"][0]["steps"][-1]["roles"] == ["sink"]
    assert result["units"][1]["questions_completed"] == ["b"]
    assert result["instrument"]["files"][0]["bytes"] > 0


@pytest.mark.parametrize("failure", ["fast", "zero-input", "zero-census"])
def test_invalid_instrument_is_persisted_failed(tmp_path, monkeypatch, failure):
    class BadBackend(FakeBackend):
        seconds = 0.5 if failure == "fast" else 20
        census_files = 0 if failure == "zero-census" else 1

    result = run_fake(tmp_path, monkeypatch, BadBackend, "" if failure == "zero-input" else "def run():\n    pass\n")
    assert result["done"]
    assert all(u["status"] == "failed" for u in result["units"])
    assert all(u["reason"] for u in result["units"])
    assert "<5s" in result["units"][0]["reason"] if failure == "fast" else "zero" in result["units"][0]["reason"]


def test_missing_labelled_function_is_failure(tmp_path, monkeypatch):
    with pytest.raises(ValueError, match="not declared"):
        run_fake(tmp_path, monkeypatch, content="def other():\n    pass\n")


def test_docker_command_is_serial_offline_frozen():
    cmd = trace.docker_command(Path("/export"), Path("/pin"), Path("/output"), "one", "openultrasast:dev", 900, 60)
    assert cmd[cmd.index("--memory") + 1] == "3g"
    assert cmd[cmd.index("--network") + 1] == "none"
    assert cmd[cmd.index("--pull") + 1] == "never"
    assert "/export:/frozen:ro" in cmd and "/pin:/case:ro" in cmd
    assert cmd[-1] == "60"
    assert "/frozen/benchmarks/learn/engine_trace_worker.py" in cmd
    assert "benchmarks/learn/engine_trace_worker.py" in trace.SNAPSHOT_FILES


def test_summary_separates_unanswered_pairs(tmp_path):
    rows = [
        trace.unit_record(unit(), "path"),
        trace.unit_record(unit("b", label=0), "asked-nothing"),
        trace.unit_record(unit("c", pair="other"), "failed"),
        trace.unit_record(unit("d", pair="other", label=0), "timeout"),
    ]
    trace.write_json(tmp_path / "one.json", {"done": True, "units": rows})
    result = trace.summary(tmp_path)
    assert result["pair_asymmetry"] == {"vulnerable-only": 1, "neither": 1}
    assert result["pair_asymmetry_both_answered"] == {"vulnerable-only": 1}
    assert result["by_family"]["injection"]["failed"] == 1


def test_frozen_source_overlays_uncommitted_files_without_live_mount(tmp_path, monkeypatch):
    root = tmp_path / "repo"
    root.mkdir()
    for name in trace.SNAPSHOT_FILES:
        dest = root / name
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text("snapshot " + name)
    monkeypatch.setattr(trace, "ROOT", root)
    commands = []

    def run(cmd, **kwargs):
        commands.append(cmd)
        return subprocess.CompletedProcess(cmd, 0, "head-sha\n", "")

    monkeypatch.setattr(trace.subprocess, "run", run)
    frozen = tmp_path / "frozen"
    proof = trace.freeze_source(frozen)
    assert commands[0][-3:] == ["archive", "HEAD", "src"]
    for name in trace.SNAPSHOT_FILES:
        assert (frozen / name).read_bytes() == (root / name).read_bytes()
        assert len(proof["overlay_sha256"][name]) == 64


def test_question_watchdog_retains_completed_rows_without_joern(tmp_path):
    import sys

    code = (
        "import json,time\n"
        'print(json.dumps({"__question__":"a"}),flush=True)\n'
        'print(json.dumps({"id":"a","rows":[]}),flush=True)\n'
        'print(json.dumps({"__question__":"b"}),flush=True)\n'
        "time.sleep(60)\n"
    )
    backend = worker.MeasuredBackend(10, 0.05, lambda: None)
    done = backend._run([sys.executable, "-c", code], timeout=10, cwd=tmp_path)
    assert backend.timed_out
    assert backend.asked == ["a", "b"]
    assert backend.answers == {"a": []}
    assert done.returncode != 0
    assert backend.stages[0]["seconds"] < 5
    assert "implausibly fast" in backend.fatal


def test_pair_mapping_handles_harvest_normalized_crlf(tmp_path):
    import hashlib
    from types import SimpleNamespace

    file = tmp_path / "excerpt.py"
    file.write_bytes(b"def run():\r\n    pass\r\n")
    normalized = file.read_text().encode()
    pin = hashlib.sha1(b"blob " + str(len(normalized)).encode() + b"\0" + normalized).hexdigest()
    case = SimpleNamespace(name="case", repo="owner/repo", relpath="app.py", vuln_file=file, fixed_file=file)
    inputs = object.__new__(trace.Inputs)
    inputs.pairs = [case]
    inputs.hashes = {}
    assert inputs.pair_for(unit(pin=pin)) == (case, "vuln")


def test_missing_batch_answers_are_timeout_not_empty(tmp_path, monkeypatch):
    class TimeoutBackend(FakeBackend):
        def build(self, root, language):
            graph = root / "graph.bin"
            graph.write_bytes(b"CPG bytes")

            def batch(query, requests):
                self.timed_out = True
                self.asked.append("a")
                return None

            return CpgResult(graph, lambda q, p: None, run_batch_once=batch)

    result = run_fake(tmp_path, monkeypatch, TimeoutBackend)
    assert [r["status"] for r in result["units"]] == ["timeout", "timeout"]
    assert result["units"][0]["questions_asked"] == ["a"]
    assert result["units"][1]["questions_asked"] == []


def test_dry_run_makes_no_subprocess_calls(tmp_path, monkeypatch, capsys):
    import sys

    u = unit()
    examples = [{**u, "id": "a", "candidate": "app.py::run", "source": "population-v1"}]
    units = tmp_path / "units.jsonl"
    units.write_text(json.dumps(u) + "\n")
    example_file = tmp_path / "examples.jsonl"
    example_file.write_text(json.dumps(examples[0]) + "\n")
    monkeypatch.setattr(
        sys, "argv", [str(SCRIPT), "--dry-run", "--units", str(units), "--examples", str(example_file), "--out", str(tmp_path / "results")]
    )
    monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("dry-run spawned a subprocess"))
    assert trace.main() == 0
    plan = json.loads(capsys.readouterr().out)
    assert plan["runnable_pins"] == 1 and plan["units"] == 1
    assert plan["estimated_seconds"] == 182
    assert not (tmp_path / "results").exists()


def test_family_filtered_results_merge_without_losing_prior_units(tmp_path):
    first = trace.make_plan([unit(), unit("b", family="path")], tmp_path, only_family="injection")[0]
    trace.save_pin(tmp_path / trace.pin_name(first), {**first, "done": True})
    second = trace.make_plan([unit(), unit("b", family="path")], tmp_path, only_family="path")[0]
    trace.save_pin(tmp_path / trace.pin_name(second), {**second, "done": True})
    assert trace.make_plan([unit(), unit("b", family="path")], tmp_path) == []


def test_failed_launch_persists_failed_pin_and_resume_skips_it(tmp_path, monkeypatch):
    from types import SimpleNamespace

    class Inputs:
        def materialize(self, pin, target):
            target.mkdir()
            (target / "app.py").write_text("def run():\n    pass\n")

    commands = []

    def run(cmd, **kwargs):
        commands.append(cmd)
        return subprocess.CompletedProcess(cmd, 0, "", "")

    monkeypatch.setattr(trace.subprocess, "run", run)
    pin = trace.make_plan([unit()], tmp_path)[0]
    args = SimpleNamespace(image="openultrasast:dev", deadline=900, question_deadline=120)
    assert trace.run_pin(pin, tmp_path / "frozen", {}, Inputs(), tmp_path, args)
    result = json.loads((tmp_path / trace.pin_name(pin)).read_text())
    assert result["units"][0]["status"] == "failed"
    assert "produced no result" in result["units"][0]["reason"]
    assert [c[1] for c in commands if c[0] == "docker"] == ["run", "rm"]
    assert trace.make_plan([unit()], tmp_path) == []


def import_container_worker():
    """Import in a fresh interpreter so host imports cannot mask the boundary."""
    script = SCRIPT.with_name("engine_trace_worker.py")
    code = """
import json, pathlib, runpy, sys
sys.modules['yaml'] = None
sys.modules['boto3'] = None
script = pathlib.Path(sys.argv[1])
size = len(script.read_bytes())
runpy.run_path(str(script), run_name='trace_import_check')
print(json.dumps({'bytes': size, 'modules': sorted(
    name for name in sys.modules if name == 'openultrasast' or name.startswith('openultrasast.')
)}))
"""
    result = subprocess.run([sys.executable, "-c", code, str(script)], capture_output=True, text=True, timeout=30)
    assert result.returncode == 0, result.stderr
    return json.loads(result.stdout)


def test_trace_container_import_without_yaml_or_boto3():
    assert import_container_worker()["bytes"] > 0


def test_trace_container_transitive_import_boundary():
    # Explicit inventory: additions need review, particularly learn.examples/plane.*.
    assert import_container_worker()["modules"] == [
        "openultrasast",
        "openultrasast.contracts",
        "openultrasast.cpg",
        "openultrasast.cpg.artifact",
        "openultrasast.cpg.backend",
        "openultrasast.cpg.session",
        "openultrasast.model",
        "openultrasast.model.contracts",
        "openultrasast.model.trace",
    ]


@pytest.mark.parametrize("existing", [False, True])
def test_trace_dry_run_unsupported_units_write_nothing(tmp_path, monkeypatch, capsys, existing):
    out = tmp_path / "results"
    if existing:
        out.mkdir()
        (out / ".lock").write_bytes(b"existing lock")
        (out / "checkpoint.json").write_text('{"done": false}')
        (out / "previous.json").write_text('{"done": true, "units": []}')
    before = {p.name: (p.read_bytes(), p.stat().st_mtime_ns) for p in out.glob("*")}
    units = [unit("ts", file="app.ts"), unit("py", pin="python")]
    examples = [{**u, "id": u["unit"], "candidate": u["file"] + "::run", "source": "population-v1"} for u in units]
    units_file, examples_file = tmp_path / "units.jsonl", tmp_path / "examples.jsonl"
    units_file.write_text("\n".join(map(json.dumps, units)))
    examples_file.write_text("\n".join(map(json.dumps, examples)))
    monkeypatch.setattr(
        sys, "argv", [str(SCRIPT), "--dry-run", "--units", str(units_file), "--examples", str(examples_file), "--out", str(out)]
    )
    monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("dry run spawned subprocess"))
    monkeypatch.setattr(trace, "write_json", lambda *a: pytest.fail("dry run wrote JSON"))
    assert trace.main() == 0
    plan = json.loads(capsys.readouterr().out)
    assert plan["pins"] == 2 and plan["runnable_pins"] == 1
    assert out.exists() == existing
    assert {p.name: (p.read_bytes(), p.stat().st_mtime_ns) for p in out.glob("*")} == before


def test_trace_real_unsupported_run_then_typescript_opt_in_replans(tmp_path, monkeypatch, capsys):
    u = unit(file="app.ts")
    example = {**u, "id": "a", "candidate": "app.ts::run", "source": "population-v1"}
    units_file, examples_file, out = tmp_path / "units.jsonl", tmp_path / "examples.jsonl", tmp_path / "results"
    units_file.write_text(json.dumps(u) + "\n")
    examples_file.write_text(json.dumps(example) + "\n")
    argv = [str(SCRIPT), "--units", str(units_file), "--examples", str(examples_file), "--out", str(out)]
    monkeypatch.setattr(sys, "argv", argv)
    monkeypatch.setattr(trace, "freeze_source", lambda target: {})
    monkeypatch.setattr(trace.subprocess, "run", lambda *a, **kw: pytest.fail("unsupported run invoked Docker"))
    assert trace.main() == 0
    result_file = out / trace.pin_name(u)
    first = json.loads(result_file.read_text())
    assert first["done"] and first["units"][0]["status"] == "unsupported"
    assert trace.make_plan(trace.join_units([u], [example]), out) == []

    called = []

    def run_pin(pin, source, provenance, inputs, results, args):
        called.append(pin)
        trace.save_pin(
            results / trace.pin_name(pin), {**pin, "done": True, "units": [trace.unit_record(r, "asked-nothing") for r in pin["units"]]}
        )
        return False

    def inspect_image(command, **kwargs):
        assert command[:3] == ["docker", "image", "inspect"]
        return subprocess.CompletedProcess(command, 0, "sha256:fake\n", "")

    monkeypatch.setattr(trace, "run_pin", run_pin)
    monkeypatch.setattr(trace.subprocess, "run", inspect_image)
    monkeypatch.setattr(sys, "argv", [*argv, "--include-typescript"])
    assert trace.main() == 0
    assert len(called) == 1 and called[0]["units"][0]["supported"]
    assert called[0]["selection"] != first["selection"]
    assert json.loads(result_file.read_text())["units"][0]["status"] == "asked-nothing"
    assert trace.make_plan(trace.join_units([u], [example], include_typescript=True), out, include_typescript=True) == []


def test_trace_resume_identity_includes_selection_even_if_support_unchanged(tmp_path):
    first = trace.make_plan([unit()], tmp_path)[0]
    trace.save_pin(tmp_path / trace.pin_name(first), {**first, "done": True})
    assert trace.make_plan([unit()], tmp_path) == []
    assert len(trace.make_plan([unit()], tmp_path, include_typescript=True)) == 1