vulncheck/app/services/ai_service.py

"""
AI remediation service (OpenRouter, OpenAI-compatible).

On-demand generator that turns a vulnerability + its host context into
concrete, OS-aware remediation steps. Uses OpenRouter's OpenAI-compatible
REST API directly via httpx (no extra SDK dependency).

Config (env first, then settings table, so it works headless or via UI):
  OPENROUTER_API_KEY   — required to enable the feature
  OPENROUTER_MODEL     — default "openrouter/free" (free auto-router)
  OPENROUTER_FALLBACKS — optional comma list for route=fallback

Free tier: ~10 requests/day across free models — fine for on-demand use.
"""
import logging
import os
from typing import List, Optional

import httpx
from sqlalchemy.orm import Session

from app.models.setting import Setting

logger = logging.getLogger(__name__)

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
DEFAULT_MODEL = "openrouter/free"
HTTP_TIMEOUT = 60.0

SETTING_KEY = "openrouter_api_key"
SETTING_MODEL = "openrouter_model"
SETTING_ENABLED = "ai_remediation_enabled"


class AIServiceError(Exception):
    """Raised when AI remediation cannot be produced."""


def _cfg(db: Session, env_name: str, setting_key: str, default: str = "") -> str:
    val = os.getenv(env_name, "").strip()
    if val:
        return val
    try:
        from app.auth.setting_crypto import read_setting_value
        sv = read_setting_value(db, setting_key)
        if sv:
            return str(sv).strip()
    except Exception:
        s = db.query(Setting).filter(Setting.key == setting_key).first()
        if s and s.value:
            return str(s.value).strip()
    return default


def is_enabled(db: Session) -> bool:
    """True when an API key is configured (env or settings)."""
    return bool(_cfg(db, "OPENROUTER_API_KEY", SETTING_KEY))


def _build_messages(*, cve_id, title, description, package, installed,
                    fixed, os_name, scanner_remediation) -> List[dict]:
    sys = (
        "You are a senior security engineer. Given a vulnerability and the "
        "affected host, produce concise, ACTIONABLE remediation guidance for "
        "the specific operating system. Prefer concrete commands in fenced "
        "code blocks (apt/dnf/zypper for Linux distros, PowerShell/winget/MSI "
        "for Windows). Include: 1) the fix (upgrade/patch/config), 2) exact "
        "commands for THIS OS, 3) a verification step, 4) a mitigation if no "
        "patch is available. Be brief — no preamble, no marketing."
    )
    # EOL/EOS pseudo-findings (cve_id starts with EOL- / NESSUS-PLUGIN-) are
    # not patchable CVEs — there is no fix, the product is out of support.
    # Give the model EOL-specific instructions so the answer is an upgrade/
    # replacement plan, not a "apply the patch" hallucination.
    is_eol = bool(cve_id) and (
        cve_id.upper().startswith("EOL-") or cve_id.upper().startswith("NESSUS-PLUGIN-")
    )
    if is_eol:
        sys = (
            "You are a senior IT-security engineer advising on END-OF-LIFE / "
            "END-OF-SUPPORT (EOL/EOS) software. The product below no longer "
            "receives security patches — there is NO CVE patch to apply, so do "
            "NOT suggest 'apply the update'. Produce a concise upgrade/migration "
            "plan for THIS operating system:\n"
            "1) State the EOL/EOS situation and the risk of staying on it.\n"
            "2) The supported target release/edition to move to (name the "
            "current supported version and its support timeline if known).\n"
            "3) Concrete upgrade/replace commands for THIS OS (apt/dnf/zypper "
            "dist-upgrade or repo swap on Linux; winget/MSI/installer or OS "
            "in-place upgrade on Windows), plus where to download the supported "
            "build.\n"
            "4) Interim COMPENSATING CONTROLS / containment while the upgrade is "
            "pending (network isolation/segmentation, restrict exposure, disable "
            "the component, WAF/firewall rules, increased monitoring).\n"
            "5) A verification step. Be brief, no preamble."
        )
    lines = [
        ("EOL/EOS finding ID: " if is_eol else "CVE / ID: ") + str(cve_id),
        f"Title: {title or '—'}",
        f"Affected OS: {os_name or 'unknown'}",
        f"Product / package: {package or '—'}",
        f"Installed version: {installed or '—'}",
        (f"Latest supported version: {fixed or '—'}" if is_eol
         else f"Fixed version: {fixed or '—'}"),
    ]
    if scanner_remediation:
        lines.append(f"Scanner-suggested remediation: {scanner_remediation}")
    if description:
        lines.append(f"\nDescription:\n{description[:1500]}")
    if is_eol:
        lines.append(
            "\nThis is an end-of-life / out-of-support product, NOT a patchable "
            "CVE. Give the upgrade/migration plan + interim compensating controls."
        )
    else:
        lines.append("\nGive the remediation now.")
    return [
        {"role": "system", "content": sys},
        {"role": "user", "content": "\n".join(lines)},
    ]


def generate_remediation(
    db: Session,
    *,
    cve_id: str,
    title: Optional[str] = None,
    description: Optional[str] = None,
    package: Optional[str] = None,
    installed: Optional[str] = None,
    fixed: Optional[str] = None,
    os_name: Optional[str] = None,
    scanner_remediation: Optional[str] = None,
) -> dict:
    """Call OpenRouter and return {"content": str, "model": str}.

    Synchronous + blocking — the caller must run it OFF the event loop
    (asyncio.to_thread) so it never freezes the GUI.
    """
    api_key = _cfg(db, "OPENROUTER_API_KEY", SETTING_KEY)
    if not api_key:
        raise AIServiceError(
            "OpenRouter not configured — set OPENROUTER_API_KEY (env or Settings)."
        )
    model = _cfg(db, "OPENROUTER_MODEL", SETTING_MODEL, DEFAULT_MODEL)
    fallbacks = _cfg(db, "OPENROUTER_FALLBACKS", "openrouter_fallbacks", "")

    body = {
        "model": model,
        "messages": _build_messages(
            cve_id=cve_id, title=title, description=description, package=package,
            installed=installed, fixed=fixed, os_name=os_name,
            scanner_remediation=scanner_remediation,
        ),
    }
    if fallbacks:
        models = [model] + [m.strip() for m in fallbacks.split(",") if m.strip()]
        body["models"] = models
        body["route"] = "fallback"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        # OpenRouter attribution headers (optional but recommended).
        "HTTP-Referer": "https://truevuln.local",
        "X-Title": "TrueVuln",
    }
    try:
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.post(OPENROUTER_URL, headers=headers, json=body)
    except httpx.HTTPError as e:
        raise AIServiceError(f"OpenRouter request failed: {e}") from e

    if resp.status_code == 401:
        raise AIServiceError("OpenRouter rejected the API key (401).")
    if resp.status_code == 402:
        raise AIServiceError(
            "OpenRouter quota/credits exhausted (402) — free-tier daily cap hit "
            "or a paid model needs credits."
        )
    if resp.status_code >= 400:
        raise AIServiceError(f"OpenRouter error {resp.status_code}: {resp.text[:300]}")

    try:
        data = resp.json()
        content = data["choices"][0]["message"]["content"]
        used_model = data.get("model", model)
    except (KeyError, IndexError, ValueError) as e:
        raise AIServiceError(f"Unexpected OpenRouter response shape: {e}") from e

    if not content or not content.strip():
        raise AIServiceError("OpenRouter returned an empty response.")
    return {"content": content.strip(), "model": used_model}