llm-security/scanners/lib/py-ast-taint.py

#!/usr/bin/env python3
"""py-ast-taint.py — PARSE-ONLY Python taint helper for the AST scanner.

Reads a single target .py path (argv[1]), parses it with ast.parse, and walks
the tree for variable-level taint (data flow from an input source to a
dangerous sink) within each function/module scope. It prints one JSON object
to stdout:

    {"status": "ok", "findings": [
        {"rule": "...", "severity": "...", "line": int,
         "source": "...", "sink": "...", "message": "..."}]}

On a parse error it prints {"status": "error", "message": "..."} and exits 2.

SAFETY INVARIANT: this helper only PARSES the target (ast.parse). It never
exec/eval/compile/imports the target, so analysing hostile code has no side
effects. The only filesystem access is reading the target as text.
"""

import ast
import json
import sys

# Sinks: dotted name (or bare builtin) -> (rule, severity, category)
CODE_EXEC_SINKS = {
    "exec": ("AST-CODE-EXEC", "critical"),
    "eval": ("AST-CODE-EXEC", "critical"),
    "os.system": ("AST-CMD-EXEC", "critical"),
    "os.popen": ("AST-CMD-EXEC", "critical"),
}
NET_SINKS = {
    "requests.post": ("AST-NET-EXFIL", "critical"),
    "requests.put": ("AST-NET-EXFIL", "critical"),
    "requests.patch": ("AST-NET-EXFIL", "critical"),
}

READ_SOURCE_CALLS = {
    "os.getenv": "os.getenv",
    "input": "input",
    "requests.get": "requests.get",
    "requests.request": "requests.request",
}


def dotted(node):
    """Return the dotted name for a Name/Attribute chain, else None."""
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        base = dotted(node.value)
        return base + "." + node.attr if base else None
    return None


def open_mode(call):
    """Best-effort 'mode' string of an open(...) call (default 'r')."""
    if len(call.args) >= 2 and isinstance(call.args[1], ast.Constant):
        return str(call.args[1].value)
    for kw in call.keywords:
        if kw.arg == "mode" and isinstance(kw.value, ast.Constant):
            return str(kw.value.value)
    return "r"


def is_write_open(call):
    return any(c in open_mode(call) for c in ("w", "a", "x", "+"))


def source_label(value):
    """If `value` is a taint source expression, return a label, else None."""
    if isinstance(value, ast.Subscript):
        if dotted(value.value) == "os.environ":
            return "os.environ"
        return source_label(value.value)
    if isinstance(value, ast.Attribute):
        d = dotted(value)
        if d == "os.environ":
            return "os.environ"
        if d and d.startswith("sys.stdin"):
            return "sys.stdin"
        return None
    if isinstance(value, ast.Call):
        fd = dotted(value.func)
        if fd in READ_SOURCE_CALLS:
            return READ_SOURCE_CALLS[fd]
        if fd == "open":
            return None if is_write_open(value) else "open"
        if fd and fd.startswith("sys.stdin"):
            return "sys.stdin"
        # Chained access on a source, e.g. open(p).read() or sys.stdin.read()
        if isinstance(value.func, ast.Attribute):
            return source_label(value.func.value)
    return None


def assigned_names(target):
    """Yield bound Name ids for an assignment target (Name / Tuple / List)."""
    if isinstance(target, ast.Name):
        yield target.id
    elif isinstance(target, (ast.Tuple, ast.List)):
        for elt in target.elts:
            yield from assigned_names(elt)


def walk_scope(body):
    """Pre-order walk of a scope's nodes, treating a nested function/class/lambda
    as opaque (its body belongs to a separate scope and is analysed on its own)."""
    stack = list(body)
    while stack:
        node = stack.pop(0)
        yield node
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
            continue  # nested scope — do not descend
        stack[:0] = list(ast.iter_child_nodes(node))


def tainted_arg(call, tainted):
    """Return (name, info) for the first directly-passed tainted Name arg."""
    candidates = list(call.args) + [kw.value for kw in call.keywords]
    for arg in candidates:
        if isinstance(arg, ast.Name) and arg.id in tainted:
            return arg.id, tainted[arg.id]
    return None


def sink_for(call, write_handles):
    """Return (rule, severity, sink_label) if `call` is a sink, else None."""
    fd = dotted(call.func)
    if fd in CODE_EXEC_SINKS:
        rule, sev = CODE_EXEC_SINKS[fd]
        return rule, sev, fd
    if fd in NET_SINKS:
        rule, sev = NET_SINKS[fd]
        return rule, sev, fd
    if fd and fd.startswith("subprocess."):
        return "AST-CMD-EXEC", "critical", fd
    if isinstance(call.func, ast.Attribute) and call.func.attr == "write":
        recv = call.func.value
        if isinstance(recv, ast.Name) and recv.id in write_handles:
            return "AST-FILE-WRITE", "high", "file.write"
    return None


def analyze_scope(body):
    tainted = {}          # name -> (lineno, source_label)
    write_handles = set()  # names bound to open(..., 'w'/'a')
    findings = []
    for node in walk_scope(body):
        if isinstance(node, ast.Assign):
            src = source_label(node.value)
            if src:
                for name in (n for t in node.targets for n in assigned_names(t)):
                    tainted[name] = (node.lineno, src)
            if isinstance(node.value, ast.Call) and dotted(node.value.func) == "open" \
                    and is_write_open(node.value):
                for name in (n for t in node.targets for n in assigned_names(t)):
                    write_handles.add(name)
        if isinstance(node, ast.Call):
            sink = sink_for(node, write_handles)
            if sink:
                hit = tainted_arg(node, tainted)
                if hit:
                    name, (src_line, src_label) = hit
                    rule, sev, sink_label = sink
                    findings.append({
                        "rule": rule,
                        "severity": sev,
                        "line": getattr(node, "lineno", 0),
                        "source": src_label,
                        "sink": sink_label,
                        "message": (
                            "Tainted value from %s (line %d) reaches %s via `%s`."
                            % (src_label, src_line, sink_label, name)
                        ),
                    })
    return findings


def iter_scopes(tree):
    yield tree.body
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            yield node.body


def main():
    if len(sys.argv) < 2:
        print(json.dumps({"status": "error", "message": "missing target path"}))
        sys.exit(2)
    path = sys.argv[1]
    try:
        with open(path, "r", encoding="utf-8", errors="replace") as fh:
            src = fh.read()
    except OSError as exc:
        print(json.dumps({"status": "error", "message": "read error: %s" % exc}))
        sys.exit(2)
    try:
        tree = ast.parse(src)
    except (SyntaxError, ValueError) as exc:
        print(json.dumps({"status": "error", "message": "parse error: %s" % exc}))
        sys.exit(2)

    findings = []
    seen = set()
    for body in iter_scopes(tree):
        for f in analyze_scope(body):
            key = (f["rule"], f["line"], f["source"], f["sink"])
            if key in seen:
                continue
            seen.add(key)
            findings.append(f)
    print(json.dumps({"status": "ok", "findings": findings}))


if __name__ == "__main__":
    main()