External/securelens-backend
1
0
Fork 0
mirror of https://github.com/Rarebuffalo/securelens-backend.git synced 2026-06-19 07:00:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

updated the model

Browse Source
This commit is contained in:
rarebuffalo
2026-04-25 20:49:17 +05:30
parent 324ebe8955
commit 139c8d982b
9 changed files with 120 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.env.example
View File

@@ -23,5 +23,5 @@ PATH_CHECK_TIMEOUT=3
DATABASE_URL=postgresql+asyncpg://securelens:securelens@localhost:5433/securelens
# AI Integration
OPENAI_API_KEY=
GEMINI_API_KEY=

2
app/config.py
View File

@@ -22,7 +22,7 @@ class Settings(BaseSettings):
jwt_algorithm: str = "HS256"
jwt_expiry_minutes: int = 1440
openai_api_key: str | None = None
gemini_api_key: str | None = None
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")

18
app/routers/__init__.py
View File

@@ -1,17 +1 @@
from .auth import router as auth
from .health import router as health
from .history import router as history
from .scan import router as scan
from .apikey import router as apikey
from .report import router as report
from .code_scan import router as code_scan
__all__ = [
"auth",
"health",
"history",
"scan",
"apikey",
"report",
"code_scan"
]
# Empty init file to allow correct module importing

33
app/routers/code_scan.py
View File

@@ -5,7 +5,7 @@ from fastapi import APIRouter, HTTPException
from typing import Dict, Any
from app.schemas.code_scan import CodeScanRequest, CodeScanResponse, CodeChatRequest, CodeChatResponse
from app.services.code_scanner.orchestrator import CodeScanOrchestrator, client
from app.services.code_scanner.orchestrator import CodeScanOrchestrator
from app.config import settings
logger = logging.getLogger(__name__)
@@ -59,32 +59,41 @@ async def analyze_codebase(request: CodeScanRequest):
@router.post("/code-scan/chat", response_model=CodeChatResponse)
async def chat_with_scan(request: CodeChatRequest):
if not settings.openai_api_key:
raise HTTPException(status_code=400, detail="AI Chat is disabled because OPENAI_API_KEY is not configured.")
if not settings.gemini_api_key:
raise HTTPException(status_code=400, detail="AI Chat is disabled because GEMINI_API_KEY is not configured.")
from google import genai
with open("/app/models.txt", "w") as f:
# Just writing a placeholder, list_models is different in new SDK
f.write("AVAILABLE MODELS: migrated to new SDK")
scan_data = scan_store.get(request.scan_id)
if not scan_data:
raise HTTPException(status_code=404, detail="Scan ID not found or expired.")
system_prompt = (
prompt = (
"You are SecureLens AI, an expert application security assistant. "
"You are helping a developer understand a security scan report for their codebase. "
f"Here is the context of the scan for the repository {scan_data.repo_url}:\n"
f"Summary: {scan_data.summary}\n"
f"Vulnerabilities: {json.dumps([v.model_dump() for v in scan_data.issues])}\n\n"
f"User Message: {request.message}\n\n"
"Answer the user's questions clearly, concisely, and professionally. Provide code fixes if requested."
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": request.message}
],
temperature=0.5,
from google import genai
from google.genai import types
client = genai.Client(api_key=settings.gemini_api_key)
response = await client.aio.models.generate_content(
model='gemini-2.5-flash',
contents=prompt,
config=types.GenerateContentConfig(
temperature=0.5,
)
)
reply = response.choices[0].message.content or "No response from AI."
reply = response.text or "No response from AI."
return CodeChatResponse(reply=reply)
except Exception as e:
logger.error(f"AI Chat Error: {str(e)}")

2
app/routers/scan.py
View File

@@ -99,7 +99,7 @@ async def scan_website(
score = calculate_score(all_issues)
layers = calculate_layer_statuses(all_issues)
if settings.openai_api_key and all_issues:
if settings.gemini_api_key and all_issues:
issues_dict_list = [i.model_dump() for i in all_issues]
ai_data = await enhance_security_issues(issues_dict_list)
enhanced_list = ai_data.get("enhanced_issues", [])

105
app/services/ai.py
View File

@@ -1,25 +1,28 @@
import json
import logging
from openai import AsyncOpenAI
import asyncio
from google import genai
from google.genai import types
from app.config import settings
logger = logging.getLogger(__name__)
api_key = settings.openai_api_key or "mock-key-for-testing"
client = AsyncOpenAI(api_key=api_key)
if settings.gemini_api_key:
# Initialize google-genai client
ai_client = genai.Client(api_key=settings.gemini_api_key)
else:
ai_client = None
async def get_gemini_model():
return 'gemini-2.5-flash'
async def enhance_security_issues(issues: list[dict]) -> dict:
"""
Takes a list of basic security issues and uses an LLM to provide:
- Contextual severity
- Natural language explanations
- Auto-generated remediation code snippets
"""
if not settings.openai_api_key:
logger.warning("OPENAI_API_KEY is not set. AI enhancements are skipped.")
if not settings.gemini_api_key:
logger.warning("GEMINI_API_KEY is not set. AI enhancements are skipped.")
return {"enhanced_issues": issues}
prompt = (
"You are a senior cybersecurity automation agent. Always respond with valid JSON.\n"
"Analyze the following security vulnerabilities:\n"
f"{json.dumps(issues, indent=2)}\n\n"
"Return a JSON object with a single key 'enhanced_issues' containing a list of objects. "
@@ -31,75 +34,69 @@ async def enhance_security_issues(issues: list[dict]) -> dict:
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a senior cybersecurity automation agent. Always respond with valid JSON."},
{"role": "user", "content": prompt}
],
response_format={"type": "json_object"},
temperature=0.2,
model_name = await get_gemini_model()
response = await ai_client.aio.models.generate_content(
model=model_name,
contents=prompt,
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0.2,
)
)
content = response.choices[0].message.content
if not content:
return {"enhanced_issues": issues, "ai_error": "Empty response"}
return json.loads(content)
if response.text:
return json.loads(response.text)
return {"enhanced_issues": issues, "ai_error": "Empty response"}
except Exception as e:
logger.error(f"AI Generation Error: {str(e)}")
return {"enhanced_issues": issues, "ai_error": str(e)}
async def chat_with_scan_context(scan_id: str, context_data: dict, user_message: str) -> str:
"""
Allows a user to ask a question about a specific scan's results.
"""
if not settings.openai_api_key:
return "AI Chat is disabled because OPENAI_API_KEY is not configured."
if not settings.gemini_api_key:
return "AI Chat is disabled because GEMINI_API_KEY is not configured."
system_prompt = (
prompt = (
"You are SecureLens AI, an expert cybersecurity assistant. "
"You are helping a developer understand a security scan report for their website. "
f"Here is the context of the scan: {json.dumps(context_data)}"
f"Here is the context of the scan: {json.dumps(context_data)}\n\n"
f"User Message: {user_message}"
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message}
],
temperature=0.5,
model_name = await get_gemini_model()
response = await ai_client.aio.models.generate_content(
model=model_name,
contents=prompt,
config=types.GenerateContentConfig(
temperature=0.5,
)
)
return response.choices[0].message.content or "No response from AI."
return response.text or "No response from AI."
except Exception as e:
logger.error(f"AI Chat Error: {str(e)}")
return "I encountered an error trying to process your request."
async def generate_threat_narrative(context_data: dict) -> str:
"""
Weaves multiple scan issues into a cohesive attack sequence.
"""
if not settings.openai_api_key:
return "AI Threat Narrative is disabled because OPENAI_API_KEY is not configured."
if not settings.gemini_api_key:
return "AI Threat Narrative is disabled because GEMINI_API_KEY is not configured."
system_prompt = (
prompt = (
"You are a senior cybersecurity red-teamer. Analyze the following security scan results "
"and weave them into a single, cohesive 'Threat Narrative'. Explain how an attacker might "
"chain these specific vulnerabilities together to compromise the system. "
"Keep it professional, concise (2-3 paragraphs), and actionable."
"Keep it professional, concise (2-3 paragraphs), and actionable.\n\n"
f"Context: {json.dumps(context_data)}"
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": json.dumps(context_data)}
],
temperature=0.7,
model_name = await get_gemini_model()
response = await ai_client.aio.models.generate_content(
model=model_name,
contents=prompt,
config=types.GenerateContentConfig(
temperature=0.7,
)
)
return response.choices[0].message.content or "Could not generate threat narrative."
return response.text or "Could not generate threat narrative."
except Exception as e:
logger.error(f"AI Narrative Error: {str(e)}")
return "I encountered an error trying to generate the threat narrative."

92
app/services/code_scanner/orchestrator.py
View File

@@ -1,7 +1,8 @@
import json
import logging
from typing import List, Dict, Any
from openai import AsyncOpenAI
from google import genai
from google.genai import types
from app.config import settings
from app.services.code_scanner.github_client import GitHubClient
@@ -9,26 +10,28 @@ from app.schemas.code_scan import VulnerabilityIssue
logger = logging.getLogger(__name__)
api_key = settings.openai_api_key or "mock-key-for-testing"
client = AsyncOpenAI(api_key=api_key)
if settings.gemini_api_key:
# google-genai client init
ai_client = genai.Client(api_key=settings.gemini_api_key)
else:
ai_client = None
class CodeScanOrchestrator:
def __init__(self, repo_url: str, github_token: str, branch: str = "main"):
self.repo_url = repo_url
self.branch = branch
self.github = GitHubClient(token=github_token)
# We use gemini-2.5-flash for fast and cost-effective analysis
self.model_name = 'gemini-2.5-flash'
async def triage_files(self, all_files: List[str]) -> List[str]:
"""
Uses the LLM to select which files are most likely to contain security vulnerabilities
(e.g., config files, routers, auth logic).
"""
if not settings.openai_api_key:
logger.warning("OPENAI_API_KEY is not set. Triaging all files up to a limit.")
return all_files[:10] # Hard limit for testing
if not settings.gemini_api_key:
logger.warning("GEMINI_API_KEY is not set. Triaging all files up to a limit.")
return all_files[:10]
# To avoid context limit issues, we might want to chunk this, but for now we pass the list
# We can enforce a soft limit on the string length
files_str = "\n".join(all_files)
if len(files_str) > 15000:
files_str = files_str[:15000] + "\n... (truncated)"
@@ -42,42 +45,35 @@ class CodeScanOrchestrator:
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You always respond with valid JSON."},
{"role": "user", "content": prompt}
],
response_format={"type": "json_object"},
temperature=0.1,
response = await ai_client.aio.models.generate_content(
model=self.model_name,
contents=prompt,
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0.1,
)
)
content = response.choices[0].message.content
if content:
data = json.loads(content)
if response.text:
data = json.loads(response.text)
return data.get("critical_files", [])
except Exception as e:
logger.error(f"Error triaging files: {e}")
return all_files[:10] # Fallback
return all_files[:10]
async def analyze_files(self, triaged_files: List[str]) -> List[VulnerabilityIssue]:
"""
Fetches the contents of the triaged files and uses the LLM to find vulnerabilities.
"""
vulnerabilities = []
if not settings.openai_api_key:
if not settings.gemini_api_key:
return []
# Analyze files sequentially or in batches (sequential to avoid rate limits for now)
for file_path in triaged_files:
content = await self.github.get_file_content(self.repo_url, file_path, self.branch)
if not content:
continue
# Truncate very large files
if len(content) > 20000:
content = content[:20000]
if len(content) > 30000:
content = content[:30000]
prompt = (
f"Review the following code from the file '{file_path}' for security vulnerabilities.\n"
@@ -93,19 +89,16 @@ class CodeScanOrchestrator:
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a SAST security agent. Always respond with valid JSON."},
{"role": "user", "content": prompt}
],
response_format={"type": "json_object"},
temperature=0.2,
response = await ai_client.aio.models.generate_content(
model=self.model_name,
contents=prompt,
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0.2,
)
)
resp_content = response.choices[0].message.content
if resp_content:
data = json.loads(resp_content)
if response.text:
data = json.loads(response.text)
vulns = data.get("vulnerabilities", [])
for v in vulns:
vulnerabilities.append(VulnerabilityIssue(
@@ -125,7 +118,7 @@ class CodeScanOrchestrator:
if not vulnerabilities:
return "No obvious security vulnerabilities found in the scanned files."
if not settings.openai_api_key:
if not settings.gemini_api_key:
return f"Found {len(vulnerabilities)} potential issues."
issues_data = [v.model_dump() for v in vulnerabilities]
@@ -137,15 +130,14 @@ class CodeScanOrchestrator:
)
try:
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a cybersecurity expert."},
{"role": "user", "content": prompt}
],
temperature=0.4,
response = await ai_client.aio.models.generate_content(
model=self.model_name,
contents=prompt,
config=types.GenerateContentConfig(
temperature=0.4,
)
)
return response.choices[0].message.content or "Could not generate summary."
return response.text or "Could not generate summary."
except Exception as e:
logger.error(f"Error generating summary: {e}")
return f"Found {len(vulnerabilities)} potential issues."

1
models.txt Normal file
View File

@@ -0,0 +1 @@
AVAILABLE MODELS: models/gemini-2.5-flash, models/gemini-2.5-pro, models/gemini-2.0-flash, models/gemini-2.0-flash-001, models/gemini-2.0-flash-lite-001, models/gemini-2.0-flash-lite, models/gemini-2.5-flash-preview-tts, models/gemini-2.5-pro-preview-tts, models/gemma-3-1b-it, models/gemma-3-4b-it, models/gemma-3-12b-it, models/gemma-3-27b-it, models/gemma-3n-e4b-it, models/gemma-3n-e2b-it, models/gemma-4-26b-a4b-it, models/gemma-4-31b-it, models/gemini-flash-latest, models/gemini-flash-lite-latest, models/gemini-pro-latest, models/gemini-2.5-flash-lite, models/gemini-2.5-flash-image, models/gemini-3-pro-preview, models/gemini-3-flash-preview, models/gemini-3.1-pro-preview, models/gemini-3.1-pro-preview-customtools, models/gemini-3.1-flash-lite-preview, models/gemini-3-pro-image-preview, models/nano-banana-pro-preview, models/gemini-3.1-flash-image-preview, models/lyria-3-clip-preview, models/lyria-3-pro-preview, models/gemini-3.1-flash-tts-preview, models/gemini-robotics-er-1.5-preview, models/gemini-robotics-er-1.6-preview, models/gemini-2.5-computer-use-preview-10-2025, models/deep-research-max-preview-04-2026, models/deep-research-preview-04-2026, models/deep-research-pro-preview-12-2025, models/gemini-embedding-001, models/gemini-embedding-2-preview, models/gemini-embedding-2, models/aqa, models/imagen-4.0-generate-001, models/imagen-4.0-ultra-generate-001, models/imagen-4.0-fast-generate-001, models/veo-2.0-generate-001, models/veo-3.0-generate-001, models/veo-3.0-fast-generate-001, models/veo-3.1-generate-preview, models/veo-3.1-fast-generate-preview, models/veo-3.1-lite-generate-preview, models/gemini-2.5-flash-native-audio-latest, models/gemini-2.5-flash-native-audio-preview-09-2025, models/gemini-2.5-flash-native-audio-preview-12-2025, models/gemini-3.1-flash-live-preview

2
requirements.txt
View File

@@ -14,6 +14,6 @@ pydantic[email]
pytest
pytest-asyncio
alembic
openai
google-genai
aiodns
fpdf2