Skip to content

Commit 1ef6156

Browse files
Feat/v2 agent scanner pr diff annotation (#57)
* feat: add V2 Pydantic-AI agent scanner with PR diff annotation - File-by-file scanning via Pydantic-AI Agent with structured FileScanResult output - GitHub integration for inline PR review comments - [CHANGED] line markers when scanning PR diffs or local git changes - Diff parsing helpers (get_pr_changed_line_numbers, get_local_changed_line_numbers) - Agent prompt updated to prioritise [CHANGED] lines while retaining full file context - Eval harness and 5 fixture files for regression testing * chore: add CLAUDE.md, update deps and gitignore * chore: switch default entrypoint to V2 runner * chore: fix pylint and isort violations * fix: restore exception message in scan error logs * fix: use os.path.join in test assertions for Windows compatibility
1 parent a461cd1 commit 1ef6156

17 files changed

Lines changed: 2011 additions & 6 deletions

.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,9 @@ celerybeat.pid
121121
# SageMath parsed files
122122
*.sage.py
123123

124+
# Local idea tracking
125+
NEW_IDEAS.md
126+
124127
# Environments
125128
.env
126129
.venv

CLAUDE.md

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
# CLAUDE.md
2+
3+
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
4+
5+
## Commands
6+
7+
### Install dependencies
8+
```bash
9+
pip install -r requirements.txt
10+
pip install -r requirements_ci.txt # dev/linting tools
11+
```
12+
13+
### Run tests
14+
```bash
15+
# All tests
16+
python -m unittest discover -s core_tests -p "*.py"
17+
18+
# Single test file
19+
python -m unittest core_tests.runner_test
20+
```
21+
22+
### Lint and format
23+
```bash
24+
pylint core/
25+
isort --profile=black --check-only core/ core_tests/
26+
isort --profile=black core/ core_tests/ # to fix imports
27+
```
28+
29+
### Build
30+
```bash
31+
python -m build
32+
```
33+
34+
### Run evals
35+
36+
Evals verify that the V2 agent scanner still catches expected vulnerabilities across 5 fixture files (`evals/fixtures/`). Run them after any change to `core/agent.py`, `core/code_scanner/agent_scanner.py`, or the system prompt.
37+
38+
```bash
39+
# Standard model (gpt-4o-mini) — all fixtures should pass
40+
set -a && source .env && set +a
41+
python3 evals/run_evals.py --provider openai --model gpt-4o-mini
42+
43+
# Advanced model (gpt-4o) — stricter; also requires YAML deserialization,
44+
# race condition, JWT algorithm confusion, and timing attack findings
45+
python3 evals/run_evals.py --provider openai --model gpt-4o
46+
47+
# Single fixture only
48+
python3 evals/run_evals.py --provider openai --model gpt-4o-mini --fixture auth_service.py
49+
```
50+
51+
Findings are split into two tiers in `evals/expected_findings.json`:
52+
- **standard** — required from any model (SQL injection, XSS, path traversal, pickle, hardcoded secrets, etc.)
53+
- **advanced** — only required when running gpt-4o (YAML code execution, race conditions, JWT algorithm confusion, timing attacks)
54+
55+
Exit code 0 = all fixtures at or above the 80% threshold. Exit code 1 = one or more failed.
56+
### Run the CLI locally
57+
```bash
58+
# V1 runner (sends all files to AI at once)
59+
python3 -m core.runner --provider openai
60+
61+
# V2 runner (file-by-file via Pydantic-AI agent)
62+
python3 -m core.runner_v2 --provider openai
63+
```
64+
65+
## Architecture
66+
67+
CodeScanAI is a CLI tool that scans codebases for security vulnerabilities using AI models.
68+
69+
### Two parallel codepaths
70+
71+
There are two scanner implementations that share the same CLI argument surface (`core/utils/argument_parser.py`):
72+
73+
1. **V1 (`core/runner.py``core/code_scanner/code_scanner.py`)**: Aggregates all file content into a single code summary, sends it to the AI provider in one call, returns a markdown string. Uses the provider abstraction layer.
74+
75+
2. **V2 (`core/runner_v2.py``core/code_scanner/agent_scanner.py`)**: Iterates file-by-file, runs a Pydantic-AI `Agent` synchronously on each, and streams structured `FileScanResult` output to stdout. Also supports posting inline PR review comments via `GithubIntegration`. This is the more feature-rich path.
76+
77+
The active entrypoint is `core.runner_v2:main` (V2), set in `pyproject.toml` under `[project.scripts]`.
78+
79+
### Provider abstraction (V1 only)
80+
81+
`core/providers/base_ai_provider.py` defines the `BaseAIProvider` interface with a single `scan_code(code_summary)` method. Concrete implementations:
82+
- `OpenAIProvider` — uses `openai` SDK
83+
- `GoogleGeminiAIProvider` — uses `google-generativeai`
84+
- `CustomAIProvider` — HTTP requests to a self-hosted server (Ollama, etc.)
85+
86+
`core/utils/provider_creator.py` maps CLI `--provider` values to provider classes.
87+
88+
### Pydantic-AI agent (V2 only)
89+
90+
`core/agent.py` defines structured output types (`Vulnerability`, `FileScanResult`) and factory functions. It also holds pre-configured system prompts for different scan modes: `SECURITY_AGENT_PROMPT`, `PERFORMANCE_AGENT_PROMPT`, `CLEAN_CODE_AGENT_PROMPT`. Custom providers route through the OpenAI-compatible interface via `OPENAI_BASE_URL`.
91+
92+
### GitHub integration
93+
94+
`core/utils/github_integration.py` (`GithubIntegration`) is used only in V2. It posts inline PR review comments using PyGithub. Falls back to a regular issue comment if the line isn't in the PR diff.
95+
96+
`core/utils/file_extractor.py` handles both local git-diff file discovery and PR file listing via the GitHub API, shared by both V1 and V2.
97+
98+
### Scan modes
99+
100+
Both runners support three modes driven by CLI args:
101+
- **Full scan** (default): walks `--directory` and scans all files
102+
- **Changes only** (`--changes_only`): scans files changed in local git repo
103+
- **PR scan** (`--repo` + `--pr_number` + `--github_token`): fetches changed files from a GitHub PR

core/agent.py

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
"""
2+
Defines the structured output types, agent factory, and pre-configured system prompts
3+
used by the V2 Pydantic-AI scanner.
4+
"""
5+
6+
from typing import Optional, Type
7+
8+
from pydantic import BaseModel, Field
9+
from pydantic_ai import Agent
10+
11+
12+
class Vulnerability(BaseModel):
13+
"""Represents a single security vulnerability found in a file."""
14+
15+
line_number: Optional[int] = Field(
16+
default=None,
17+
description=(
18+
"The exact line number where the issue is found inside the file. "
19+
"Omit if the issue is architectural or spans multiple lines."
20+
),
21+
)
22+
description: str = Field(
23+
description="A detailed description of the issue and why it is a security risk."
24+
)
25+
remediation: str = Field(
26+
description="Actionable suggestion or code snippet on how to fix this specific vulnerability."
27+
)
28+
severity: str = Field(description="Severity measure (Low, Medium, High, Critical)")
29+
vulnerability_type: str = Field(
30+
description="The category of issue (e.g. SQL Injection, Big-O inefficiency, etc.)"
31+
)
32+
33+
34+
class FileScanResult(BaseModel):
35+
"""Structured result returned by the agent for a single scanned file."""
36+
37+
vulnerabilities: list[Vulnerability] = Field(
38+
description="List of issues found in the file. Empty if zero issues are found."
39+
)
40+
41+
42+
def get_pydantic_ai_model(provider: str, model: Optional[str]) -> str:
43+
"""Map a CLI provider name and optional model string to a Pydantic-AI model identifier."""
44+
if provider == "openai":
45+
return f"openai:{model or 'gpt-4o-mini'}"
46+
if provider == "gemini":
47+
return f"gemini:{model or 'gemini-1.5-flash'}"
48+
if provider == "custom":
49+
# Falls back to the OpenAI-compatible interface via OPENAI_BASE_URL
50+
return f"openai:{model or 'custom-model'}"
51+
return "openai:gpt-4o-mini"
52+
53+
54+
def create_agent(
55+
model_str: str, system_prompt: str, result_type: Type[BaseModel] = FileScanResult
56+
) -> Agent:
57+
"""
58+
Creates and returns a Pydantic-AI Agent configured for a custom, laser-focused task.
59+
By passing different `system_prompt` and `result_type` schemas, you can deploy
60+
multiple types of agents.
61+
"""
62+
return Agent(
63+
model_str,
64+
result_type=result_type,
65+
system_prompt=system_prompt,
66+
)
67+
68+
69+
# --- Define pre-configured Agent Prompts for laser-focused tasks ---
70+
71+
SECURITY_AGENT_PROMPT = (
72+
"You are an expert in software security analysis, adept at identifying and explaining "
73+
"potential vulnerabilities in code. "
74+
"You will be given complete code snippets from various applications. "
75+
"EVERY line of the source code is prefixed with its exact line number "
76+
"(e.g. `14: def foo():`). "
77+
"Your task is to analyze the provided code, pinpoint potential security risks, "
78+
"and offer clear suggestions for enhancing the application's security posture. "
79+
"Focus on the critical issues that could impact the overall security of the application. "
80+
"You MUST be exhaustive. Carefully audit the entire script from top to bottom "
81+
"and return EVERY vulnerability you find. Do not stop at the first issue. "
82+
"If any are found, use the explicitly provided line numbers to pinpoint the defect "
83+
"where possible. For architectural or multi-line issues, you may omit the line number. "
84+
"Also, strictly provide an actionable `remediation` that makes suggestions on how to "
85+
"rewrite or fix the code securely. "
86+
"If no vulnerabilities are found, return an empty list. "
87+
"When scanning a pull request or diff, some lines will be marked with `[CHANGED]` "
88+
"after the line number (e.g. `14: [CHANGED] def foo():`). "
89+
"These lines are newly added or modified in the change under review. "
90+
"Prioritise your analysis on `[CHANGED]` lines, but use the full file context — "
91+
"imports, surrounding functions, class definitions, and data flow — to assess "
92+
"whether those changes introduce or worsen a vulnerability."
93+
)
94+
95+
PERFORMANCE_AGENT_PROMPT = (
96+
"You are a Senior Staff Software Engineer laser-focused on performance optimization. "
97+
"Analyze the following code for memory leaks, O(N^2) bottlenecks, or CPU inefficiencies. "
98+
"Pinpoint exact line numbers and return a list of performance issues. "
99+
"If none are found, return an empty list."
100+
)
101+
102+
CLEAN_CODE_AGENT_PROMPT = (
103+
"You are an expert in code refactoring and Clean Code methodologies. "
104+
"Analyze the code for anti-patterns, confusing variable names, massive functions, "
105+
"or high cyclomatic complexity. "
106+
"Pinpoint exact line numbers and return a list of maintainability issues. "
107+
"If the code is perfectly clean, return an empty list."
108+
)

core/code_scanner/agent_scanner.py

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
"""
2+
V2 agent-based scanner. Scans source files one at a time using a Pydantic-AI Agent
3+
and streams structured FileScanResult output to stdout.
4+
"""
5+
6+
import logging
7+
import os
8+
9+
from core.agent import (
10+
SECURITY_AGENT_PROMPT,
11+
FileScanResult,
12+
create_agent,
13+
get_pydantic_ai_model,
14+
)
15+
from core.utils.file_extractor import (
16+
get_changed_files_in_pr,
17+
get_changed_files_in_repo,
18+
get_local_changed_line_numbers,
19+
get_pr_changed_line_numbers,
20+
)
21+
from core.utils.github_integration import GithubIntegration
22+
23+
logging.basicConfig(
24+
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
25+
)
26+
27+
28+
class AgentScanner:
29+
"""
30+
Scans source code via the Pydantic AI agent, file by file.
31+
Optionally posts inline review comments to GitHub PRs.
32+
"""
33+
34+
def __init__(self, args) -> None:
35+
self.args = args
36+
model_str = get_pydantic_ai_model(args.provider, args.model)
37+
38+
# Set OPENAI_BASE_URL for custom providers so Pydantic-AI targets the correct backend.
39+
if args.provider == "custom" and args.host:
40+
host_url = f"{args.host}:{args.port}" if args.port else args.host
41+
if args.endpoint:
42+
host_url += args.endpoint
43+
os.environ["OPENAI_BASE_URL"] = host_url
44+
if args.token:
45+
os.environ["OPENAI_API_KEY"] = args.token
46+
47+
self.agent = create_agent(
48+
model_str=model_str,
49+
system_prompt=SECURITY_AGENT_PROMPT,
50+
result_type=FileScanResult,
51+
)
52+
self.github_integration = (
53+
GithubIntegration(args)
54+
if args.repo and args.pr_number and args.github_token
55+
else None
56+
)
57+
58+
def scan(self):
59+
"""
60+
Scans the code by identifying files based on PR context or local directory
61+
and iterates through them using the Pydantic AI agent.
62+
"""
63+
if self.args.changes_only or (self.args.repo and self.args.pr_number):
64+
return self._scan_changes()
65+
return self._scan_files()
66+
67+
def _scan_changes(self):
68+
try:
69+
if self.args.repo and self.args.pr_number:
70+
changed_files = get_changed_files_in_pr(
71+
self.args.repo, self.args.pr_number, self.args.github_token
72+
)
73+
changed_line_map = get_pr_changed_line_numbers(
74+
self.args.repo, self.args.pr_number, self.args.github_token
75+
)
76+
else:
77+
changed_files = get_changed_files_in_repo(self.args.directory)
78+
changed_line_map = None
79+
except ValueError as e:
80+
logging.error(e)
81+
return
82+
83+
if not changed_files:
84+
logging.info("No changes detected.")
85+
return
86+
87+
for filename in changed_files:
88+
filepath = os.path.join(self.args.directory, filename)
89+
if changed_line_map is not None:
90+
changed_lines = changed_line_map.get(filename, set())
91+
else:
92+
changed_lines = get_local_changed_line_numbers(self.args.directory, filename)
93+
self._scan_single_file(filepath, display_name=filename, changed_lines=changed_lines)
94+
95+
def _scan_files(self):
96+
file_paths = []
97+
for root, _, files in os.walk(self.args.directory):
98+
for file in files:
99+
file_paths.append(os.path.join(root, file))
100+
101+
for filepath in file_paths:
102+
self._scan_single_file(
103+
filepath, display_name=os.path.relpath(filepath, self.args.directory)
104+
)
105+
106+
def _scan_single_file(self, file_path: str, display_name: str, changed_lines: set = None):
107+
"""Scan a single file and print any vulnerabilities found."""
108+
if not os.path.isfile(file_path):
109+
logging.warning("Skipping %s: Not a valid file or not found locally.", file_path)
110+
return
111+
112+
try:
113+
with open(file_path, "r", encoding="utf-8") as f:
114+
content = f.read()
115+
except Exception as e: # pylint: disable=broad-exception-caught
116+
logging.warning("Skipping %s: %s", file_path, e)
117+
return
118+
119+
if not content.strip():
120+
return
121+
122+
logging.info("Scanning file: %s ...", display_name)
123+
124+
def _format_line(idx, line):
125+
lineno = idx + 1
126+
if changed_lines and lineno in changed_lines:
127+
return f"{lineno}: [CHANGED] {line}"
128+
return f"{lineno}: {line}"
129+
130+
numbered_content = "\n".join([_format_line(idx, line) for idx, line in enumerate(content.splitlines())])
131+
132+
try:
133+
result = self.agent.run_sync(f"File: {display_name}\n\n{numbered_content}")
134+
scan_result = result.data
135+
136+
if scan_result.vulnerabilities:
137+
print(f"\n--- Vulnerabilities found in {display_name} ---")
138+
md_output = ""
139+
for vuln in scan_result.vulnerabilities:
140+
line_info = f"Line {vuln.line_number}: " if vuln.line_number else ""
141+
md_output += f" - **{line_info}[{vuln.severity}] {vuln.vulnerability_type}**\n"
142+
md_output += f" - **Issue**: {vuln.description}\n"
143+
md_output += f" - **Fix**: {vuln.remediation}\n"
144+
print(md_output)
145+
146+
if self.github_integration:
147+
for vuln in scan_result.vulnerabilities:
148+
comment_body = (
149+
f"**[{vuln.severity.upper()} SEVERITY] {vuln.vulnerability_type}**"
150+
f"\n\n{vuln.description}"
151+
f"\n\n**Suggested Fix:**\n{vuln.remediation}"
152+
)
153+
self.github_integration.post_inline_comment(
154+
path=display_name,
155+
line=vuln.line_number,
156+
body=comment_body,
157+
)
158+
else:
159+
logging.info("No vulnerabilities found in %s.", display_name)
160+
161+
except Exception as e: # pylint: disable=broad-exception-caught
162+
logging.error("Error scanning %s: %s", display_name, e)

0 commit comments

Comments
 (0)