vulncheck/app/integrations/infomaniak_ai_client.py

"""
Infomaniak AI API Client (Legacy Wrapper)
This file now serves as a wrapper around the Universal AI Client.
"""
import logging
import json
from typing import Optional, Dict, Any, List
from sqlalchemy.orm import Session
from fastapi import Depends
from app.database import get_db
from app.models.setting import Setting
from app.integrations.ai_client import UniversalAIClient, AIError

logger = logging.getLogger(__name__)

# Maintain alias for compatibility
InfomaniakAIClient = UniversalAIClient
InfomaniakAIError = AIError

def get_infomaniak_ai_client(db: Session = Depends(get_db)) -> UniversalAIClient:
    """
    Factory function for dependency injection.
    """
    config_setting = db.query(Setting).filter(Setting.key == "ai_config").first()
    if config_setting and config_setting.value:
        try:
            config = json.loads(config_setting.value)
            return UniversalAIClient.from_config(config)
        except Exception:
            logger.error("Invalid AI config in DB")
    
    return UniversalAIClient()

# Core service layer remains similar but uses the universal client
from app.models.ai_analysis import AIAnalysis
from app.models.vulnerability import Vulnerability
from app.models.asset import Asset
from datetime import datetime
from sqlalchemy import desc

class AIAnalysisService:
    """Service for CVE analysis with caching"""

    def __init__(self, db: Session, ai_client: UniversalAIClient):
        self.db = db
        self.ai_client = ai_client

    def get_or_create_analysis(
        self,
        vulnerability_id: int,
        force_refresh: bool = False
    ) -> AIAnalysis:
        if not force_refresh:
            cached = self.db.query(AIAnalysis).filter(
                AIAnalysis.vulnerability_id == vulnerability_id,
                AIAnalysis.cache_expires_at > datetime.now()
            ).first()

            if cached:
                return cached

        vuln = self.db.query(Vulnerability).filter(Vulnerability.id == vulnerability_id).first()
        if not vuln:
            raise ValueError(f"Vulnerability {vulnerability_id} not found")

        analysis_data = self.ai_client.analyze_cve(
            cve_id=vuln.cve_id,
            package_name=vuln.package_name,
            package_version=vuln.package_version,
            cvss_score=vuln.cvss_score
        )

        existing = self.db.query(AIAnalysis).filter(AIAnalysis.vulnerability_id == vulnerability_id).first()
        if existing:
            for key, value in analysis_data.items():
                setattr(existing, key, value)
            self.db.commit()
            self.db.refresh(existing)
            return existing
        else:
            ai_analysis = AIAnalysis(vulnerability_id=vulnerability_id, **analysis_data)
            self.db.add(ai_analysis)
            self.db.commit()
            self.db.refresh(ai_analysis)
            return ai_analysis

    def get_priority_recommendations(self, limit: int = 20, severity: Optional[str] = None) -> Dict[str, Any]:
        # Same scope the reports and the dashboard use: real CVEs on assets
        # that still exist. Without it the audit ranked findings on
        # decommissioned hosts and treated EOL- pseudo-CVEs as vulnerabilities.
        from app.services.report_scope import scoped
        query = scoped(self.db).filter(Vulnerability.status == "open")
        if severity:
            query = query.filter(Vulnerability.severity == severity)

        # ORDER BY cvss_score DESC put the WORST candidates first: Postgres
        # sorts NULLs before everything on DESC, so every finding without a
        # score — fresh Chrome CVEs carry none at all — was handed to the model
        # as the top of the list, ahead of a scored 9.8. Rank by the scores the
        # product computes for exactly this question, and put nulls last.
        top_vulns = (query.order_by(
            Vulnerability.priority_score.desc().nullslast(),
            Vulnerability.cpr_score.desc().nullslast(),
            Vulnerability.cvss_score.desc().nullslast(),
        # Over-fetch, because the dedup below happens after the LIMIT: one CVE
        # spread over 30 hosts would otherwise leave the model a handful of
        # distinct findings instead of `limit` of them.
        ).limit(limit * 5).all())
        if not top_vulns:
            return {"recommendations": [], "global_strategy": "No open vulnerabilities found."}

        # One entry per CVE. The same CVE on 30 hosts used to fill the whole
        # list, so the model ranked one problem thirty times and never saw the
        # other 29.
        vuln_data = []
        seen_cves = set()
        for v in top_vulns:
            if len(vuln_data) >= limit:
                break
            if v.cve_id in seen_cves:
                continue
            seen_cves.add(v.cve_id)
            # The model was given CVE id, host, CVSS and package — less than the
            # dashboard shows. It could not weigh what actually decides urgency:
            # known exploitation (KEV/EUVD), exploit probability (EPSS), and the
            # asset's own criticality. Those are the fields that make an 7.5
            # outrank a 9.8.
            vuln_data.append({
                "cve_id": v.cve_id,
                "hostname": v.asset.hostname if v.asset else "Unknown",
                "cvss_score": v.cvss_score,
                "package_name": v.package_name,
                "severity": v.severity.value if v.severity else None,
                "priority_score": v.priority_score,
                "epss_percentile": v.epss_percentile,
                "known_exploited_cisa_kev": bool(v.kev_listed),
                "listed_enisa_euvd": bool(v.euvd_listed),
                "fixed_version": v.fixed_version,
                "days_open": ((datetime.now() - v.detected_at).days
                              if v.detected_at else None),
            })

        # Prompt for structured prioritization
        # "top 10" was hardcoded while `limit` defaulted to 20 — the model was
        # told to discard half its input for no stated reason, and at limit<10
        # it was asked for more items than it had been given.
        want = min(10, len(vuln_data))
        prompt = f"""You are a vulnerability management expert. Rank the {want} findings below that need attention first.

Rank by exploitability and blast radius, not by CVSS alone:
- known_exploited_cisa_kev = true outranks a higher CVSS that is not exploited.
- listed_enisa_euvd = true is the EU equivalent signal.
- epss_percentile is the probability of exploitation in the next 30 days (0-1).
- days_open shows how long the finding has been unaddressed.
- fixed_version = null means no patch exists yet, so advise mitigation instead.
- cvss_score may be null when the vendor published no score; judge such a
  finding on its severity and the signals above rather than skipping it.

Findings ({len(vuln_data)} distinct CVEs):
{json.dumps(vuln_data)}

Return exactly {want} recommendations, most urgent first, priority numbered
from 1. Use only cve_id values from the list above — never invent one.

You MUST respond with a valid JSON object ONLY, following this exact structure:
{{
  "global_strategy": "Brief strategic summary of the risk landscape and focus areas.",
  "recommendations": [
    {{
      "priority": 1,
      "cve_id": "CVE-XXXX-XXXX",
      "asset_hostname": "hostname",
      "justification": "Why this is critical...",
      "remediation_hint": "Brief fix action..."
    }}
  ]
}}
"""
        try:
            response_text = self.ai_client._call_llm(prompt, json_mode=True)
            
            # Clean DeepSeek <think> tags if present
            if "<think>" in response_text:
                import re
                response_text = re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL).strip()

            # Robust JSON Extraction
            try:
                # 1. Try removing markdown code blocks
                clean_text = response_text
                if "```json" in clean_text:
                    clean_text = clean_text.split("```json")[1].split("```")[0].strip()
                elif "```" in clean_text:
                    clean_text = clean_text.split("```")[1].split("```")[0].strip()
                
                return json.loads(clean_text)
            except (json.JSONDecodeError, IndexError, TypeError):
                # 2. Aggressive Fallback: Find first '{' and last '}'
                try:
                    start_idx = response_text.find("{")
                    end_idx = response_text.rfind("}")
                    if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
                        json_str = response_text[start_idx : end_idx + 1]
                        return json.loads(json_str)
                except Exception:
                    pass
                
                logger.warning(f"Failed to parse AI response. Raw output: {response_text[:200]}...")
                return {"global_strategy": response_text, "recommendations": []}
        except Exception as e:
            logger.error(f"AI prioritization failed: {e}")
            return {"error": str(e), "recommendations": []}