#!/usr/bin/env python3
"""llm_prescan.py — deterministic LLM-application integration pre-scan.

ARION service artifact. Scans application source for the signatures of
unsafe LLM integration — model output flowing into exec/SQL/HTTP/HTML
sinks without review, request input concatenated into prompts and
system messages, secrets interpolated into prompt context, message
roles taken from the request, agent loops with no iteration budget,
dynamic dispatch on model-chosen names, full prompt/transcript
logging, and files that call a model with zero output-validation
vocabulary at all. Line-wise taint (name assignment only), fully
offline, no inference — ironic and deliberate.

Findings (base severity):
  LLM-OUTPUT-EXEC       HIGH    a value derived from an LLM call
                                reaches eval/exec/os.system/subprocess
                                execSync/Runtime.exec/ProcessBuilder/
                                new Function/vm.runIn*/pickle.loads/
                                __import__/import_module — the model's
                                output is executable code, and the model
                                is steered by whoever controls the input
  LLM-DISPATCH-DYNAMIC  HIGH    getattr(x, tainted) or obj[tainted]( —
                                the model picks the attribute/method to
                                invoke; any callable in the namespace is
                                reachable
  LLM-OUTPUT-SQL        HIGH    model-derived text into .execute/
                                .executemany/.raw/.query/$queryRaw —
                                generated SQL runs as-is
  LLM-OUTPUT-FETCH      MEDIUM  model-derived URL/host into requests.*,
                                fetch, axios, urlopen, http.Get — the
                                model chooses the target (SSRF channel)
  LLM-OUTPUT-HTML       MEDIUM  model-derived text into
                                dangerouslySetInnerHTML/innerHTML/
                                insertAdjacentHTML/marked/markdown —
                                generated markup renders unsanitised;
                                LOW when the file carries a sanitizer
                                (DOMPurify/bleach/sanitize-html)
  PROMPT-USER-CONCAT    MEDIUM  request input interpolated (or assigned
                                wholesale) into a prompt/system/content
                                build — the caller steers the model;
                                prompt structure, secrets and downstream
                                tools inherit the injection
  SECRET-IN-CONTEXT     HIGH    env/vault/key/secret/token/password
                                value interpolated into a prompt/
                                messages/content string — credentials
                                leave for the provider and every log
                                the prompt lands in
  ROLE-FROM-REQ         MEDIUM  a message role taken from request input
                                — the caller claims system/developer
                                authority
  UNBOUNDED-AGENT-LOOP  MEDIUM  while True/while(true)/for(;;)/go for{
                                /rust loop{ containing an LLM call with
                                no max_iterations/steps/turns/budget
                                vocabulary — a runaway loop burns
                                tokens and tool calls until killed
  LLM-LOG-CONTENT       LOW     a logger/console line emitting the
                                prompt/messages/transcript variable —
                                user PII and the system prompt land in
                                logs
  LLM-NO-VALIDATION     LOW     file-level: the file calls a model and
                                contains zero output-validation vocab
                                (zod/pydantic/jsonschema/ajv/joi/parse/
                                validate/schema/response_format/
                                OutputParser) — output is trusted
                                blind; LOW, informational

Context rules:
  * comment lines and block-comment bodies are stripped first
  * import/require lines never flag
  * test/spec/fixture paths downgrade one tier
  * taint is name-assignment only: `x = <llm call>` and names built
    from tainted names (.content/.text/attr/index/f-string/concat);
    no cross-function or cross-file flow
  * all findings require LLM vocabulary somewhere in the file — a
    sink next to a coincidence-named variable does not flag
  * files carrying a sanitizer (DOMPurify/bleach/sanitize-html/xss)
    downgrade LLM-OUTPUT-HTML to LOW

Scope: deterministic pre-scan for human review — flags the line shapes
where model I/O crosses a trust boundary; it cannot see system-level
prompt assembly, provider-side guardrails, or middleware. Clean output
!= a safe integration.

  llm_prescan.py <file-or-dir> [--json]
  llm_prescan.py --selftest
Stdlib only. Exit 0 = no HIGH/MEDIUM findings, 1 = findings, 2 = usage.
"""
import json
import os
import re
import sys

SEV_ORDER = {"HIGH": 0, "MEDIUM": 1, "LOW": 2, "INFO": 3}

TARGET_EXT = (
    ".py", ".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs",
    ".php", ".rb", ".go", ".java", ".kt", ".cs", ".groovy",
    ".rs", ".scala", ".dart", ".vue", ".svelte",
)

SKIP_DIRS = {
    "node_modules", ".git", "vendor", "dist", "build", "target",
    "__pycache__", ".venv", "venv", "coverage", ".next", "out",
}

IMPORT_LINE = re.compile(
    r"^\s*(?:from\s+[\w.]+\s+import\b|import\s+[\w.{*]|"
    r"(?:const|let|var)\s+[\w${},\s]*=\s*require\s*\(|"
    r"\buse\s+[\w:]+|require\s*\(['\"][\w./-]+['\"]\s*\)?\s*;?\s*$|"
    r"#include|package\s+[\w.]+;?\s*$|using\s+[\w.]+;)")

TEST_PATH = re.compile(
    r"(^|/)(test|tests|spec|specs|fixtures?|mocks?|__tests__|__mocks__|"
    r"e2e|cypress|playwright)(/|$)|(_test|_spec|\.test|\.spec)\.",
    re.I)

# --- vocabulary -----------------------------------------------------------

# strong markers — the file actually wires an SDK: import lines,
# constructor calls, module-attribute use, or model-call shapes.
# Loose keywords (a prose 'LLM', provider names inside a regex that
# scans FOR api keys) never gate a file open. Hostnames excluded —
# 'api.openai.com' is not 'openai.' module use.
LLM_STRONG = re.compile(
    r"chat\.completions\.create|\.responses\.create|"
    r"messages\.create\s*\(|\.generate_content\s*\(|"
    r"generateContent\s*\(|litellm\.completion|"
    r"completions\.create|ChatCompletion\.create|"
    r"chain\.(?:run|invoke|ainvoke)\s*\(|agent\.(?:run|arun)\s*\(|"
    r"ollama\.(?:chat|generate)\s*\(|\.converse\s*\(|"
    r"(?:\bchat|\.chat)\s*\(\s*(?:messages|history|prompt)|"
    r"client\.(?:chat|complete)\s*\(|call_model\s*\(|"
    r"import\s+(?:openai|anthropic|langchain|litellm|ollama|"
    r"google\.generativeai|cohere|mistralai|groq|openrouter|"
    r"replicate|vertexai|llama_index|llamaindex|instructor|"
    r"transformers|huggingface_hub|guidance|dspy|haystack|"
    r"semantic_kernel|autogen|crewai)|"
    r"from\s+['\"]?(?:openai|@anthropic-ai|@langchain|langchain|"
    r"litellm|ollama|@google/generative-ai|@google/genai|"
    r"cohere-ai|@mistralai|groq-sdk|openrouter|replicate|"
    r"@instructor|@huggingface|@aws-sdk/client-bedrock|"
    r"@langchain/[\w-]+)|"
    r"require\(\s*['\"](?:openai|@anthropic-ai|langchain|litellm|"
    r"ollama|cohere-ai|groq-sdk|replicate)|"
    r"\b(?:OpenAI|Anthropic|AzureOpenAI|ChatOpenAI|ChatAnthropic|"
    r"ChatOllama|BedrockChat|ChatBedrock|CohereClient|"
    r"MistralClient|Groq|OpenRouter|PromptTemplate|"
    r"ChatPromptTemplate|HumanMessage|SystemMessage|AIMessage|"
    r"LLMChain|ChatGoogleGenerativeAI|ChatVertexAI|"
    r"ChatMistralAI|ChatGroq|ChatCohere)\s*\(|"
    r"(?<![\w./])(?:openai|anthropic|litellm|ollama|genai|cohere|"
    r"mistralai|groq|openrouter|langchain|llama_index|llamaindex|"
    r"instructor|transformers|semantic_kernel|autogen|crewai|"
    r"dspy)\.", re.I)

# call shapes whose result is model output
LLM_CALL = re.compile(
    r"chat\.completions\.create|\.responses\.create|"
    r"messages\.create\s*\(|\.generate_content\s*\(|"
    r"generateContent\s*\(|litellm\.completion|\.completion\s*\(|"
    r"completions\.create|ChatCompletion\.create|"
    r"\.(?:a?invoke|predict|apredict)\s*\(|chain\.run\s*\(|"
    r"agent\.(?:run|arun)\s*\(|\bllm[\.\(]|"
    r"ollama\.(?:chat|generate)\s*\(|\.converse\s*\(|"
    r"model\.generate|client\.(?:chat|complete)\s*\(|"
    r"get_chat_completion|completion\s*\(\s*model|"
    r"\bmodel\s*\(\s*messages|call_model\s*\(|"
    r"(?:\bchat|\.chat)\s*\(\s*(?:messages|history|prompt)", re.I)

# exec / code-eval sinks
EXEC_SINK = re.compile(
    r"\beval\s*\(|\bexec\s*\(|execSync\s*\(|execFileSync\s*\(|"
    r"os\.system\s*\(|os\.popen\s*\(|"
    r"subprocess\.(?:call|run|Popen|check_output|check_call)\s*\(|"
    r"child_process|spawnSync\s*\(|"
    r"Runtime\.getRuntime\(\)\.exec|ProcessBuilder\s*\(|"
    r"new\s+Function\s*\(|vm\.runIn\w*\s*\(|"
    r"pickle\.loads\s*\(|marshal\.loads\s*\(|"
    r"__import__\s*\(|importlib\.import_module\s*\(|"
    r"shell_exec\s*\(|passthru\s*\(|pcntl_exec\s*\(|"
    r"assert\s*\(\s*['\"]`?\s*\w*\s*\$\{", )   # last alt ~never matches

# sql sinks
SQL_SINK = re.compile(
    r"\.execute(?:many|script)?\s*\(|\.raw\s*\(|\.query\s*\(|"
    r"\$queryRaw|\$executeRaw|executeQuery\s*\(|"
    r"sequelize\.query|knex\.raw|db\.run\s*\(|conn\.exec\s*\(")

# outbound-http sinks (model-chosen target = SSRF channel)
FETCH_SINK = re.compile(
    r"requests\.(?:get|post|put|head|request)\s*\(|"
    r"urllib\.request\.urlopen|urlopen\s*\(|\bfetch\s*\(|"
    r"axios\.(?:get|post|put|request)\s*\(|got\.(?:get|post)\s*\(|"
    r"http\.Get\s*\(|http\.NewRequest|http\.Post\s*\(|"
    r"needle\.(?:get|post)|superagent\.|wp_remote_get|"
    r"file_get_contents\s*\(|curl_init|RestTemplate|"
    r"\.getForObject|WebClient")

# rendered-HTML sinks
HTML_SINK = re.compile(
    r"dangerouslySetInnerHTML|\.innerHTML\b|\.outerHTML\b|"
    r"insertAdjacentHTML|\bmarked\s*\(|marked\.parse|"
    r"markdown\.markdown\s*\(|\.render\s*\(|v-html|\{\{\{")

# sanitizer vocabulary — downgrades OUTPUT-HTML
SANITIZER = re.compile(
    r"DOMPurify|\bbleach\b|sanitize[-_]?html|sanitize\s*\(|\bxss\b|"
    r"escapeHtml|html\.escape", re.I)

# dynamic dispatch on a model-chosen name
GETATTR_RX = re.compile(r"getattr\s*\(")
IDX_CALL_RX = re.compile(r"\[\s*([A-Za-z_]\w{0,39})\s*\]\s*\(")

# request-input sources (web request objects, stdin)
REQ_SRC = re.compile(
    r"\breq\.(?:query|body|param|params|args|input|form)|"
    r"request\.(?:args|GET|POST|form|data|values|json|files|"
    r"query_params|content)|\$_GET|\$_POST|\$_REQUEST|"
    r"params\s*\[\s*:|params\s*\[\s*['\"]|ctx\.(?:query|request)|"
    r"context\.params|c\.(?:QueryParam|Param|BodyParser)|"
    r"r\.URL\.Query|r\.FormValue|stdin|searchParams\.get|"
    r"event\.(?:body|queryStringParameters|input)|"
    r"req\.headers|input\s*\(", re.I)

# interpolation markers — the request value lands inside a string
INTERP = re.compile(
    r"f['\"]|f'''|f\"\"\"|\$\{|\+\s*(?:req\b|request\b|user_|input)|"
    r"\.format\s*\(|\.format_map|%\s*\(|%\s+\w|\.concat\s*\(|"
    r"sprintf\s*\(|String\.format")

# authority-position targets — request input landing in THESE is a
# boundary crossing; 'content' alone (the API's own field name) is the
# normal chat shape and does not flag
AUTHORITY_TARGET = re.compile(
    r"\b(system_?prompt|sys_?prompt|system|prompt|prompts|"
    r"instruction|instructions|messages|message_list|user_?prompt|"
    r"task_?prompt|query_?prompt|chat_history|conversation)\b", re.I)

# broader set for secret-in-context — any string that plausibly ends
# up inside a message, including bare 'content' builds
CONTEXT_TARGET = re.compile(
    r"\b(system_?prompt|sys_?prompt|system|prompt|prompts|"
    r"instruction|instructions|messages|message_list|content|"
    r"context|template|chat_history|conversation|user_?prompt)\b",
    re.I)

# secret-ish names and env/vault references
SECRET_NAME = re.compile(
    r"api_?key|secret|token|password|passwd|\bpwd\b|private_?key|"
    r"credential|auth_?token|client_?secret|access_?key|signing", re.I)
ENV_REF = re.compile(
    r"os\.environ|process\.env|getenv|ENV\s*\[|ENV\.fetch|"
    r"System\.getenv|os\.getenv|settings\.|config\.|vault|kms|"
    r"SecretManager|secretsmanager", re.I)

# message role taken from the request
ROLE_REQ = re.compile(
    r"['\"]?role['\"]?\s*[:=]\s*[^,\n}\)]*?"
    r"(?:req\.|request\.|\$_POST|\$_GET|\$_REQUEST|params\s*\[|"
    r"ctx\.|c\.QueryParam|body\.|event\.|input\s*\()", re.I)

# unbounded loop headers
LOOP_OPEN = re.compile(
    r"\bwhile\s+True\b|\bwhile\s*\(\s*true\s*\)|\bwhile\s+1\b|"
    r"for\s*\(\s*;\s*;|^\s*for\s*\{|^\s*loop\s*\{|\bloop\s+do\b")
CAP_VOCAB = re.compile(
    r"max_iter|max_steps|max_turns|max_loops|max_rounds|step_limit|"
    r"iteration_limit|max_retries|max_depth|max_calls|"
    r"MAX_ITER|MAX_STEP|MAX_TURN|budget|deadline|iteration[s]?\s*>=|"
    r"step\s*>=|turns?\s*>=|count\s*>=", re.I)

# logging sinks that would carry full content
LOG_SINK = re.compile(
    r"logging\.(?:info|debug|warn|error|warning)|"
    r"logger\.(?:info|debug|warn|error|log)|"
    r"log\.(?:info|debug|warn|error|trace)|"
    r"console\.(?:log|debug|info)|fmt\.Print|println|"
    r"System\.out\.print|print_r\s*\(|var_dump")
CONTENT_VAR = re.compile(
    r"\b(messages|message_list|system_prompt|full_prompt|prompt|"
    r"chat_history|conversation|transcript|ctx_messages)\b", re.I)

# output-validation vocabulary — file-level gate for NO-VALIDATION
VALID_VOCAB = re.compile(
    r"\bzod\b|pydantic|jsonschema|json_schema|safeParse|safe_parse|"
    r"\bajv\b|\bjoi\b|\byup\b|marshmallow|guardrail|"
    r"\.model_validate|response_format|OutputParser|"
    r"output_parser|\.parse\s*\(|validate\s*\(|validator|"
    r"schema\b|json\.loads|JSON\.parse|TypeAdapter|struct\.decode|"
    r"serde|jsonschema\.validate", re.I)

PER_CODE_CAP = 40
LOOP_WINDOW = 60       # lines after a loop header to inspect


class Finding:
    def __init__(self, sev, code, where, line, msg):
        self.sev, self.code, self.where = sev, code, where
        self.line, self.msg = line, msg

    def as_dict(self):
        return {"severity": self.sev, "code": self.code,
                "file": self.where, "line": self.line, "msg": self.msg}


# --- helpers --------------------------------------------------------------
def strip_comments(lines):
    """Remove full-line comments and /* */ block bodies. Keeps count."""
    out, in_block = [], False
    for ln in lines:
        s = ln
        if in_block:
            if "*/" in s:
                s = s.split("*/", 1)[1]
                in_block = False
            else:
                out.append("")
                continue
        while "/*" in s:
            if "*/" in s.split("/*", 1)[1]:
                s = s.replace(s[s.index("/*"):s.index("*/") + 2], " ", 1)
            else:
                s = s.split("/*", 1)[0]
                in_block = True
        st = s.strip()
        if st.startswith(("//", "#", "*", "<!--")):
            out.append("")
            continue
        out.append(s)
    return out


def downgrade(sev):
    return "MEDIUM" if sev == "HIGH" else "LOW"


ASSIGN_RX = re.compile(
    r"^\s*(?:const|let|var|final|static|auto|val|String|std::string|"
    r"self\.)?\s*\.?\s*([A-Za-z_]\w{0,39})\s*"
    r"(?::\s*[\w<>\[\],. |]+\s*)?=\s*(.+?);?\s*$")
SELF_ASSIGN = re.compile(r"^\s*self\.([A-Za-z_]\w{0,39})\s*=\s*(.+)$")

UNWRAP_ATTR = re.compile(
    r"\.(?:content|text|message|choices|output|result|answer|"
    r"response|candidates|data|completion|generated_text|"
    r"tool_calls|function_call)\b")


def rhs_taints(rhs, tainted):
    """True if the RHS references any tainted name (attr/index/format)."""
    for name in tainted:
        if re.search(r"\b%s\b" % re.escape(name), rhs):
            return True
    return False


def names_in_line(ln, tainted):
    """Yield tainted names referenced on a line."""
    for name in tainted:
        if re.search(r"\b%s\b" % re.escape(name), ln):
            yield name


def scan_file(path, rel):
    try:
        raw = open(path, encoding="utf-8", errors="replace").read()
    except OSError:
        return []
    is_test = bool(TEST_PATH.search(rel))
    lines = strip_comments(raw.splitlines())
    text = "\n".join(lines)

    # file-level gates — no real LLM wiring: not an LLM file, all silent
    if not LLM_STRONG.search(text):
        return []

    file_has_sanitizer = bool(SANITIZER.search(text))
    file_has_valid = bool(VALID_VOCAB.search(text))
    findings = []
    counts = {}
    tainted = set()

    def emit(sev, code, ln, msg):
        if is_test:
            sev = downgrade(sev)
        if counts.get(code, 0) >= PER_CODE_CAP:
            return
        counts[code] = counts.get(code, 0) + 1
        findings.append(Finding(sev, code, rel, ln, msg))

    # ---- pass 1: taint propagation + per-line findings ------------------
    for i, ln in enumerate(lines, 1):
        st = ln.strip()
        if not st or IMPORT_LINE.match(ln):
            continue

        # taint assignment: name = <llm call> or name = <tainted...>
        m = ASSIGN_RX.match(ln) or SELF_ASSIGN.match(ln)
        if m:
            name, rhs = m.group(1), m.group(2)
            if LLM_CALL.search(rhs) or rhs_taints(rhs, tainted):
                tainted.add(name)

        line_taints = list(names_in_line(ln, tainted))

        # -- sinks consuming model output --
        if line_taints:
            if EXEC_SINK.search(ln):
                emit("HIGH", "LLM-OUTPUT-EXEC", i,
                     "model output (%s) reaches an exec/eval sink — "
                     "generated code runs as-is"
                     % ", ".join(sorted(line_taints)[:3]))
            if GETATTR_RX.search(ln):
                emit("HIGH", "LLM-DISPATCH-DYNAMIC", i,
                     "getattr() on model output (%s) — the model picks "
                     "the attribute/method to invoke"
                     % ", ".join(sorted(line_taints)[:3]))
            if SQL_SINK.search(ln):
                emit("HIGH", "LLM-OUTPUT-SQL", i,
                     "model output (%s) into a SQL sink — generated "
                     "queries execute verbatim"
                     % ", ".join(sorted(line_taints)[:3]))
            if FETCH_SINK.search(ln):
                emit("MEDIUM", "LLM-OUTPUT-FETCH", i,
                     "model output (%s) into an HTTP fetch — the model "
                     "chooses the target (SSRF channel)"
                     % ", ".join(sorted(line_taints)[:3]))
            if HTML_SINK.search(ln):
                sev = "LOW" if file_has_sanitizer else "MEDIUM"
                emit(sev, "LLM-OUTPUT-HTML", i,
                     "model output (%s) rendered into HTML/markdown — "
                     "generated markup lands in the DOM"
                     % ", ".join(sorted(line_taints)[:3]))

        # dict-dispatch on model output: tools[resp.action](...)
        for mm in IDX_CALL_RX.finditer(ln):
            if mm.group(1) in tainted:
                emit("HIGH", "LLM-DISPATCH-DYNAMIC", i,
                     "[%s](...) dict-dispatch on model output — the "
                     "model chooses which function runs" % mm.group(1))
                break

        # -- request input into prompt construction --
        if REQ_SRC.search(ln) and AUTHORITY_TARGET.search(ln):
            if INTERP.search(ln):
                emit("MEDIUM", "PROMPT-USER-CONCAT", i,
                     "request input interpolated into a prompt/"
                     "instruction build — the caller steers the model")
            elif m and AUTHORITY_TARGET.search(m.group(1)) \
                    and REQ_SRC.search(m.group(2)):
                emit("MEDIUM", "PROMPT-USER-CONCAT", i,
                     "request input assigned wholesale into a prompt/"
                     "messages variable — caller controls the content")

        # -- secrets into prompt context --
        if (CONTEXT_TARGET.search(ln) and INTERP.search(ln)
                and SECRET_NAME.search(ln)
                and not re.match(r"^\s*[\w.]+\s*\(", st)):
            env = ENV_REF.search(ln)
            emit("HIGH", "SECRET-IN-CONTEXT", i,
                 "%s interpolated into a prompt/message string — "
                 "credentials ship to the provider and every log"
                 % ("env/vault value" if env else "secret-named value"))

        # -- message role from the request --
        if ROLE_REQ.search(ln):
            emit("MEDIUM", "ROLE-FROM-REQ", i,
                 "message role taken from request input — the caller "
                 "can claim system/developer authority")

        # -- full-content logging --
        if LOG_SINK.search(ln) and CONTENT_VAR.search(ln):
            emit("LOW", "LLM-LOG-CONTENT", i,
                 "log line emits prompt/messages — user PII and the "
                 "system prompt land in logs")

        # -- unbounded agent loop --
        if LOOP_OPEN.search(ln):
            window = "\n".join(lines[i - 1:i - 1 + LOOP_WINDOW])
            if LLM_CALL.search(window) and not CAP_VOCAB.search(window):
                emit("MEDIUM", "UNBOUNDED-AGENT-LOOP", i,
                     "unbounded loop containing an LLM call with no "
                     "iteration/step/budget cap — a runaway burns "
                     "tokens until killed")

    # ---- file-level: model calls with zero validation vocab -------------
    if not file_has_valid:
        # point at the first line that actually wires the SDK, else the
        # first code line — never an import or docstring
        emit_at = None
        for i, ln in enumerate(lines, 1):
            if not ln.strip() or IMPORT_LINE.match(ln):
                continue
            if LLM_STRONG.search(ln):
                emit_at = i
                break
            if emit_at is None:
                emit_at = i
        emit("LOW", "LLM-NO-VALIDATION", emit_at or 1,
             "file references an LLM SDK and carries zero output-"
             "validation vocabulary (no schema/parse/validate) "
             "— output trusted blind")

    return findings


def find_targets(root):
    if os.path.isfile(root):
        return [root]
    out = []
    for dp, dns, fns in os.walk(root):
        dns[:] = [d for d in dns if d not in SKIP_DIRS]
        for fn in fns:
            if fn.endswith(TARGET_EXT):
                out.append(os.path.join(dp, fn))
    return sorted(out)


def render(findings, targets):
    findings.sort(key=lambda f: (SEV_ORDER[f.sev], f.where, f.line))
    rows = ["# llm_prescan report", "",
            "files scanned: %d — findings: %d" % (len(targets),
                                                  len(findings)), ""]
    cur = None
    for f in findings:
        if f.sev != cur:
            cur = f.sev
            rows.append("## %s" % cur)
        rows.append("- `%s` %s:%d — %s" % (f.code, f.where, f.line,
                                           f.msg))
    if not findings:
        rows.append("no HIGH/MEDIUM/LOW signatures found — this is a "
                    "pre-scan; clean output != a safe integration")
    return "\n".join(rows)


# --- selftest -------------------------------------------------------------
SELFTEST = [
    # --- exec / dispatch / sql / fetch / html sinks ---
    ("a.py",
     "import openai\n"
     "resp = openai.chat.completions.create(model='g', messages=m)\n"
     "code = resp.choices[0].message.content\n"
     "exec(code)\n",
     {"LLM-OUTPUT-EXEC", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "resp = openai.chat.completions.create(model='g', messages=m)\n"
     "eval(resp.choices[0].message.content)\n",
     {"LLM-OUTPUT-EXEC", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "const r = await client.chat.completions.create({model,messages});\n"
     "const out = r.choices[0].message.content;\n"
     "eval(out);\n",
     {"LLM-OUTPUT-EXEC", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "resp = client.chat.completions.create(model='g', messages=m)\n"
     "cmd = resp.choices[0].message.content\n"
     "subprocess.run(cmd, shell=True)\n",
     {"LLM-OUTPUT-EXEC", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import anthropic\n"
     "r = client.messages.create(model='c', messages=m)\n"
     "tool = getattr(worker, r.content[0].text)\n"
     "tool()\n",
     {"LLM-DISPATCH-DYNAMIC", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import Anthropic from '@anthropic-ai/sdk';\n"
     "const r = await claude.messages.create({model,messages});\n"
     "const fn = r.content[0].name;\n"
     "tools[fn](args);\n",
     {"LLM-DISPATCH-DYNAMIC", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "q = llm.predict(text=question)\n"
     "cur.execute(q)\n",
     {"LLM-OUTPUT-SQL", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "sql = chain.invoke({'q': q})\n"
     "rows = db.query(sql)\n",
     {"LLM-OUTPUT-SQL", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai, requests\n"
     "u = llm.predict(text=q)\n"
     "requests.get(u)\n",
     {"LLM-OUTPUT-FETCH", "LLM-NO-VALIDATION"}),
    ("a.tsx",
     "import OpenAI from 'openai';\n"
     "const r = await client.chat.completions.create({model,messages});\n"
     "const html = r.choices[0].message.content;\n"
     "return <div dangerouslySetInnerHTML={{__html: html}} />;\n",
     {"LLM-OUTPUT-HTML", "LLM-NO-VALIDATION"}),
    ("a.tsx",
     "import OpenAI from 'openai';\nimport DOMPurify from 'dompurify';\n"
     "const r = await client.chat.completions.create({model,messages});\n"
     "const html = r.choices[0].message.content;\n"
     "el.innerHTML = DOMPurify.sanitize(html);\n",
     {"LLM-OUTPUT-HTML", "LLM-NO-VALIDATION"}),  # HTML downgraded LOW
    # --- prompt construction ---
    ("a.py",
     "import openai\n"
     "q = request.args.get('q')\n"
     "prompt = f'Answer this: {q}'\n",
     {"LLM-NO-VALIDATION"}),           # local var, not req-in-str
    ("a.py",
     "import openai\n"
     "prompt = f'Answer this: {request.args.get(\"q\")}'\n",
     {"PROMPT-USER-CONCAT", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "const system = `You are helpful. User says: ${req.body.text}`;\n",
     {"PROMPT-USER-CONCAT", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "messages = request.json['messages']\n",
     {"PROMPT-USER-CONCAT", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "const prompt = TEMPLATE.format(topic);\n",
     {"LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "prompt = 'Summarise the report'\n",
     {"LLM-NO-VALIDATION"}),
    # --- secrets into context ---
    ("a.py",
     "import openai\n"
     "prompt = f'Use key {os.environ[\"API_KEY\"]} to call the api'\n",
     {"SECRET-IN-CONTEXT", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "const content = `token=${process.env.AUTH_TOKEN} q=${q}`;\n",
     {"SECRET-IN-CONTEXT", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "client = openai.OpenAI(api_key=os.environ['OPENAI_API_KEY'])\n",
     {"LLM-NO-VALIDATION"}),         # init, not prompt — only file-level
    # --- role from request ---
    ("a.js",
     "import OpenAI from 'openai';\n"
     "messages.push({role: req.body.role, content: t});\n",
     {"ROLE-FROM-REQ", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "role = 'user'\n"
     "messages.append({'role': role, 'content': t})\n",
     {"LLM-NO-VALIDATION"}),
    # --- loops ---
    ("a.py",
     "import openai\n"
     "while True:\n"
     "    r = client.chat.completions.create(model='g', messages=m)\n"
     "    if done(r): break\n",
     {"UNBOUNDED-AGENT-LOOP", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "step = 0\n"
     "while True:\n"
     "    step += 1\n"
     "    if step >= max_steps: break\n"
     "    r = client.chat.completions.create(model='g', messages=m)\n",
     {"LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "for (;;) {\n"
     "  const r = await client.chat.completions.create({model,messages});\n"
     "  if (r.done) break;\n"
     "}\n",
     {"UNBOUNDED-AGENT-LOOP", "LLM-NO-VALIDATION"}),
    # --- logging ---
    ("a.py",
     "import openai\n"
     "logging.info('sending %s', messages)\n"
     "r = client.chat.completions.create(model='g', messages=messages)\n",
     {"LLM-LOG-CONTENT", "LLM-NO-VALIDATION"}),
    ("a.js",
     "import OpenAI from 'openai';\n"
     "console.log('prompt:', system_prompt);\n"
     "const r = await client.chat.completions.create({model,messages});\n",
     {"LLM-LOG-CONTENT", "LLM-NO-VALIDATION"}),
    # --- validation gate ---
    ("a.py",
     "import openai\n"
     "r = client.chat.completions.create(model='g', messages=m)\n"
     "data = json.loads(r.choices[0].message.content)\n",
     set()),                           # json.loads counts as parse
    ("a.py",
     "import openai\n"
     "r = client.chat.completions.create(model='g', messages=m)\n"
     "print(r)\n",
     {"LLM-NO-VALIDATION"}),
    ("a.ts",
     "import OpenAI from 'openai';\nimport { z } from 'zod';\n"
     "const r = await client.chat.completions.create({model,messages});\n"
     "const o = schema.safeParse(r);\n",
     set()),
    # --- FP guards ---
    ("a.py",
     "import os\n"
     "resp = session.get(url)\n"
     "print(resp)\n",
     set()),                           # no LLM vocab at all
    ("a.py",
     "# resp = openai.chat.completions.create(model='g', messages=m)\n"
     "# exec(resp)\n",
     set()),                           # comments stripped
    ("a.py",
     "import openai\n"
     "resp = client.chat.completions.create(model='g', messages=m)\n"
     "# exec(resp)\n",
     {"LLM-NO-VALIDATION"}),           # sink commented out
    ("test/test_a.py",
     "import openai\n"
     "r = client.chat.completions.create(model='g', messages=m)\n"
     "exec(r.choices[0].message.content)\n",
     {"LLM-OUTPUT-EXEC", "LLM-NO-VALIDATION"}),  # EXEC downgraded MED
    ("a.js",
     "import { chat } from './llm';\n"
     "const r = await chat(messages);\n"
     "const u = r.url || fallback;\n"
     "fetch(u);\n",
     {"LLM-OUTPUT-FETCH", "LLM-NO-VALIDATION"}),
    ("a.py",
     "import openai\n"
     "resp = client.chat.completions.create(model='g', messages=m)\n"
     "db.execute('SELECT 1')\n",       # constant SQL, no taint
     {"LLM-NO-VALIDATION"}),
]


def run_selftest():
    import tempfile
    fails = []
    for i, (name, src, expect) in enumerate(SELFTEST):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, name)
            os.makedirs(os.path.dirname(p), exist_ok=True)
            open(p, "w").write(src)
            got = {f.code for f in scan_file(p, name)}
            missing = expect - got
            extra = got - expect
            if missing or extra:
                fails.append(("v%d %s" % (i, name),
                              "missing=%s extra=%s" % (missing, extra)))
    if fails:
        for f_ in fails:
            print("SELFTEST FAIL", f_)
        return 1
    print("SELFTEST PASS (%d vectors)" % len(SELFTEST))
    return 0


def main():
    args = [a for a in sys.argv[1:] if a != "--json"]
    want_json = "--json" in sys.argv
    if "--selftest" in sys.argv:
        return run_selftest()
    if not args:
        print(__doc__)
        return 2
    root = args[0]
    targets = find_targets(root)
    findings = []
    for p in targets:
        rel = os.path.relpath(p, root) if os.path.isdir(root) \
            else os.path.basename(p)
        findings += scan_file(p, rel)
    if want_json:
        print(json.dumps({"findings": [f.as_dict() for f in findings],
                          "scanned": len(targets)}, indent=1))
    else:
        print(render(findings, targets))
    return 1 if any(f.sev in ("HIGH", "MEDIUM") for f in findings) else 0


if __name__ == "__main__":
    sys.exit(main())