Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 43f8ad1

Browse files
samcodesignsamcodesign
committed
Refactor base provider and implement model listing
- Updated BaseProvider to include a model_name parameter and a static method for listing available models. - Enhanced GeminiProvider to utilize the new model_name feature and added a method to list available models from the Gemini API. - Updated requirements.txt to include the questionary library for CLI interactions. - Refactored run_benchmark.py to integrate a new CLI module for better task management and user interaction. - Removed outdated task directories and their associated files for hardcoded keys and command injection vulnerabilities. - Added new implementations for tasks with improved security practices, including environment variable usage and subprocess handling. - Created a CLI for SecureDev-Bench to facilitate task discovery and execution. - Implemented reporting functionality to generate JSON and Markdown reports of benchmark results.
1 parent 35af6d7 commit 43f8ad1

25 files changed

Lines changed: 345 additions & 240 deletions

File tree

‎agent.py‎

Lines changed: 11 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -4,64 +4,49 @@
44
from pathlib import Path
55
from dotenv import load_dotenv
66

7-
# Load environment variables from .env file
87
load_dotenv()
98

10-
# Import our new provider classes
119
from providers.gemini_provider import GeminiProvider
12-
# from providers.openai_provider import OpenAIProvider # You could add this later
10+
# from providers.openai_provider import OpenAIProvider
1311

14-
# This dictionary makes the agent easily extensible.
15-
# To add a new model, just add it here.
1612
SUPPORTED_PROVIDERS = {
1713
"gemini": GeminiProvider,
1814
# "openai": OpenAIProvider,
1915
}
2016

2117
def get_api_key(provider_name: str) -> str:
22-
"""Gets the API key from environment variables based on the provider name."""
2318
env_var_name = f"{provider_name.upper()}_API_KEY"
2419
api_key = os.environ.get(env_var_name)
2520
if not api_key:
26-
raise ValueError(f"Please set the {env_var_name} environment variable.")
21+
raise ValueError(f"API key not found. Please set {env_var_name} in your .env file.")
2722
return api_key
2823

2924
def main():
30-
"""
31-
Main function to parse arguments, select a provider, and fix code.
32-
"""
3325
parser = argparse.ArgumentParser(description="AI Security Agent")
3426
parser.add_argument("file_path", type=str, help="Path to the Python file to fix.")
3527
parser.add_argument(
36-
"--provider",
37-
type=str,
38-
default="gemini",
39-
choices=SUPPORTED_PROVIDERS.keys(),
40-
help="The AI provider to use."
28+
"--provider", type=str, default="gemini", choices=SUPPORTED_PROVIDERS.keys()
29+
)
30+
# --- NEW: Add an argument for the specific model name ---
31+
parser.add_argument(
32+
"--model", type=str, default=None, help="The specific model to use (e.g., gemini-1.5-pro)."
4133
)
4234
args = parser.parse_args()
4335

44-
print(f"[Agent]: Using provider: {args.provider}")
36+
print(f"[Agent]: Using provider: {args.provider}, Model: {args.model or 'default'}")
4537

4638
try:
47-
# 1. Select the provider class from our dictionary
4839
provider_class = SUPPORTED_PROVIDERS[args.provider]
49-
50-
# 2. Get the correct API key for that provider
5140
api_key = get_api_key(args.provider)
5241

53-
# 3. Instantiate the provider
54-
provider = provider_class(api_key=api_key)
42+
# --- NEW: Pass the model name to the provider's constructor ---
43+
provider = provider_class(api_key=api_key, model_name=args.model)
5544

56-
# 4. Read the vulnerable code
5745
file_path = Path(args.file_path)
5846
vulnerable_code = file_path.read_text()
59-
60-
# 5. Ask the provider to fix the code
6147
corrected_code = provider.fix_code(vulnerable_code)
6248

63-
# 6. Write the fix back to the file
64-
print(f"[Agent]: Overwriting file with corrected code from {args.provider}...")
49+
print(f"[Agent]: Overwriting file with corrected code...")
6550
file_path.write_text(corrected_code)
6651
print("[Agent]: Fix complete.")
6752

‎providers/base_provider.py‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,24 @@
11
from abc import ABC, abstractmethod
22

33
class BaseProvider(ABC):
4-
"""
5-
Abstract base class for all AI model providers.
6-
It ensures that every provider has a 'fix_code' method.
7-
"""
8-
def __init__(self, api_key: str):
4+
"""Abstract base class for all AI model providers."""
5+
6+
def __init__(self, api_key: str, model_name: str = None):
97
if not api_key:
108
raise ValueError(f"{self.__class__.__name__} requires an API key.")
119
self.api_key = api_key
10+
self.model_name = model_name
1211

1312
@abstractmethod
1413
def fix_code(self, vulnerable_code: str) -> str:
14+
"""Takes a string of vulnerable code and returns the fixed code."""
15+
pass
16+
17+
@staticmethod
18+
@abstractmethod
19+
def list_models(api_key: str) -> list[str]:
1520
"""
16-
Takes a string of vulnerable code and returns the fixed code.
17-
This method must be implemented by all subclasses.
21+
A static method that connects to the provider's API and returns a list
22+
of available model names that are suitable for code generation.
1823
"""
1924
pass

‎providers/gemini_provider.py‎

Lines changed: 33 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import google.generativeai as genai
22
from .base_provider import BaseProvider
33

4-
# This is the core instruction set for the AI model.
54
SYSTEM_PROMPT = """
65
You are an expert security programmer. Your task is to analyze a given Python code file for security vulnerabilities and fix them.
76
Do not explain the vulnerability. Do not add any comments or introductory text.
@@ -11,29 +10,45 @@
1110
class GeminiProvider(BaseProvider):
1211
"""Provider for Google's Gemini models."""
1312

14-
def __init__(self, api_key: str):
15-
super().__init__(api_key)
13+
def __init__(self, api_key: str, model_name: str = None):
14+
super().__init__(api_key, model_name)
1615
genai.configure(api_key=self.api_key)
17-
self.model = genai.GenerativeModel('gemini-2.5-flash') # Or another Gemini model
16+
17+
# Use the provided model name, or a sensible default.
18+
self.model_name = model_name or 'gemini-1.5-flash'
19+
20+
print(f"[Agent/Gemini]: Initializing model: {self.model_name}")
21+
self.model = genai.GenerativeModel(self.model_name)
1822

1923
def fix_code(self, vulnerable_code: str) -> str:
20-
print("[Agent/Gemini]: Sending code to Gemini for analysis...")
24+
# This method remains the same
25+
print(f"[Agent/Gemini]: Sending code to {self.model_name} for analysis...")
2126
try:
22-
# Gemini uses a specific format for system prompts + user input
2327
full_prompt = f"{SYSTEM_PROMPT}\n\n--- VULNERABLE CODE ---\n{vulnerable_code}"
24-
2528
response = self.model.generate_content(full_prompt)
26-
27-
# The response text contains the code
2829
corrected_code = response.text
29-
30-
# Clean up markdown formatting, which models often add
31-
if corrected_code.startswith("```python"):
32-
corrected_code = corrected_code[len("```python\n"):-len("```")]
33-
34-
return corrected_code
35-
30+
cleaned_code = corrected_code.strip().strip('`')
31+
if cleaned_code.startswith("python"):
32+
cleaned_code = cleaned_code[len("python\n"):].strip()
33+
return cleaned_code
3634
except Exception as e:
3735
print(f"[Agent/Gemini]: ERROR - Failed to get response from Gemini: {e}")
38-
# Return original code on failure to avoid breaking the file
39-
return vulnerable_code
36+
return e
37+
38+
@staticmethod
39+
def list_models(api_key: str) -> list[str]:
40+
"""Connects to the Gemini API to discover available models."""
41+
print("[Harness]: Discovering available Gemini models...")
42+
try:
43+
genai.configure(api_key=api_key)
44+
available_models = []
45+
for m in genai.list_models():
46+
# We only want models that support content generation
47+
if 'generateContent' in m.supported_generation_methods:
48+
# The API returns 'models/gemini-1.5-pro', we just want 'gemini-1.5-pro'
49+
model_name = m.name.split('/')[-1]
50+
available_models.append(model_name)
51+
return available_models
52+
except Exception as e:
53+
print(f"Warning: Could not fetch Gemini models. Error: {e}")
54+
return []

‎requirements.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
# Dependencies for running the benchmark harness and agent
22
python-dotenv
3-
google-generativeai
3+
google-generativeai
4+
questionary

‎run_benchmark.py‎

Lines changed: 16 additions & 142 deletions
Original file line numberDiff line numberDiff line change
@@ -1,145 +1,19 @@
1-
import subprocess
2-
import json
3-
import time
4-
import os
5-
import shutil
6-
from pathlib import Path
1+
#!/usr/bin/env python
72
import sys
8-
9-
def simulate_agent_fix(task_dir: Path):
10-
"""Simulates a successful agent fix."""
11-
print("[Harness]: Simulating agent fix...")
12-
app_path = task_dir / "app.py"
13-
solution_path = task_dir / "solution.py"
14-
if not solution_path.exists():
15-
print("[Harness]: WARNING - No solution.py found.")
16-
return
17-
shutil.copy(solution_path, app_path)
18-
print("[Harness]: Agent simulation complete.")
19-
20-
def parse_pytest_report(report_path: Path) -> dict:
21-
"""Parses the JSON report from pytest."""
22-
results = {
23-
"security_tests": {"passed": 0, "failed": 0, "total": 0},
24-
"functional_tests": {"passed": 0, "failed": 0, "total": 0},
25-
"overall_passed": True
26-
}
27-
if not report_path.exists():
28-
results["overall_passed"] = False
29-
return results
30-
with open(report_path, 'r') as f:
31-
report = json.load(f)
32-
if report.get("summary", {}).get("failed", 0) > 0:
33-
results["overall_passed"] = False
34-
for test in report.get("tests", []):
35-
category = "functional_tests" if "test_app.py" in test.get("nodeid") else "security_tests"
36-
results[category]["total"] += 1
37-
if test.get("outcome") == "passed":
38-
results[category]["passed"] += 1
39-
else:
40-
results[category]["failed"] += 1
41-
return results
42-
43-
def run_task(task_id: str):
44-
"""Runs a full evaluation for a given task."""
45-
print(f"--- Running Task: {task_id} ---")
46-
start_time = time.time()
47-
original_task_dir = Path(task_id)
48-
temp_dir = Path(f"/tmp/{task_id}_{int(start_time)}").resolve()
49-
final_result, scorecard = "FAILED", {}
50-
image_name = f"securedev-bench-{task_id.lower()}"
51-
container_name = f"{image_name}-test"
52-
53-
try:
54-
if not original_task_dir.exists():
55-
print(f"[Harness]: ERROR - Task directory not found: {original_task_dir}")
56-
return
57-
shutil.copytree(original_task_dir, temp_dir)
58-
print(f"[Harness]: Created isolated environment at {temp_dir}")
59-
60-
# 2. Run the real agent via a command-line call
61-
print("[Harness]: Running real agent...")
62-
agent_path = Path(__file__).parent / "agent.py"
63-
target_file_path = temp_dir / "app.py"
64-
provider_to_test = "gemini" # You can change this to "openai" later
65-
66-
try:
67-
# This builds the command: python agent.py /path/to/app.py --provider gemini
68-
subprocess.run(
69-
[
70-
sys.executable, str(agent_path),
71-
str(target_file_path),
72-
"--provider", provider_to_test
73-
],
74-
check=True, capture_output=True, text=True
75-
)
76-
except subprocess.CalledProcessError as e:
77-
# This block catches errors if the agent.py script itself fails
78-
print("[Harness]: The agent script returned an error!")
79-
print("--- Agent STDOUT ---")
80-
print(e.stdout)
81-
print("--- Agent STDERR ---")
82-
print(e.stderr)
83-
raise e # Stop the task if the agent fails
84-
85-
print("[Harness]: Building Docker container...")
86-
subprocess.run(["docker", "build", "-t", image_name, "."], cwd=temp_dir, check=True, capture_output=True)
87-
88-
print("[Harness]: Running tests inside container...")
89-
run_result = subprocess.run(
90-
[
91-
"docker", "run",
92-
"-e", "API_KEY=DUMMY_KEY_FOR_TESTING", # <--- ADD THIS LINE
93-
f"--name={container_name}", image_name
94-
],
95-
cwd=temp_dir, capture_output=True, text=True
96-
)
97-
print(f"\n--- Container logs for {container_name}: ---")
98-
print(run_result.stdout)
99-
if run_result.stderr:
100-
print("--- Stderr ---")
101-
print(run_result.stderr)
102-
print("--- End of logs ---")
103-
104-
if run_result.returncode != 0:
105-
raise subprocess.CalledProcessError(run_result.returncode, run_result.args, run_result.stdout, run_result.stderr)
106-
107-
report_path_in_container = f"{container_name}:/usr/src/app/report.json"
108-
report_path_on_host = temp_dir / "report.json"
109-
subprocess.run(["docker", "cp", report_path_in_container, str(report_path_on_host)], check=True, capture_output=True)
110-
111-
scorecard = parse_pytest_report(report_path_on_host)
112-
if scorecard.get("overall_passed"):
113-
final_result = "SUCCESS"
114-
115-
except subprocess.CalledProcessError as e:
116-
print(f"\n--- [Harness]: A critical command failed! ---")
117-
final_result = "HARNESS_FAILURE"
118-
finally:
119-
print(f"[Harness]: Cleaning up container '{container_name}'...")
120-
subprocess.run(["docker", "rm", container_name], capture_output=True, check=False)
121-
if temp_dir.exists():
122-
shutil.rmtree(temp_dir)
123-
print("[Harness]: Cleaned up temporary environment.")
124-
125-
duration = time.time() - start_time
126-
print(f"\n--- Task {task_id} Finished ---")
127-
print(f"FINAL RESULT: {final_result}")
128-
print(f"Duration: {duration:.2f}s")
129-
print("\nSCORECARD:")
130-
if scorecard:
131-
sec = scorecard.get('security_tests', {})
132-
func = scorecard.get('functional_tests', {})
133-
print(f" - Correctness (Security): {sec.get('passed', 0)}/{sec.get('total', 0)} PASSED")
134-
print(f" - Functionality (Tests): {func.get('passed', 0)}/{func.get('total', 0)} PASSED")
135-
else:
136-
print(" - No results generated.")
137-
print("--------------------------------\n")
3+
from dotenv import load_dotenv
4+
from securedev_bench.cli import main
1385

1396
if __name__ == "__main__":
140-
tasks = [d for d in os.listdir('.') if os.path.isdir(d) and d.startswith('task-')]
141-
if not tasks:
142-
print("No task directories found. Exiting.")
143-
else:
144-
for task in sorted(tasks):
145-
run_task(task)
7+
# Load environment variables from .env file at the very start
8+
load_dotenv()
9+
10+
# Call the main function from our CLI module
11+
try:
12+
main()
13+
except (KeyboardInterrupt):
14+
# Handle user cancellation gracefully
15+
print("\nBenchmark run cancelled by user.")
16+
sys.exit(0)
17+
except Exception as e:
18+
print(f"\nAn unexpected error occurred: {e}")
19+
sys.exit(1)

0 commit comments

Comments
 (0)