Building an AI-Powered Code Vulnerability Scanner

Static analysis tools like Semgrep, Bandit, and CodeQL catch well-known patterns. They miss logic bugs, insecure data flows that span multiple files, and context-dependent vulnerabilities that require understanding intent. A language model can fill that gap — not as a replacement for traditional SAST, but as a second layer that reasons about what the code actually does.

How It Works

The architecture is straightforward: chunk source files by function or class, send each chunk to a language model with a structured prompt, parse the JSON output, and deduplicate findings across chunks. The key insight is that most real vulnerabilities are local — they live within a single function or a tight call chain. You don't need to send an entire codebase at once.

The pipeline has four stages:

1. Parse — extract functions/classes from source files using AST
2. Prompt — send each chunk with a security-focused system prompt
3. Parse output — extract structured findings (severity, CWE, line, description)
4. Report — deduplicate and rank by severity

Parsing Source Code into Chunks

Using Python's built-in ast module, we can extract function and class bodies with their line numbers:

import ast
import textwrap
from dataclasses import dataclass
from pathlib import Path

@dataclass
class CodeChunk:
    filepath: str
    name: str
    start_line: int
    source: str

def extract_chunks(filepath: str) -> list[CodeChunk]:
    source = Path(filepath).read_text()
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return []

    lines = source.splitlines()
    chunks = []

    for node in ast.walk(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            continue
        start = node.lineno - 1
        end = node.end_lineno
        body = "\n".join(lines[start:end])
        chunks.append(CodeChunk(
            filepath=filepath,
            name=node.name,
            start_line=node.lineno,
            source=textwrap.dedent(body),
        ))

    return chunks

For Go or other languages, tree-sitter is the right tool — it has Python bindings and parsers for 40+ languages. The same pattern applies: walk the AST, extract named blocks, record their line ranges.

The Vulnerability Prompt

The system prompt is the most important part. Vague prompts produce vague findings. You want structured JSON output with a fixed schema so parsing is deterministic:

import json
import httpx

SYSTEM_PROMPT = """You are a security code reviewer. Analyze the provided code for vulnerabilities.

Return ONLY a JSON object with this exact schema:
{
  "findings": [
    {
      "severity": "critical|high|medium|low|info",
      "cwe": "CWE-XXX",
      "line": <integer, relative to chunk start>,
      "title": "<short title>",
      "description": "<what is wrong and why>",
      "remediation": "<concrete fix>"
    }
  ]
}

If no vulnerabilities are found, return {"findings": []}.
Focus on: injection flaws, broken auth, insecure deserialization, hardcoded secrets,
path traversal, SSRF, XXE, missing input validation, race conditions."""

def scan_chunk(chunk: CodeChunk, api_key: str, model: str = "gpt-4.1") -> list[dict]:
    prompt = f"File: {chunk.filepath}\nFunction: {chunk.name} (starts at line {chunk.start_line})\n\n```
{% endraw %}
python\n{chunk.source}\n
{% raw %}
```"

    resp = httpx.post(
        "https://api.openai.com/v1/chat/completions",
        headers={"Authorization": f"Bearer {api_key}"},
        json={
            "model": model,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            "response_format": {"type": "json_object"},
            "temperature": 0,
        },
        timeout=30,
    )
    resp.raise_for_status()

    raw = resp.json()["choices"][0]["message"]["content"]
    parsed = json.loads(raw)
    findings = parsed.get("findings", [])

    # Adjust line numbers to absolute positions
    for f in findings:
        f["line"] = chunk.start_line + f.get("line", 1) - 1
        f["filepath"] = chunk.filepath

    return findings

Setting temperature: 0 and response_format: {"type": "json_object"} is non-negotiable for production use. Without them, you'll get markdown fences around the JSON and inconsistent schemas.

Scanning a Repository

Wire everything together with a concurrency-aware runner. Rate limits on LLM APIs are the main bottleneck, so use a semaphore:

import asyncio
from pathlib import Path

async def scan_repo(repo_path: str, api_key: str, concurrency: int = 5) -> list[dict]:
    py_files = list(Path(repo_path).rglob("*.py"))
    sem = asyncio.Semaphore(concurrency)
    all_findings = []

    async def scan_file(filepath: str):
        chunks = extract_chunks(filepath)
        async with sem:
            for chunk in chunks:
                # Run sync HTTP call in thread pool
                findings = await asyncio.get_event_loop().run_in_executor(
                    None, scan_chunk, chunk, api_key
                )
                all_findings.extend(findings)

    await asyncio.gather(*[scan_file(str(f)) for f in py_files])

    # Sort by severity
    severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
    all_findings.sort(key=lambda x: severity_order.get(x.get("severity", "info"), 5))

    return all_findings

# Usage
if __name__ == "__main__":
    findings = asyncio.run(scan_repo("./myproject", api_key="sk-..."))
    for f in findings:
        print(f"[{f['severity'].upper()}] {f['filepath']}:{f['line']} — {f['title']}")

On a mid-sized Python codebase (150 files, ~300 functions), this runs in under 3 minutes with concurrency=5 and costs roughly $0.40–$0.80 in API tokens depending on the model.

Reducing False Positives

Raw LLM findings have a false positive rate of 20–40% in practice. Two techniques help significantly:

1. Two-pass validation — after the first scan, send findings back to the model and ask: "Is this actually exploitable given the full context?" This cuts false positives roughly in half.

2. Cross-reference with known patterns — before reporting a finding, check it against a SAST tool. If Bandit also flags the same line, confidence goes up. If only the LLM flags it, mark it as needs_review in your report.

For a security team workflow, integrating this with your existing checklist process makes findings actionable. We use a structured approach to vulnerability triage that's documented in our security hardening checklists — the same classification logic applies here.

Limitations You Need to Know

This approach does not replace traditional SAST for two reasons:

• No dataflow tracking — the LLM sees a chunk, not a call graph. Taint analysis across files requires a real AST/IR.
• Non-deterministic — even at temperature 0, different runs can produce different findings. Don't rely on it for compliance evidence.

Use it as a complement: run Semgrep or Bandit first for fast, deterministic pattern matching, then run the LLM scanner on files that changed in a PR for deeper reasoning. This hybrid approach gives you the best of both worlds without the cost of scanning everything every time.

The Takeaway

A language model can catch vulnerability classes that pattern-matching tools miss — business logic flaws, incorrect use of cryptography, subtly wrong authentication checks. The implementation is genuinely simple: parse to AST, chunk by function, prompt with a strict JSON schema, deduplicate.

The real work is in the prompt engineering and the validation layer. A model that returns "possible injection" with no line number or remediation is useless in a review workflow. Force structure from the start, and build the two-pass validation early — it's much harder to retrofit.

---

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