Lab Specification — Module S06: Secure Code Review Harnesses

pip install semgrep pydantic chromadb

from pydantic import BaseModel, Field
from typing import Literal
import subprocess, json

class RawFinding(BaseModel):
    file: str
    start_line: int
    end_line: int
    cwe: str
    cwe_name: str
    severity: Literal["critical", "high", "medium", "low"]
    rule_id: str
    snippet: str
    source: Literal["semgrep", "codeql", "ast"] = "semgrep"

def run_semgrep_on_diff(changed_files: list[str]) -> list[RawFinding]:
    """Run Semgrep on the changed files only."""
    cmd = ["semgrep", "--config", "auto", "--json"] + changed_files
    result = subprocess.run(cmd, capture_output=True, text=True)
    data = json.loads(result.stdout)
    findings = []
    for r in data.get("results", []):
        findings.append(RawFinding(
            file=r["path"],
            start_line=r["start"]["line"],
            end_line=r["end"]["line"],
            cwe=r.get("extra", {}).get("metadata", {}).get("cwe", ["CWE-Unknown"])[0],
            cwe_name=r.get("extra", {}).get("message", "Unknown"),
            severity=r.get("extra", {}).get("severity", "medium"),
            rule_id=r["check_id"],
            snippet=r.get("extra", {}).get("lines", ""),
        ))
    return findings

docker run -d -p 8080:80 --name dvwa vulnerables/web-dvwa
# Clone the DVWA source for static scanning
git clone https://github.com/digininja/DVWA.git
semgrep --config auto --json DVWA/ > dvwa-raw.json

import json, os

async def triage_finding(raw: RawFinding, enclosing_fn: str) -> TriagedFinding:
    prompt = f"""
    You are a security triage analyst. Evaluate this finding.

    Rule: {raw.rule_id} ({raw.cwe})
    Location: {raw.file}:{raw.start_line}
    Matched code:
    ```
    {raw.snippet}
    ```
    Enclosing function:
    ```
    {enclosing_fn}
    ```

    Questions (answer each WITH EVIDENCE):
    1. Is the tainted input reachable from untrusted data in THIS code path?
    2. Is there a sanitizer, validator, or framework control between source and sink?
    3. True positive or false positive?
    4. Confidence: high / medium / low?

    Respond as JSON: {{"reachable": bool, "sanitized": bool, "verdict": "tp"|"fp", "confidence": "high"|"medium"|"low", "reasoning": str}}
    """
    result = await llm_complete(prompt)
    return TriagedFinding(**json.loads(result), original=raw)

def record_feedback(finding_id: str, verdict: str, reasoning: str, store):
    """Human verdict stored for future few-shot retrieval."""
    store.add(
        documents=[f"{verdict}: {reasoning}"],
        metadatas=[{"finding_id": finding_id, "verdict": verdict}],
        ids=[finding_id]
    )

def build_triage_prompt_with_feedback(raw: RawFinding, store) -> str:
    """Retrieve past human verdicts on similar findings as few-shot examples."""
    similar = store.query(query_texts=[raw.snippet], n_results=3)
    examples = "\n".join(similar["documents"][0]) if similar["documents"][0] else "none"
    return f"Prior judgements on similar findings:\n{examples}\n\nNow evaluate: ..."

import ast, hashlib, chromadb

class CodeUnit(BaseModel):
    id: str
    file: str
    symbol: str
    signature: str
    source: str
    callers: list[str] = Field(default_factory=list)
    callees: list[str] = Field(default_factory=list)

def extract_functions(file_path: str) -> list[CodeUnit]:
    """Parse a Python file to CodeUnits via AST."""
    with open(file_path) as f:
        tree = ast.parse(f.read())
    units = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            source = ast.get_source_segment(open(file_path).read(), node)
            unit_id = hashlib.sha256(f"{file_path}:{node.name}:{node.lineno}".encode()).hexdigest()[:16]
            units.append(CodeUnit(
                id=unit_id, file=file_path, symbol=node.name,
                signature=f"{node.name}({', '.join(a.arg for a in node.args.args)})",
                source=source, callers=[], callees=[]
            ))
    return units

def build_call_graph(units: list[CodeUnit]) -> list[CodeUnit]:
    """Populate callers/callees by scanning source for call references."""
    # For each unit, find which other units' symbols it calls, and vice versa.
    ...

def retrieve_context(finding: RawFinding, index, budget_tokens: int = 6000) -> str:
    """For a finding, retrieve enclosing fn + callers + callees + sanitizers, within budget."""
    # 1. Find the enclosing function
    enclosing = index.find_enclosing(finding.file, finding.start_line)
    context_parts = [enclosing]
    # 2. Direct callers (source) and callees (sink)
    context_parts += index.get_callers(enclosing.id)
    context_parts += index.get_callees(enclosing.id)
    # 3. Sanitizer-like functions (semantic search)
    context_parts += index.semantic_search("input validation sanitizer escape", k=3)
    # 4. Truncate to budget (prioritize callers/callees/sanitizers over transitive)
    return truncate_to_budget(context_parts, budget_tokens)

async def generate_patch(finding: TriagedFinding, context: str) -> str:
    prompt = f"""
    Resolve this security finding with a minimal, targeted patch.

    Finding: {finding.original.cwe} — {finding.original.cwe_name}
    Location: {finding.original.file}:{finding.original.start_line}
    Code:
    ```
    {finding.original.snippet}
    ```
    Cross-file context:
    {context}

    Constraints:
    - Minimal diff. Change only what is necessary to resolve the finding.
    - Do not refactor, rename, or reformat unrelated code.
    - Preserve all existing behavior and test coverage.
    - Output a unified diff.
    """
    return await llm_complete(prompt)

def apply_and_verify(patch_diff: str, repo_path: str, finding: RawFinding) -> dict:
    """Apply the patch and run all automated gates."""
    # Gate 1: apply patch
    apply_result = subprocess.run(["git", "apply"], input=patch_diff, text=True, cwd=repo_path)
    if apply_result.returncode != 0:
        return {"gate": "apply", "passed": False}

    # Gate 2: linter clean
    if not run_linter(repo_path):
        return {"gate": "linter", "passed": False}

    # Gate 3: Semgrep re-scan — finding resolved, no new findings
    new_findings = run_semgrep_on_diff([finding.file])
    still_present = any(f.start_line == finding.start_line and f.rule_id == finding.rule_id for f in new_findings)
    if still_present or len(new_findings) > 0:
        return {"gate": "semgrep_rescan", "passed": False, "still_present": still_present, "new": len(new_findings)}

    # Gate 4: test suite passes
    if not run_tests(repo_path):
        return {"gate": "tests", "passed": False}

    return {"gate": "all", "passed": True}

def open_draft_pr(patch_diff: str, finding: TriagedFinding, verification: dict):
    """Simulate opening a draft PR with the patch, finding, and verification evidence."""
    pr_body = f"""
    ## Automated Security Fix

    **Finding**: {finding.original.cwe} — {finding.original.cwe_name}
    **Location**: {finding.original.file}:{finding.original.start_line}
    **Confidence**: {finding.confidence}

    **Verification**:
    - Linter: PASS
    - Semgrep re-scan: PASS (finding resolved, no new findings)
    - Tests: PASS

    ---
    Human review required before merge. This PR is a DRAFT.
    """
    print(f"[DRAFT PR] {pr_body}")
    print(f"[PATCH]\n{patch_diff}")

# Lab Specification — Module S06: Secure Code Review Harnesses

**Course**: 2A — Building AI Harnesses for Cybersecurity
**Module**: S06 — Secure Code Review Harnesses
**Duration**: 120–150 minutes (four labs, one per sub-section)
**Environment**: Python 3.11+, Semgrep, CodeQL CLI (optional), ChromaDB, Pydantic. A git repo with a vulnerable branch for lab targets. Docker for DVWA (lab 2). An LLM API key for triage and patch generation.

---

## Learning objectives

1. Build a layered code review pipeline (AST → Semgrep → LLM triage → synthesis) that ingests a PR diff and emits structured findings.
2. Implement LLM false-positive triage with confidence scoring and a feedback loop; measure precision/recall against DVWA.
3. Build a function-level semantic codebase index and retrieve cross-file taint context for a finding under a context budget.
4. Engineer an autofix loop (patch → verify → draft PR) with all automated gates; document the approval gate.

---

## Phase 1 — Layered Pipeline on a PR Diff (35 min)

### 1.1 Set up

```bash
pip install semgrep pydantic chromadb
```

Create a git repo with two branches: `main` (clean) and `vulnerable` (introduces 5 known weaknesses: SQL injection, hardcoded secret, path traversal, XSS, insecure deserialization).

### 1.2 Implement the diff-driven Semgrep layer

```python
from pydantic import BaseModel, Field
from typing import Literal
import subprocess, json

class RawFinding(BaseModel):
    file: str
    start_line: int
    end_line: int
    cwe: str
    cwe_name: str
    severity: Literal["critical", "high", "medium", "low"]
    rule_id: str
    snippet: str
    source: Literal["semgrep", "codeql", "ast"] = "semgrep"

def run_semgrep_on_diff(changed_files: list[str]) -> list[RawFinding]:
    """Run Semgrep on the changed files only."""
    cmd = ["semgrep", "--config", "auto", "--json"] + changed_files
    result = subprocess.run(cmd, capture_output=True, text=True)
    data = json.loads(result.stdout)
    findings = []
    for r in data.get("results", []):
        findings.append(RawFinding(
            file=r["path"],
            start_line=r["start"]["line"],
            end_line=r["end"]["line"],
            cwe=r.get("extra", {}).get("metadata", {}).get("cwe", ["CWE-Unknown"])[0],
            cwe_name=r.get("extra", {}).get("message", "Unknown"),
            severity=r.get("extra", {}).get("severity", "medium"),
            rule_id=r["check_id"],
            snippet=r.get("extra", {}).get("lines", ""),
        ))
    return findings
```

### 1.3 Generate the diff and run the pipeline

1. `git diff main...vulnerable --name-only` → list of changed files.
2. Run `run_semgrep_on_diff(changed_files)`.
3. Verify: Semgrep produces raw findings on the vulnerable branch (expect 5+).

### Deliverable
- [ ] Semgrep runs against the diff (changed files only), not the whole repo
- [ ] Raw findings normalized to the `RawFinding` schema
- [ ] Diff-driven: re-running on `main` produces zero findings

---

## Phase 2 — LLM False-Positive Triage against DVWA (35 min)

### 2.1 Run Semgrep on DVWA and measure raw precision

```bash
docker run -d -p 8080:80 --name dvwa vulnerables/web-dvwa
# Clone the DVWA source for static scanning
git clone https://github.com/digininja/DVWA.git
semgrep --config auto --json DVWA/ > dvwa-raw.json
```

DVWA has known vulnerabilities with known locations. Count raw findings. Compute raw precision against DVWA's labeled vulnerability set (the intentional vulnerabilities are in `vulnerabilities/`).

### 2.2 Implement LLM triage

```python
import json, os

async def triage_finding(raw: RawFinding, enclosing_fn: str) -> TriagedFinding:
    prompt = f"""
    You are a security triage analyst. Evaluate this finding.

    Rule: {raw.rule_id} ({raw.cwe})
    Location: {raw.file}:{raw.start_line}
    Matched code:
    ```
    {raw.snippet}
    ```
    Enclosing function:
    ```
    {enclosing_fn}
    ```

    Questions (answer each WITH EVIDENCE):
    1. Is the tainted input reachable from untrusted data in THIS code path?
    2. Is there a sanitizer, validator, or framework control between source and sink?
    3. True positive or false positive?
    4. Confidence: high / medium / low?

    Respond as JSON: {{"reachable": bool, "sanitized": bool, "verdict": "tp"|"fp", "confidence": "high"|"medium"|"low", "reasoning": str}}
    """
    result = await llm_complete(prompt)
    return TriagedFinding(**json.loads(result), original=raw)
```

### 2.3 Measure precision before and after triage

1. Compute raw precision: of raw findings, how many match DVWA's labeled vulnerabilities.
2. Run LLM triage on all raw findings.
3. Compute triaged precision: of HIGH-confidence findings only, how many are real.

### 2.4 Implement the feedback loop

```python
def record_feedback(finding_id: str, verdict: str, reasoning: str, store):
    """Human verdict stored for future few-shot retrieval."""
    store.add(
        documents=[f"{verdict}: {reasoning}"],
        metadatas=[{"finding_id": finding_id, "verdict": verdict}],
        ids=[finding_id]
    )

def build_triage_prompt_with_feedback(raw: RawFinding, store) -> str:
    """Retrieve past human verdicts on similar findings as few-shot examples."""
    similar = store.query(query_texts=[raw.snippet], n_results=3)
    examples = "\n".join(similar["documents"][0]) if similar["documents"][0] else "none"
    return f"Prior judgements on similar findings:\n{examples}\n\nNow evaluate: ..."
```

### Deliverable
- [ ] Raw Semgrep precision measured against DVWA
- [ ] LLM triage implemented; high-confidence precision measured
- [ ] Improvement documented (raw precision → triaged precision)
- [ ] Feedback loop records human verdicts to the vector store

---

## Phase 3 — Semantic Codebase Memory (30 min)

### 3.1 Build the function-level index

```python
import ast, hashlib, chromadb

class CodeUnit(BaseModel):
    id: str
    file: str
    symbol: str
    signature: str
    source: str
    callers: list[str] = Field(default_factory=list)
    callees: list[str] = Field(default_factory=list)

def extract_functions(file_path: str) -> list[CodeUnit]:
    """Parse a Python file to CodeUnits via AST."""
    with open(file_path) as f:
        tree = ast.parse(f.read())
    units = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            source = ast.get_source_segment(open(file_path).read(), node)
            unit_id = hashlib.sha256(f"{file_path}:{node.name}:{node.lineno}".encode()).hexdigest()[:16]
            units.append(CodeUnit(
                id=unit_id, file=file_path, symbol=node.name,
                signature=f"{node.name}({', '.join(a.arg for a in node.args.args)})",
                source=source, callers=[], callees=[]
            ))
    return units

def build_call_graph(units: list[CodeUnit]) -> list[CodeUnit]:
    """Populate callers/callees by scanning source for call references."""
    # For each unit, find which other units' symbols it calls, and vice versa.
    ...
```

### 3.2 Implement the retrieval query

```python
def retrieve_context(finding: RawFinding, index, budget_tokens: int = 6000) -> str:
    """For a finding, retrieve enclosing fn + callers + callees + sanitizers, within budget."""
    # 1. Find the enclosing function
    enclosing = index.find_enclosing(finding.file, finding.start_line)
    context_parts = [enclosing]
    # 2. Direct callers (source) and callees (sink)
    context_parts += index.get_callers(enclosing.id)
    context_parts += index.get_callees(enclosing.id)
    # 3. Sanitizer-like functions (semantic search)
    context_parts += index.semantic_search("input validation sanitizer escape", k=3)
    # 4. Truncate to budget (prioritize callers/callees/sanitizers over transitive)
    return truncate_to_budget(context_parts, budget_tokens)
```

### 3.3 Verify cross-file retrieval

1. Create a 3-file taint flow: `handler.py` (source) → `repository.py` (sink) → `validators.py` (sanitizer).
2. Place a finding in `repository.py`.
3. Verify retrieval pulls `handler.py` (caller), `validators.py` (sanitizer), enabling the triage model to judge it a false positive.

### Deliverable
- [ ] Function-level index built from AST extraction
- [ ] Call graph populated (callers/callees)
- [ ] Retrieval query returns cross-file context within the token budget
- [ ] Cross-file false positive correctly resolvable with retrieved context

---

## Phase 4 — Autofix Loop with Verification Gates (30 min)

### 4.1 Implement patch generation

```python
async def generate_patch(finding: TriagedFinding, context: str) -> str:
    prompt = f"""
    Resolve this security finding with a minimal, targeted patch.

    Finding: {finding.original.cwe} — {finding.original.cwe_name}
    Location: {finding.original.file}:{finding.original.start_line}
    Code:
    ```
    {finding.original.snippet}
    ```
    Cross-file context:
    {context}

    Constraints:
    - Minimal diff. Change only what is necessary to resolve the finding.
    - Do not refactor, rename, or reformat unrelated code.
    - Preserve all existing behavior and test coverage.
    - Output a unified diff.
    """
    return await llm_complete(prompt)
```

### 4.2 Implement the verification gates

```python
def apply_and_verify(patch_diff: str, repo_path: str, finding: RawFinding) -> dict:
    """Apply the patch and run all automated gates."""
    # Gate 1: apply patch
    apply_result = subprocess.run(["git", "apply"], input=patch_diff, text=True, cwd=repo_path)
    if apply_result.returncode != 0:
        return {"gate": "apply", "passed": False}

    # Gate 2: linter clean
    if not run_linter(repo_path):
        return {"gate": "linter", "passed": False}

    # Gate 3: Semgrep re-scan — finding resolved, no new findings
    new_findings = run_semgrep_on_diff([finding.file])
    still_present = any(f.start_line == finding.start_line and f.rule_id == finding.rule_id for f in new_findings)
    if still_present or len(new_findings) > 0:
        return {"gate": "semgrep_rescan", "passed": False, "still_present": still_present, "new": len(new_findings)}

    # Gate 4: test suite passes
    if not run_tests(repo_path):
        return {"gate": "tests", "passed": False}

    return {"gate": "all", "passed": True}
```

### 4.3 Open a draft PR (simulated)

```python
def open_draft_pr(patch_diff: str, finding: TriagedFinding, verification: dict):
    """Simulate opening a draft PR with the patch, finding, and verification evidence."""
    pr_body = f"""
    ## Automated Security Fix

    **Finding**: {finding.original.cwe} — {finding.original.cwe_name}
    **Location**: {finding.original.file}:{finding.original.start_line}
    **Confidence**: {finding.confidence}

    **Verification**:
    - Linter: PASS
    - Semgrep re-scan: PASS (finding resolved, no new findings)
    - Tests: PASS

    ---
    Human review required before merge. This PR is a DRAFT.
    """
    print(f"[DRAFT PR] {pr_body}")
    print(f"[PATCH]\n{patch_diff}")
```

### Deliverable
- [ ] Patch generation produces a minimal targeted diff
- [ ] All 4 automated gates implemented (apply, linter, semgrep, tests)
- [ ] A patch that fails any gate is discarded/revised
- [ ] Draft PR (simulated) opened only for fully-verified patches
- [ ] Approval gate documented as the required final step (no auto-merge)

---

## Stretch goals

1. **Add CodeQL as a third deterministic layer** between Semgrep and the LLM. Cross-dedup Semgrep and CodeQL findings via `dedup_key` before triage.
2. **Measure F1** (not just precision and recall) across a sweep of confidence thresholds. Plot the precision-recall curve and pick the operating point.
3. **Patch quality scoring**: implement a heuristic (diff size, style regressions, coverage delta) and hold back patches below the quality threshold even if all gates pass.