feat: integrate LiteLLM for provider-agnostic AI (supports Gemini, OpenAI, Claude, Ollama)

app/config.py

@@ -1,9 +1,32 @@
+"""
+Application Configuration
+==========================
+
+All configuration is loaded from environment variables (or a .env file).
+
+AI Provider Configuration
+--------------------------
+SecureLens supports multiple AI providers via LiteLLM.
+Set AI_MODEL and AI_API_KEY in your .env to choose a provider:
+
+  Provider     | AI_MODEL example                     | Where to get key
+  -------------|--------------------------------------|-----------------------------
+  Gemini       | gemini/gemini-2.0-flash              | aistudio.google.com
+  OpenAI       | gpt-4o-mini                          | platform.openai.com
+  Claude       | claude-3-5-haiku-20241022            | console.anthropic.com
+  OpenRouter   | openrouter/google/gemini-2.0-flash   | openrouter.ai
+  Ollama       | ollama/llama3.1                      | ollama.com (local, no key)
+
+If you only set GEMINI_API_KEY (legacy), it will still work — the app
+automatically maps it to the Gemini provider for backward compatibility.
+"""
+
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class Settings(BaseSettings):
     app_name: str = "SecureLens AI"
-    app_version: str = "1.0.0"
+    app_version: str = "1.1.0"
     debug: bool = False
 
     host: str = "0.0.0.0"
@@ -22,13 +45,39 @@ class Settings(BaseSettings):
     jwt_algorithm: str = "HS256"
     jwt_expiry_minutes: int = 1440
 
+    # -------------------------------------------------------------------------
+    # AI Provider Settings (new, provider-agnostic)
+    # -------------------------------------------------------------------------
+    # AI_MODEL: the LiteLLM model string (see table in module docstring above)
+    ai_model: str = "gemini/gemini-2.0-flash"
+
+    # AI_API_KEY: the API key for the chosen provider.
+    # Leave blank for Ollama (local, no key needed).
+    ai_api_key: str | None = None
+
+    # -------------------------------------------------------------------------
+    # Legacy Gemini key — kept for backward compatibility.
+    # If AI_API_KEY is not set but GEMINI_API_KEY is, we use that automatically.
+    # -------------------------------------------------------------------------
     gemini_api_key: str | None = None
 
+    # Threat Intelligence API keys (Step 3)
+    virustotal_api_key: str | None = None
+    abuseipdb_api_key: str | None = None
+
     model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
 
     @property
     def cors_origin_list(self) -> list[str]:
         return [origin.strip() for origin in self.cors_origins.split(",") if origin.strip()]
 
+    @property
+    def effective_ai_key(self) -> str | None:
+        """
+        Returns the resolved AI API key.
+        Prefers AI_API_KEY; falls back to GEMINI_API_KEY for backward compatibility.
+        """
+        return self.ai_api_key or self.gemini_api_key
+
 
 settings = Settings()

app/routers/scan.py

@@ -22,6 +22,7 @@ from app.services.scanner.dns import DNSScanner
 from app.services.scanner.ports import PortScanner
 from app.services.scoring import calculate_layer_statuses, calculate_score
 from app.services.ai import enhance_security_issues
+from app.services.threat_intel import get_threat_intel_summary
 from app.utils.validators import validate_url
 
 logger = logging.getLogger(__name__)
@@ -78,6 +79,8 @@ async def scan_website(
         
         dns_task = asyncio.create_task(dns_scanner.scan(url))
         port_task = asyncio.create_task(port_scanner.scan(url))
+        # Step 3: Run threat intel lookup concurrently — zero extra latency
+        threat_intel_task = asyncio.create_task(get_threat_intel_summary(url))
 
         async with httpx.AsyncClient(
             timeout=httpx.Timeout(settings.scan_timeout),
@@ -95,11 +98,12 @@ async def scan_website(
         # Await infrastructure scans
         all_issues.extend(await dns_task)
         all_issues.extend(await port_task)
+        threat_intel = await threat_intel_task
 
         score = calculate_score(all_issues)
         layers = calculate_layer_statuses(all_issues)
 
-        if settings.gemini_api_key and all_issues:
+        if settings.effective_ai_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", [])
@@ -144,6 +148,7 @@ async def scan_website(
             layers=layers,
             issues=all_issues,
             created_at=created_at,
+            threat_intel=threat_intel,  # Step 3: attach threat intelligence
         )
 
     except httpx.HTTPError as e:

app/services/ai.py

@@ -1,24 +1,108 @@
+"""
+AI Service Layer — Provider-Agnostic via LiteLLM
+==================================================
+
+Why LiteLLM?
+  Previously every AI call used the google-genai SDK directly, which meant
+  the entire codebase was hard-wired to Gemini. Switching to OpenAI or
+  Claude would require rewriting every file that touched AI.
+
+  LiteLLM is a thin translation layer. You call one function, it handles
+  the right SDK under the hood based on the model string you pass:
+    - "gpt-4o-mini"                  → OpenAI
+    - "claude-3-5-haiku-20241022"    → Anthropic
+    - "gemini/gemini-2.0-flash"      → Google Gemini
+    - "ollama/llama3.1"              → local Ollama instance
+    - "openrouter/..."               → OpenRouter
+
+  Now you only need to change two env vars (AI_MODEL, AI_API_KEY) to switch
+  providers — no code changes required.
+
+Public API (used by the rest of the app):
+  call_ai(prompt, temperature, json_mode)  → str
+  enhance_security_issues(issues)           → dict
+  chat_with_scan_context(...)              → str
+  generate_threat_narrative(context_data)  → str
+"""
+
 import json
 import logging
-import asyncio
-from google import genai
-from google.genai import types
 from app.config import settings
 
 logger = logging.getLogger(__name__)
 
-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.0-flash'
+# Core LiteLLM wrapper
+# ---------------------------------------------------------------------------
+
+async def call_ai(
+    prompt: str,
+    temperature: float = 0.3,
+    json_mode: bool = False,
+) -> str:
+    """
+    The single entry-point for all AI calls in SecureLens.
+
+    Parameters
+    ----------
+    prompt      : The full prompt string to send to the model.
+    temperature : Creativity level (0 = deterministic, 1 = creative).
+    json_mode   : If True, instruct the model to return valid JSON only.
+                  This maps to response_format={"type":"json_object"} on
+                  providers that support it (OpenAI, Gemini via LiteLLM).
+
+    Returns
+    -------
+    The model's text response as a plain string. Empty string on failure.
+    """
+    import litellm
+
+    api_key = settings.effective_ai_key
+    model = settings.ai_model
+
+    if not api_key and not model.startswith("ollama/"):
+        logger.warning("No AI API key configured. Skipping AI call.")
+        return ""
+
+    messages = [{"role": "user", "content": prompt}]
+
+    kwargs: dict = {
+        "model": model,
+        "messages": messages,
+        "temperature": temperature,
+        "api_key": api_key,
+    }
+
+    # JSON mode: supported natively by OpenAI and LiteLLM proxied Gemini.
+    # For providers that don't support it, LiteLLM silently ignores the flag.
+    if json_mode:
+        kwargs["response_format"] = {"type": "json_object"}
+
+    try:
+        response = await litellm.acompletion(**kwargs)
+        return response.choices[0].message.content or ""
+    except Exception as e:
+        logger.error(f"LiteLLM call failed [model={model}]: {e}")
+        return ""
+
+
+# ---------------------------------------------------------------------------
+# Domain-specific AI functions
+# ---------------------------------------------------------------------------
 
 async def enhance_security_issues(issues: list[dict]) -> dict:
-    if not settings.gemini_api_key:
+    """
-        logger.warning("GEMINI_API_KEY is not set. AI enhancements are skipped.")
+    Takes a raw list of scanner-detected issues and enriches each one with:
+      - contextual_severity : AI-assessed severity in the real-world context
+      - explanation         : Plain-English description of the risk
+      - remediation_snippet : Concrete code or config fix
+
+    Returns a dict {"enhanced_issues": [...]} that mirrors the original list
+    with the three new fields merged in.
+    """
+    if not settings.effective_ai_key:
+        logger.warning("AI enhancements skipped — no AI API key set.")
         return {"enhanced_issues": issues}
 
     prompt = (
@@ -28,75 +112,61 @@ async def enhance_security_issues(issues: list[dict]) -> dict:
         "Return a JSON object with a single key 'enhanced_issues' containing a list of objects. "
         "Each object MUST correspond to one of the original issues and have the following keys: "
         "'issue' (exact string of the original issue), "
-        "'contextual_severity' (Low, Medium, High, Critical), "
+        "'contextual_severity' (Low, Medium, High, or Critical), "
-        "'explanation' (a 1-2 sentence non-technical explanation), "
+        "'explanation' (a 1-2 sentence non-technical explanation of the real risk), "
-        "'remediation_snippet' (Actionable code snippet, e.g. Nginx config, or 'N/A')."
+        "'remediation_snippet' (an actionable code snippet or config fix, or 'N/A')."
     )
 
+    raw = await call_ai(prompt, temperature=0.2, json_mode=True)
+    if not raw:
+        return {"enhanced_issues": issues, "ai_error": "Empty response from AI"}
+
     try:
-        model_name = await get_gemini_model()
+        return json.loads(raw)
-        response = await ai_client.aio.models.generate_content(
+    except json.JSONDecodeError as e:
-            model=model_name,
+        logger.error(f"Failed to parse AI JSON response: {e}\nRaw: {raw[:500]}")
-            contents=prompt,
+        return {"enhanced_issues": issues, "ai_error": "JSON parse error"}
-            config=types.GenerateContentConfig(
+
-                response_mime_type="application/json",
-                temperature=0.2,
-            )
-        )
-        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:
-    if not settings.gemini_api_key:
+    """
-        return "AI Chat is disabled because GEMINI_API_KEY is not configured."
+    Powers the conversational chat feature for web scans.
+
+    The full scan context (score, layers, issues) is injected into the prompt
+    so the model can answer specific questions about the scan results.
+    """
+    if not settings.effective_ai_key:
+        return "AI Chat is disabled because no AI API key is configured."
 
     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)}\n\n"
+        f"Here is the context of the scan:\n{json.dumps(context_data, indent=2)}\n\n"
-        f"User Message: {user_message}"
+        f"Developer's question: {user_message}\n\n"
+        "Answer clearly and professionally. Reference specific findings from the scan when relevant."
     )
 
-    try:
+    result = await call_ai(prompt, temperature=0.5)
-        model_name = await get_gemini_model()
+    return result or "I couldn't generate a response. Please try again."
-        response = await ai_client.aio.models.generate_content(
+
-            model=model_name,
-            contents=prompt,
-            config=types.GenerateContentConfig(
-                temperature=0.5,
-            )
-        )
-        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:
-    if not settings.gemini_api_key:
+    """
-        return "AI Threat Narrative is disabled because GEMINI_API_KEY is not configured."
+    Generates a 2-3 paragraph red-team style threat narrative.
+
+    Explains how an attacker could chain the discovered vulnerabilities
+    together to compromise the system. Used in the PDF report.
+    """
+    if not settings.effective_ai_key:
+        return "AI Threat Narrative is disabled because no AI API key is configured."
 
     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.\n\n"
-        f"Context: {json.dumps(context_data)}"
+        f"Scan Context:\n{json.dumps(context_data, indent=2)}"
     )
 
-    try:
+    result = await call_ai(prompt, temperature=0.7)
-        model_name = await get_gemini_model()
+    return result or "Could not generate threat narrative."
-        response = await ai_client.aio.models.generate_content(
-            model=model_name,
-            contents=prompt,
-            config=types.GenerateContentConfig(
-                temperature=0.7,
-            )
-        )
-        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."

app/services/code_scanner/orchestrator.py

@@ -1,36 +1,50 @@
+"""
+Code Scan Orchestrator
+=======================
+
+Coordinates the three phases of an agentic code security scan:
+  1. Triage   — Ask the AI which files are worth scanning.
+  2. Analyze  — Send each file's code to the AI for OWASP vulnerability review.
+  3. Summarize— Generate an executive summary of all findings.
+
+Previously this used the google-genai SDK directly. It now delegates all AI
+calls to app.services.ai.call_ai(), which is provider-agnostic via LiteLLM.
+This means switching from Gemini to Claude (or any other model) automatically
+applies to the code scanner without any changes here.
+"""
+
 import json
 import logging
-from typing import List, Dict, Any
-from google import genai
-from google.genai import types
 import asyncio
+from typing import List
 
 from app.config import settings
+from app.services.ai import call_ai
 from app.services.code_scanner.github_client import GitHubClient
 from app.schemas.code_scan import VulnerabilityIssue
 
 logger = logging.getLogger(__name__)
 
-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.0-flash for high rate limits and stability
-        self.model_name = 'gemini-2.0-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 
+        Phase 1 — AI-driven file triage.
+
+        Sends the full file tree to the LLM and asks it to select the
+        most security-critical files (e.g. auth handlers, DB queries,
+        config files). Caps at 5 files to stay within token budgets.
+
+        Falls back to the first 5 files if the AI call fails or no key
+        is configured.
         """
-        if not settings.gemini_api_key:
+        if not settings.effective_ai_key:
-            logger.warning("GEMINI_API_KEY is not set. Triaging all files up to a limit.")
+            logger.warning("No AI key set. Falling back to first 5 files.")
             return all_files[:5]
 
         files_str = "\n".join(all_files)
@@ -40,85 +54,89 @@ class CodeScanOrchestrator:
         prompt = (
             "You are a Senior Application Security Engineer. I have a repository with the following files:\n"
             f"{files_str}\n\n"
-            "Select the most critical files to review for security vulnerabilities (e.g., SAST, hardcoded secrets, SQLi, Auth bypass). "
+            "Select the most critical files to review for security vulnerabilities "
-            "Return a JSON object with a single key 'critical_files' containing a list of the exact file paths. "
+            "(e.g. authentication, database access, config, API routes, secrets handling). "
-            "Do not select more than 5 files."
+            "Return a JSON object with a single key 'critical_files' containing a list of "
+            "the exact file paths from the list above. Do not select more than 5 files."
         )
 
         try:
-            response = await ai_client.aio.models.generate_content(
+            raw = await call_ai(prompt, temperature=0.1, json_mode=True)
-                model=self.model_name,
+            if raw:
-                contents=prompt,
+                data = json.loads(raw)
-                config=types.GenerateContentConfig(
-                    response_mime_type="application/json",
-                    temperature=0.1,
-                )
-            )
-            if response.text:
-                data = json.loads(response.text)
                 return data.get("critical_files", [])
         except Exception as e:
-            logger.error(f"Error triaging files: {e}")
+            logger.error(f"File triage failed: {e}")
 
         return all_files[:5]
 
     async def analyze_files(self, triaged_files: List[str]) -> List[VulnerabilityIssue]:
-        if not settings.gemini_api_key:
+        """
+        Phase 2 — Per-file SAST analysis.
+
+        Downloads each file's source code from GitHub and sends it to
+        the AI for a focused OWASP Top-10 vulnerability review.
+
+        Concurrency is throttled with a semaphore to avoid hitting
+        provider rate limits (max 5 simultaneous AI requests).
+        """
+        if not settings.effective_ai_key:
             return []
 
         vulnerabilities = []
-        semaphore = asyncio.Semaphore(5)  # Max 5 concurrent requests to avoid rate limits
+        # Limit concurrent AI calls to avoid rate-limiting
+        semaphore = asyncio.Semaphore(5)
 
-        async def process_file(file_path: str):
+        async def process_file(file_path: str) -> List[VulnerabilityIssue]:
-            # Skip massive dependency lock files as they are too slow and unhelpful for SAST
+            # Skip lock files — huge, slow, zero security signal
-            if file_path.endswith('package-lock.json') or file_path.endswith('yarn.lock'):
+            if file_path.endswith(("package-lock.json", "yarn.lock", "poetry.lock")):
                 return []
 
-            content = await self.github.get_file_content(self.repo_url, file_path, self.branch)
+            content = await self.github.get_file_content(
+                self.repo_url, file_path, self.branch
+            )
             if not content:
                 return []
 
+            # Cap file size to avoid token overflows
             if len(content) > 30000:
                 content = content[:30000]
 
             prompt = (
-                f"Review the following code from the file '{file_path}' for security vulnerabilities.\n"
+                f"Review the following code from '{file_path}' for security vulnerabilities.\n"
-                "Focus on OWASP Top 10: SQLi, XSS, Hardcoded Secrets, IDOR, Misconfigurations, etc.\n\n"
+                "Focus on OWASP Top 10: SQL Injection, XSS, Hardcoded Secrets, IDOR, "
+                "Insecure Deserialization, Broken Auth, Misconfigurations, SSRF, etc.\n\n"
                 f"CODE:\n{content}\n\n"
                 "Return a JSON object with a key 'vulnerabilities' containing a list of objects. "
-                "Each object MUST have the following keys: "
+                "Each object MUST have the following keys:\n"
-                "'severity' (Critical, High, Medium, Low), "
+                "  'severity'     : Critical | High | Medium | Low\n"
-                "'issue' (A short title), "
+                "  'issue'        : Short title of the vulnerability\n"
-                "'explanation' (1-2 sentences explaining the vulnerability), "
+                "  'explanation'  : 1-2 sentences explaining the risk\n"
-                "'suggested_fix' (Code snippet or clear instructions to fix), "
+                "  'suggested_fix': Code snippet or clear instruction to fix it\n"
-                "'line_number' (integer or null if general)."
+                "  'line_number'  : Integer line number, or null if not applicable\n"
+                "If no vulnerabilities are found, return {\"vulnerabilities\": []}."
             )
 
             file_vulns = []
             async with semaphore:
                 try:
-                    response = await ai_client.aio.models.generate_content(
+                    raw = await call_ai(prompt, temperature=0.2, json_mode=True)
-                        model=self.model_name,
+                    if raw:
-                        contents=prompt,
+                        data = json.loads(raw)
-                        config=types.GenerateContentConfig(
+                        for v in data.get("vulnerabilities", []):
-                            response_mime_type="application/json",
+                            file_vulns.append(
-                            temperature=0.2,
+                                VulnerabilityIssue(
-                        )
+                                    file_path=file_path,
-                    )
+                                    severity=v.get("severity", "Medium"),
-                    if response.text:
+                                    issue=v.get("issue", "Unknown Issue"),
-                        data = json.loads(response.text)
+                                    explanation=v.get("explanation", ""),
-                        vulns = data.get("vulnerabilities", [])
+                                    suggested_fix=v.get("suggested_fix"),
-                        for v in vulns:
+                                    line_number=v.get("line_number"),
-                            file_vulns.append(VulnerabilityIssue(
+                                )
-                                file_path=file_path,
+                            )
-                                severity=v.get("severity", "Medium"),
-                                issue=v.get("issue", "Unknown Issue"),
-                                explanation=v.get("explanation", ""),
-                                suggested_fix=v.get("suggested_fix"),
-                                line_number=v.get("line_number")
-                            ))
                 except Exception as e:
-                    logger.error(f"Error analyzing file {file_path}: {e}")
+                    logger.error(f"Analysis failed for {file_path}: {e}")
+
             return file_vulns
 
         results = await asyncio.gather(*(process_file(f) for f in triaged_files))
@@ -128,29 +146,26 @@ class CodeScanOrchestrator:
         return vulnerabilities
 
     async def generate_summary(self, vulnerabilities: List[VulnerabilityIssue]) -> str:
-        if not vulnerabilities:
+        """
-            return "No obvious security vulnerabilities found in the scanned files."
+        Phase 3 — Executive summary.
 
-        if not settings.gemini_api_key:
+        Asks the AI to distill all findings into a 2-3 paragraph summary
-            return f"Found {len(vulnerabilities)} potential issues."
+        suitable for a security report or management briefing.
+        """
+        if not vulnerabilities:
+            return "No security vulnerabilities were identified in the scanned files."
+
+        if not settings.effective_ai_key:
+            return f"Found {len(vulnerabilities)} potential issue(s) across the scanned files."
 
         issues_data = [v.model_dump() for v in vulnerabilities]
         prompt = (
-            "You are a Senior AppSec Manager. Summarize the following list of vulnerabilities found in a recent scan. "
+            "You are a Senior AppSec Manager. Summarize the following list of vulnerabilities "
-            "Provide a 2-3 paragraph executive summary of the repository's security posture. "
+            "found in a recent automated security scan. Provide a 2-3 paragraph executive summary "
-            "Keep it professional and highlight the most critical risks.\n\n"
+            "of the repository's overall security posture. Highlight the most critical risks "
-            f"{json.dumps(issues_data)}"
+            "and recommend the immediate priorities. Keep it professional and actionable.\n\n"
+            f"Findings:\n{json.dumps(issues_data, indent=2)}"
         )
 
-        try:
+        result = await call_ai(prompt, temperature=0.4)
-            response = await ai_client.aio.models.generate_content(
+        return result or f"Found {len(vulnerabilities)} potential issue(s)."
-                model=self.model_name,
-                contents=prompt,
-                config=types.GenerateContentConfig(
-                    temperature=0.4,
-                )
-            )
-            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."

requirements.txt

@@ -17,3 +17,7 @@ alembic
 google-genai
 aiodns
 fpdf2
+# ---- Step 2: LiteLLM for provider-agnostic AI calls ----
+litellm
+# ---- Step 3: Threat Intelligence ----
+# (VirusTotal and AbuseIPDB are queried via httpx, no extra SDK needed)