slice_context.py

#!/usr/bin/env python3
"""
slice_context.py — Convert a PDG slice graph into structured vulnerability context.

Takes a slice dict produced by preprocess_slice_pdg.py (or v3) and returns
a structured summary of the vulnerability pattern: what the dangerous sink is,
where the input comes from, whether a guard is present, and natural-language
context ready for injection into an LLM harness-generation prompt.

The goal is to pre-compute the hard analysis step (what is dangerous and why)
so that a small sparse model (e.g. Qwen 35B-A3B) only has to do the easy step
(generate code from a specification) rather than reason about data flow itself.

Usage (standalone):
    python slice_context.py path/to/function.ll
    python slice_context.py path/to/function.ll --json

Usage (as library):
    from slice_context import summarize_slice, format_for_llm
    summary = summarize_slice(graph_dict, fn_name="process_packet")
    prompt_block = format_for_llm(summary, score=0.82)
"""

import json
import sys
from pathlib import Path

# ---------------------------------------------------------------------------
# Opcode constants (must match OPCODE_VOCAB in preprocess_slice_pdg.py)
# ---------------------------------------------------------------------------

IDX_CONTEXT   = 0
IDX_ARGUMENT  = 1
IDX_MOCK      = 75    # external function / global variable reference
IDX_CONST_INT = 76
IDX_CONST_FP  = 77
IDX_UNDEF     = 78
IDX_UNKNOWN   = 79

ICMP_OPCODES = frozenset({
    46,              # icmp (unspecified predicate)
    80, 81,          # eq, ne
    82, 83, 84, 85,  # slt, sle, sgt, sge
    86, 87, 88, 89,  # ult, ule, ugt, uge
})

# Subsets of ICMP_OPCODES by structural meaning
_BOUNDS_CHECK_OPCODES = frozenset({82, 83, 84, 85, 86, 87, 88, 89})  # slt/sle/sgt/sge/ult/ule/ugt/uge
_NULL_CHECK_OPCODES   = frozenset({80, 81})                           # eq, ne

_ICMP_LABEL: dict[int, str] = {
    80: "eq", 81: "ne",
    82: "slt", 83: "sle", 84: "sgt", 85: "sge",
    86: "ult", 87: "ule", 88: "ugt", 89: "uge",
    46: "icmp",
}

# Human-readable display names for IR opcodes that appear as sinks.
# Keyed by the internal name used in _SINK_INFO / sink_fn_names.
_DISPLAY_NAMES: dict[str, str] = {
    "getelementptr": "array/ptr-subscript",
    "alloca":        "vla-stack-alloc",
}

# Dangerous sink descriptions: fn_name → (what it does, what to fuzz)
_SINK_INFO: dict[str, tuple[str, str]] = {
    "strcpy":   ("copies string to dest without bounds check",
                 "source string length vs dest buffer size"),
    "strncpy":  ("copies at most n bytes — check dest capacity vs n",
                 "n relative to dest size; source not null-terminated when n < len(src)"),
    "strcat":   ("appends string without bounds check",
                 "remaining dest capacity vs source length"),
    "strncat":  ("appends at most n bytes",
                 "n relative to remaining dest capacity"),
    "memcpy":   ("copies n bytes from src to dest — no overlap or bounds check",
                 "n relative to dest buffer size; n=0, n=SIZE_MAX, n=dest_size+1"),
    "memmove":  ("moves n bytes — allows overlap, no bounds check",
                 "n relative to dest size"),
    "memset":   ("fills n bytes — no bounds check",
                 "n relative to buffer size"),
    "bcopy":    ("copies n bytes src→dest (deprecated) — no bounds check",
                 "n relative to dest size"),
    "sprintf":  ("formats into dest without length limit",
                 "formatted output length vs dest size; user-controlled format string"),
    "snprintf": ("formats at most n bytes — verify n matches dest",
                 "n relative to dest size; off-by-one at n"),
    "vsprintf": ("formats variadic args without length limit",
                 "formatted output length vs dest size"),
    "vsnprintf":("formats variadic args, at most n bytes",
                 "n relative to dest size"),
    "gets":     ("reads line with no length limit — inherently unsafe",
                 "any non-empty input; function is unconditionally vulnerable"),
    "fgets":    ("reads at most n-1 chars — check n matches buffer",
                 "n relative to buffer size; input exactly n-1 chars with no newline"),
    "scanf":    ("reads formatted input — %%s without width is unbounded",
                 "input longer than any %%s field width; format string control"),
    "sscanf":   ("reads from string — %%s without width is unbounded",
                 "input longer than any %%s field width"),
    "fscanf":   ("reads from file — %%s without width is unbounded",
                 "input longer than any %%s field width"),
    "read":     ("reads up to count bytes — no null termination",
                 "count relative to buffer size; count=0; count=SIZE_MAX"),
    "recv":     ("receives up to len bytes — no null termination",
                 "len relative to buffer size; len=0; len=INT_MAX"),
    "recvfrom": ("receives up to len bytes from socket",
                 "len relative to buffer size; truncated messages"),
    "pread":    ("reads up to count bytes at offset",
                 "count relative to buffer size; offset overflow"),
    "malloc":   ("allocates heap — return may be null; size may overflow",
                 "size=0; size=SIZE_MAX; integer overflow in size computation"),
    "calloc":   ("allocates zeroed heap — n*size may overflow; return may be null",
                 "n*size integer overflow; n=0 or size=0; very large n"),
    "realloc":  ("reallocates — null return leaves original pointer valid but lost",
                 "new_size=0 (frees); new_size=SIZE_MAX; null return before use"),
    "free":     ("frees heap memory — double-free or use-after-free if called again",
                 "call target twice with same pointer; access after free"),
    "xmalloc":  ("malloc wrapper — check whether it aborts or returns null on failure",
                 "very large size; size=0"),
    "xrealloc": ("realloc wrapper",
                 "new_size=0; very large size"),
    "printf":   ("format string to stdout — dangerous if format is user-controlled",
                 "user-controlled format string with %%n, %%s, %%x"),
    "fprintf":  ("format string to file",
                 "user-controlled format string"),
    "syslog":   ("format string to syslog",
                 "user-controlled format string with %%n"),
    "err":      ("format string + perror to stderr",
                 "user-controlled format string"),
    "warn":     ("format string warning to stderr",
                 "user-controlled format string"),
    "getelementptr": ("pointer arithmetic with non-constant index — out-of-bounds if unchecked",
                      "index at, near, and beyond array bounds; negative index; index=SIZE_MAX"),
    "alloca":        ("stack allocation with non-constant size (VLA) — stack overflow if unchecked",
                      "size=0; size=SIZE_MAX; negative size (signed wrap); size from user input"),
}


# ---------------------------------------------------------------------------
# Core summarizer
# ---------------------------------------------------------------------------

def summarize_slice(g: dict, fn_name: str = "unknown") -> dict:
    """
    Convert a PDG slice dict into structured vulnerability context.

    Slice dict keys (preprocess_slice_pdg.py format):
      x             (N, 1) int64  — opcode ID per node
      edge_index    (2, E) int64  — [src; dst]
      edge_type     (E,)   int64  — 0=CFG, 1=DFG, 2=context
      sink_mask     (N,)   bool   — True at sink nodes (v3 only; inferred if absent)
      sink_fn_names {int: str}    — node_idx → dangerous function name at that sink

    Returns dict:
      fn_name, slice_size, n_sinks, sinks (list),
      input_channels (list), guard_count, has_guard,
      natural_language (str), harness_hint (str)
    """
    x          = g["x"]
    N          = x.shape[0]
    edge_index = g["edge_index"]
    edge_type  = g["edge_type"]
    E          = edge_index.shape[1] if (edge_index.ndim == 2 and edge_index.size > 0) else 0

    opcodes = [int(x[i, 0]) for i in range(N)]

    sink_fn_names:   dict[int, str] = g.get("sink_fn_names", {})
    source_fn_names: dict[int, str] = g.get("source_fn_names", {})
    sink_mask = g.get("sink_mask", None)

    # ---- build DFG adjacency (used for distance + guard path analysis) -----
    # fwd_dfg[src] = [dst, ...]  — DFG edges only (edge_type == 1)
    from collections import deque
    fwd_dfg: dict[int, list[int]] = {}
    rev_dfg: dict[int, list[int]] = {}
    for e in range(E):
        if int(edge_type[e]) == 1:
            s, d = int(edge_index[0, e]), int(edge_index[1, e])
            fwd_dfg.setdefault(s, []).append(d)
            rev_dfg.setdefault(d, []).append(s)

    # ---- identify sinks ------------------------------------------------
    sinks = []

    if sink_mask is not None:
        for i in range(N):
            if not sink_mask[i]:
                continue
            opc = opcodes[i]
            if opc == 63:   # call
                fn = sink_fn_names.get(i, "unknown")
                sinks.append({"node": i, "type": "dangerous_call", "fn": fn})
            elif opc == 29: # getelementptr
                sinks.append({"node": i, "type": "unguarded_gep", "fn": "getelementptr"})
            elif opc == 26: # alloca with non-constant size (VLA)
                sinks.append({"node": i, "type": "vla_alloca", "fn": "alloca"})
    else:
        # Fallback: scan for call, GEP, and VLA alloca nodes when sink_mask absent
        for i in range(N):
            opc = opcodes[i]
            if opc == 63 and i in sink_fn_names:
                sinks.append({"node": i, "type": "dangerous_call",
                               "fn": sink_fn_names[i]})
            elif opc in (29, 26):
                # Only flag if it has a non-constant DFG predecessor
                preds = []
                for e in range(E):
                    if int(edge_type[e]) == 1 and int(edge_index[1, e]) == i:
                        preds.append(int(edge_index[0, e]))
                const_ids = {IDX_CONST_INT, IDX_CONST_FP, IDX_UNDEF, IDX_CONTEXT}
                if any(opcodes[p] not in const_ids for p in preds):
                    fn = "getelementptr" if opc == 29 else "alloca"
                    kind = "unguarded_gep" if opc == 29 else "vla_alloca"
                    sinks.append({"node": i, "type": kind, "fn": fn})

    # ---- input channels ------------------------------------------------
    input_channels = []
    if any(opc == IDX_ARGUMENT for opc in opcodes):
        input_channels.append("function_argument")
    if any(opc == IDX_MOCK for opc in opcodes):
        input_channels.append("external_call_return")
    if not input_channels:
        input_channels.append("internal_computation")

    # ---- external input flag -------------------------------------------
    # True if a known input-source function (recv/read/fgets/...) appears as
    # a mock node in the slice — meaning network/user data reaches the sink.
    is_external_input = bool(source_fn_names)
    external_sources  = sorted(set(source_fn_names.values()))

    # ---- guard check ---------------------------------------------------
    guard_count = sum(1 for opc in opcodes if opc in ICMP_OPCODES)
    has_guard   = guard_count > 0

    # ---- guard direction -----------------------------------------------
    # Classify which kinds of comparisons guard this slice.
    # bounds_check: relational (<, <=, >, >=) — protects buffer writes
    # null_check:   equality (==, !=) — protects pointer dereferences
    bounds_check_count = sum(1 for opc in opcodes if opc in _BOUNDS_CHECK_OPCODES)
    null_check_count   = sum(1 for opc in opcodes if opc in _NULL_CHECK_OPCODES)
    # Dominant guard type; "mixed" when both present
    if bounds_check_count > 0 and null_check_count > 0:
        guard_type = "mixed"
    elif bounds_check_count > 0:
        guard_type = "bounds_check"
    elif null_check_count > 0:
        guard_type = "null_check"
    else:
        guard_type = "none"
    # Collect the specific predicate labels present
    guard_predicates = sorted({_ICMP_LABEL[opc]
                                for opc in opcodes
                                if opc in ICMP_OPCODES and opc in _ICMP_LABEL})

    # ---- sink-source hop distance --------------------------------------
    # BFS forward from all argument and input-source mock nodes through DFG edges.
    # Records minimum hop count to each sink node.
    source_nodes = {i for i, opc in enumerate(opcodes)
                    if opc == IDX_ARGUMENT or i in source_fn_names}
    sink_node_ids = {s["node"] for s in sinks}

    dist: dict[int, int] = {n: 0 for n in source_nodes}
    queue = deque(source_nodes)
    while queue:
        node = queue.popleft()
        for nxt in fwd_dfg.get(node, []):
            if nxt not in dist:
                dist[nxt] = dist[node] + 1
                queue.append(nxt)

    # Attach distance to each sink entry
    for s in sinks:
        s["distance"] = dist.get(s["node"])   # None if unreachable from sources

    min_distance = (min((s["distance"] for s in sinks if s["distance"] is not None),
                        default=None)
                    if sinks else None)

    # ---- integer truncation signal ------------------------------------
    # trunc (opcode 48) narrows an integer type (e.g. i64->i32).
    # When present in a slice that also has size-taking sinks, it is a
    # precursor pattern for integer-truncation vulnerabilities.
    _SIZE_SINKS = frozenset({
        "memcpy", "memmove", "memset", "bcopy",
        "malloc", "calloc", "realloc", "xmalloc", "xrealloc",
        "read", "recv", "recvfrom", "pread",
        "snprintf", "vsnprintf", "fgets", "alloca",
    })
    trunc_count = sum(1 for opc in opcodes if opc == 48)
    has_trunc   = trunc_count > 0 and any(
        s.get("fn") in _SIZE_SINKS for s in sinks
    )

    # ---- deduplicate sinks by function name (preserve first-seen order) ----
    from collections import Counter, OrderedDict
    sink_counts: dict[str, int] = Counter(s.get("fn") or "unknown" for s in sinks)
    seen: set[str] = set()
    unique_sinks = []
    for s in sinks:
        fn = s.get("fn") or "unknown"
        if fn not in seen:
            seen.add(fn)
            unique_sinks.append(s)

    # ---- natural language + harness hint (deduplicated) -----------------
    sink_strs  = []
    hint_parts = []

    for s in unique_sinks:
        fn      = s.get("fn") or "unknown"
        display = _DISPLAY_NAMES.get(fn, fn)
        count   = sink_counts[fn]
        info    = _SINK_INFO.get(fn)
        suffix  = f" ×{count}" if count > 1 else ""
        if info:
            what, probe = info
            sink_strs.append(f"`{display}`{suffix} ({what})")
            if f"fuzz {probe}" not in hint_parts:
                hint_parts.append(f"fuzz {probe}")
        else:
            sink_strs.append(f"`{display}`{suffix} (dangerous operation)")
            if "fuzz all arguments" not in hint_parts:
                hint_parts.append("fuzz all arguments")

    if not sinks:
        sink_strs  = ["no explicit dangerous sink identified — scored by full-graph structure"]
        hint_parts = ["fuzz all inputs broadly"]

    # ---- guard density -------------------------------------------------
    n_sinks = len(sinks)
    if n_sinks == 0:
        guard_density       = 0.0
        guard_density_label = "no sinks"
    elif not has_guard:
        guard_density       = float("inf")
        guard_density_label = "UNGUARDED"
    else:
        guard_density = n_sinks / guard_count   # sinks per guard — higher = worse
        if guard_density >= 10:
            guard_density_label = "very sparse"
        elif guard_density >= 5:
            guard_density_label = "sparse"
        elif guard_density >= 2:
            guard_density_label = "moderate"
        else:
            guard_density_label = "well-covered"

    channel_note = " or ".join(input_channels)
    if not has_guard:
        guard_note = "no comparison (icmp) in slice — sink appears UNGUARDED"
    elif n_sinks == 0:
        guard_note = f"{guard_count} comparison(s) in slice"
    else:
        guard_note = (
            f"{guard_count} guard(s) / {n_sinks} sink(s)"
            f" = {guard_density:.1f} sinks/guard ({guard_density_label})"
        )

    ext_note = ""
    if is_external_input:
        src_list = ", ".join(external_sources) if external_sources else "external source"
        ext_note = f" [network/user-controlled via {src_list}]"

    dist_note = ""
    if min_distance is not None:
        dist_note = f" Minimum source-to-sink distance: {min_distance} hop(s)."

    natural_language = (
        f"Function `{fn_name}` contains: {'; '.join(sink_strs)}. "
        f"Input originates from: {channel_note}{ext_note}. "
        f"Guard status: {guard_note}. "
        f"{dist_note}"
        f"Slice: {N} nodes, {n_sinks} sink(s) ({len(unique_sinks)} unique type(s))."
    )

    if has_trunc:
        hint_parts.append(
            f"fuzz integer truncation: supply values > INT_MAX / > UINT32_MAX as size"
        )

    harness_hint = " | ".join(hint_parts)

    return {
        "fn_name":            fn_name,
        "slice_size":         N,
        "n_sinks":            n_sinks,
        "n_unique_sinks":     len(unique_sinks),
        "sinks":              sinks,
        "sink_counts":        dict(sink_counts),
        "input_channels":     input_channels,
        "is_external_input":  is_external_input,
        "external_sources":   external_sources,
        "guard_count":        guard_count,
        "has_guard":          has_guard,
        "guard_type":         guard_type,
        "guard_predicates":   guard_predicates,
        "bounds_check_count": bounds_check_count,
        "null_check_count":   null_check_count,
        "guard_density":      guard_density,
        "guard_density_label":guard_density_label,
        "min_distance":       min_distance,
        "trunc_count":        trunc_count,
        "has_trunc":          has_trunc,
        "natural_language":   natural_language,
        "harness_hint":       harness_hint,
    }


def format_for_llm(summary: dict, score: float | None = None,
                   width: int = 60) -> str:
    """
    Format a slice summary as a labelled block for LLM prompt injection.

    Example output:
        === GNN Vulnerability Context ===
        Suspicion score : 73.2%  (SUSPICIOUS)
        Sinks           : memcpy (copies n bytes without bounds check)
        Input channels  : function_argument
        Guard status    : no comparison in slice — sink UNGUARDED
        Harness target  : fuzz n relative to dest buffer size; n=0, n=SIZE_MAX
        Slice           : 31 nodes, 1 sink
        ==================================
    """
    sep = "=" * width
    lines = [sep, "GNN Vulnerability Context"]

    if score is not None:
        flag = "SUSPICIOUS" if score >= 0.5 else "low suspicion"
        lines.append(f"Suspicion score : {score:.1%}  ({flag})")

    # Deduplicate: one label per unique sink type, with ×N count
    sink_counts = summary.get("sink_counts", {})
    seen: set[str] = set()
    sink_labels = []
    for s in summary["sinks"]:
        fn = s.get("fn", "unknown")
        if fn in seen:
            continue
        seen.add(fn)
        display = _DISPLAY_NAMES.get(fn, fn)
        info    = _SINK_INFO.get(fn)
        short   = info[0][:50] if info else "dangerous operation"
        count   = sink_counts.get(fn, 1)
        tag     = f" ×{count}" if count > 1 else ""
        sink_labels.append(f"{display}{tag} — {short}")
    lines.append("Sinks           : " + ("; ".join(sink_labels) if sink_labels
                                          else "none identified"))

    # Input channels + external input flag
    channels = ", ".join(summary["input_channels"])
    if summary.get("is_external_input"):
        srcs = ", ".join(summary.get("external_sources", []))
        channels += f"  [external_input=YES via {srcs}]" if srcs else "  [external_input=YES]"
    lines.append("Input channels  : " + channels)

    n_sinks = summary["n_sinks"]
    gc      = summary["guard_count"]
    gtype   = summary.get("guard_type", "none")
    preds   = summary.get("guard_predicates", [])
    pred_str = f" ({', '.join(preds)})" if preds else ""
    if not summary["has_guard"]:
        guard = "NO icmp in slice — sink appears UNGUARDED"
    elif n_sinks == 0:
        guard = f"{gc} comparison(s){pred_str} present"
    else:
        gd    = summary.get("guard_density", n_sinks / gc)
        label = summary.get("guard_density_label", "")
        if gtype == "bounds_check":
            gtype_str = "bounds-check"
        elif gtype == "null_check":
            gtype_str = "null-check only — may not protect buffer writes"
        elif gtype == "mixed":
            gtype_str = "bounds-check + null-check"
        else:
            gtype_str = gtype
        guard = (f"{gc} guard(s){pred_str} [{gtype_str}]"
                 f" / {n_sinks} sink(s) = {gd:.1f} sinks/guard ({label})")
    lines.append("Guard status    : " + guard)

    # Sink-source distance
    min_dist = summary.get("min_distance")
    if min_dist is not None:
        lines.append(f"Distance        : {min_dist} hop(s) source→sink")

    if summary.get("has_trunc"):
        lines.append(f"Trunc warning   : {summary['trunc_count']} integer narrowing(s) — check size args for truncation")
    lines.append("Harness target  : " + summary["harness_hint"])
    lines.append(f"Slice           : {summary['slice_size']} nodes, "
                 f"{summary['n_sinks']} sink(s)")
    lines.append(sep)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Standalone CLI (demo mode — parses IR directly without a trained model)
# ---------------------------------------------------------------------------

def _demo_cli():
    import argparse
    import json as _json
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("ir_file", help=".ll file to analyse")
    ap.add_argument("--json",  action="store_true", help="output raw JSON summary")
    ap.add_argument("--debug", action="store_true", help="show parse errors instead of silencing them")
    args = ap.parse_args()

    HERE = Path(__file__).parent
    sys.path.insert(0, str(HERE))
    from preprocess_slice_pdg import ir_to_graph_slice_pdg

    import re
    ir_text = Path(args.ir_file).read_text(errors="replace")

    # Parse the full module once to enumerate non-declaration functions.
    # Passing the full IR (not a per-function split) to ir_to_graph_slice_pdg
    # preserves all declare stubs and globals — cross-function calls that appear
    # in the same source file are visible, and llvmlite won't reject them.
    import llvmlite.binding as _llvm
    try:
        _mod = _llvm.parse_assembly(ir_text)
    except Exception as exc:
        if args.debug:
            print(f"Module parse FAILED: {exc}")
        print(f"ERROR: could not parse {args.ir_file}: {exc}")
        sys.exit(1)

    functions = [(fn.name, fn.name)
                 for fn in _mod.functions if not fn.is_declaration]

    if not functions:
        print(f"ERROR: no non-declaration functions found in {args.ir_file}")
        sys.exit(1)

    for fn_name, _ in functions:
        if args.debug:
            print(f"[{fn_name}] parse OK (full-module mode)")
        g = ir_to_graph_slice_pdg(ir_text, fn_name=fn_name)
        if g is None:
            print(f"[{fn_name}] — could not extract graph (no basic blocks)\n")
            continue

        summary = summarize_slice(g, fn_name=fn_name)

        if args.json:
            print(_json.dumps({k: v for k, v in summary.items()
                                if k != "sinks"}, indent=2))
        else:
            print(format_for_llm(summary))
            print(f"Natural language:\n  {summary['natural_language']}\n")


if __name__ == "__main__":
    _demo_cli()