Building an AI-powered code vulnerability scanner

Your security team can't review every pull request by hand. Static analysis tools like Semgrep and Bandit catch known patterns, but they miss logic flaws, insecure data flows, and context-dependent vulnerabilities. Layering a language model on top of those tools — using LLM reasoning to triage and explain findings — bridges that gap without replacing what already works.

Why Static Analysis Alone Falls Short

Traditional SAST tools operate on AST patterns and regex. They're fast, deterministic, and great for known issues — hardcoded secrets, obvious SQL injections, unvalidated inputs. But they fail at:

• Business logic flaws — a function that looks safe in isolation but is dangerous in context
• Chained vulnerabilities — where vuln A + vuln B = full compromise
• False positive noise — teams disable rules when the signal-to-noise ratio degrades

A language model can read a function, understand its intent, and reason about what an attacker could do with it. That reasoning is what you want in a scanner.

Architecture Overview

The scanner has three layers:

1. AST extraction — parse Python source files into their raw structure
2. Pattern pre-filter — run Bandit to surface known bad patterns cheaply
3. LLM analysis — send suspicious code chunks to a language model with a structured prompt; get back findings with severity, CWE, and remediation advice

We keep the LLM calls targeted. You don't want to send 10,000 lines of boilerplate to an API — you want to send the 200 lines that Bandit flagged or that heuristics marked as risky.

Building the Core Scanner

import ast
import subprocess
import json
import os
from pathlib import Path
from openai import OpenAI  # any OpenAI-compatible endpoint works

client = OpenAI(api_key=os.environ["LLM_API_KEY"], base_url=os.environ.get("LLM_BASE_URL"))

SYSTEM_PROMPT = """You are a security code reviewer. Analyze the provided Python function 
for vulnerabilities. Return a JSON object with:
- vulnerabilities: list of {cwe, severity (low/medium/high/critical), description, line_hint, remediation}
- safe: boolean (true if no significant issues found)

Be precise. Only flag real issues. Do not invent problems."""

def extract_functions(filepath: str) -> list[dict]:
    """Parse a Python file and extract all function definitions with their source."""
    source = Path(filepath).read_text()
    tree = ast.parse(source)
    lines = source.splitlines()
    functions = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            start = node.lineno - 1
            end = node.end_lineno
            func_source = "\n".join(lines[start:end])
            functions.append({
                "name": node.name,
                "lineno": node.lineno,
                "source": func_source,
            })
    return functions

def run_bandit(filepath: str) -> set[int]:
    """Run Bandit and return line numbers with findings."""
    result = subprocess.run(
        ["bandit", "-f", "json", "-q", filepath],
        capture_output=True, text=True
    )
    try:
        data = json.loads(result.stdout)
        return {r["line_number"] for r in data.get("results", [])}
    except json.JSONDecodeError:
        return set()

def analyze_function(func: dict) -> dict | None:
    """Send a function to the LLM for vulnerability analysis."""
    prompt = f"File function `{func['name']}` (line {func['lineno']}):\n\n```
{% endraw %}
python\n{func['source']}\n
{% raw %}
```"

    response = client.chat.completions.create(
        model="gpt-4o-mini",  # swap for any model you deploy
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        response_format={"type": "json_object"},
        temperature=0,
    )

    result = json.loads(response.choices[0].message.content)
    if not result.get("safe") and result.get("vulnerabilities"):
        return {"function": func["name"], "line": func["lineno"], **result}
    return None

def scan_file(filepath: str) -> list[dict]:
    """Full scan pipeline: extract → pre-filter → LLM analyze."""
    flagged_lines = run_bandit(filepath)
    functions = extract_functions(filepath)

    findings = []
    for func in functions:
        # Only send to LLM if Bandit flagged it OR function touches risky operations
        risky_keywords = {"eval", "exec", "subprocess", "pickle", "yaml.load", "os.system"}
        is_risky = (
            func["lineno"] in flagged_lines
            or any(kw in func["source"] for kw in risky_keywords)
        )
        if not is_risky:
            continue

        result = analyze_function(func)
        if result:
            findings.append(result)

    return findings

The pre-filter keeps LLM costs sane. On a typical 5,000-line Python service, you might send 15–25 functions to the model instead of 200+. At gpt-4o-mini pricing, that's under $0.01 per scan.

Integrating into CI/CD

Wrap the scanner in a CLI that exits non-zero on critical findings — this is what your CI pipeline will react to:

#!/usr/bin/env python3
import sys
import argparse
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description="AI code vulnerability scanner")
    parser.add_argument("paths", nargs="+", help="Files or directories to scan")
    parser.add_argument("--fail-on", default="high", choices=["low","medium","high","critical"])
    parser.add_argument("--output", default="text", choices=["text","json"])
    args = parser.parse_args()

    severity_order = {"low": 0, "medium": 1, "high": 2, "critical": 3}
    threshold = severity_order[args.fail_on]

    all_findings = []
    files = []
    for p in args.paths:
        path = Path(p)
        if path.is_dir():
            files.extend(path.rglob("*.py"))
        else:
            files.append(path)

    for filepath in files:
        findings = scan_file(str(filepath))
        for f in findings:
            f["file"] = str(filepath)
            all_findings.append(f)

    if args.output == "json":
        import json
        print(json.dumps(all_findings, indent=2))
    else:
        for f in all_findings:
            for vuln in f.get("vulnerabilities", []):
                sev = vuln["severity"].upper()
                print(f"[{sev}] {f['file']}:{f['line']} ({f['function']}) — {vuln['cwe']}")
                print(f"  {vuln['description']}")
                print(f"  Fix: {vuln['remediation']}\n")

    worst = max(
        (severity_order.get(v["severity"], 0) for f in all_findings for v in f.get("vulnerabilities", [])),
        default=-1
    )
    if worst >= threshold:
        print(f"❌ Scan failed: findings at or above '{args.fail_on}' severity found.")
        sys.exit(1)

    print(f"✅ Scan passed ({len(all_findings)} findings below threshold).")

if __name__ == "__main__":
    main()

Add this to your GitHub Actions workflow:

- name: AI vulnerability scan
  run: |
    pip install openai bandit
    python scanner.py src/ --fail-on high --output json > scan-results.json
  env:
    LLM_API_KEY: ${{ secrets.LLM_API_KEY }}

Tuning Signal Quality

A few things that matter in practice:

Chunking strategy — don't send entire files. Function-level chunks give the model enough context without overwhelming it with irrelevant code. If you have class methods that reference class state, include the class attributes in the prompt.

Prompt precision — vague prompts get vague answers. The system prompt above specifies the exact output schema and tells the model to be conservative. "Only flag real issues" cuts false positives significantly compared to "find all security issues."

Model selection — for code analysis, a smaller model fine-tuned on code often outperforms a larger general-purpose model. Run a benchmark against your own codebase before committing to a model. Keep a set of known-vulnerable functions as a test fixture and measure recall.

Rate limiting and cost — cache LLM results keyed by a hash of the function source. If the function didn't change between commits, reuse the previous finding rather than re-querying.

You can find a ready-to-use security hardening checklist that covers secure code review criteria — useful as a reference when tuning your model's prompt.

The Takeaway

An LLM-powered scanner doesn't replace Bandit or Semgrep — it augments them. The pattern-based tools are your first pass: fast, cheap, zero API cost. The model handles the nuanced cases those tools miss.

The key discipline is keeping LLM calls targeted. Scan everything with patterns, invoke the model only on suspicious code, cache aggressively. That approach keeps scan time under 30 seconds for most codebases and cost well under a dollar per day in CI.

The code above is production-ready in the sense that it handles errors, respects cost, and integrates with standard CI. You'll want to add retry logic, structured logging, and a proper config file for the threshold settings — but the core pipeline is solid.

---

I run AYI NEDJIMI Consultants, a cybersecurity consulting firm. We publish free security hardening checklists — PDF and Excel.