talos/backend/app/llm.py

import json
import uuid
from typing import Any

import httpx
import structlog

from app.config import settings
from app.models import Finding, Phase

log = structlog.get_logger(__name__)

SYSTEM_PROMPT = """You are a senior penetration tester analyzing raw security tool output.
Identify real vulnerabilities (not noise), assign accurate severity, and write concrete remediation steps.

Return ONLY valid JSON matching this schema:
{
  "findings": [
    {
      "severity": "critical|high|medium|low|info",
      "title": "Short specific title",
      "description": "What the issue is and why it matters",
      "evidence": "Exact snippet from the tool output that shows this",
      "recommendation": "Concrete fix steps the operator can take",
      "cve_refs": ["CVE-YYYY-XXXX"],
      "confidence": "high|medium|low"
    }
  ]
}

Rules:
- Only report real findings. Empty findings list is valid.
- "info" severity for benign observations (open standard ports without issues).
- Cite specific CVEs only when version info confirms them.
- Be concise. No marketing language. No filler.
"""

USER_TEMPLATE = """Target: {target}
Tool: {tool}
Command: {command}
Profile: {profile}

Raw output:
{output}

Analyze and return JSON."""


def truncate_output(text: str, max_chars: int | None = None) -> str:
    max_chars = max_chars or settings.llm_output_max_chars
    if len(text) <= max_chars:
        return text
    half = max_chars // 2
    return text[:half] + "\n\n...[truncated]...\n\n" + text[-half:]


async def ollama_health() -> bool:
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            r = await client.get(f"{settings.ollama_url.rstrip('/')}/api/tags")
            return r.status_code == 200
    except httpx.HTTPError:
        return False


async def _chat(messages: list[dict[str, str]], model: str) -> str:
    url = f"{settings.ollama_url.rstrip('/')}/api/chat"
    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "stream": False,
        "format": "json",
    }
    async with httpx.AsyncClient(timeout=120.0) as client:
        r = await client.post(url, json=payload)
        r.raise_for_status()
        data = r.json()
        return str(data.get("message", {}).get("content", ""))


def _parse_findings_json(
    raw: str,
    scan_id: str,
    phase: Phase,
) -> list[Finding]:
    parsed = json.loads(raw)
    items = parsed.get("findings", []) if isinstance(parsed, dict) else []
    findings: list[Finding] = []
    for item in items:
        try:
            f = Finding(
                id=str(uuid.uuid4()),
                scan_id=scan_id,
                phase_id=phase.id,
                severity=item["severity"],
                title=item["title"],
                description=item["description"],
                evidence=item.get("evidence", ""),
                recommendation=item.get("recommendation", ""),
                cve_refs=item.get("cve_refs", []),
                confidence=item.get("confidence", "medium"),
            )
            findings.append(f)
        except Exception as exc:
            log.warning("finding_validation_failed", error=str(exc), item=item)
    return findings


async def analyze_phase(
    phase: Phase,
    output: str,
    *,
    target: str,
    profile: str,
    model: str,
) -> tuple[list[Finding], bool]:
    """Returns (findings, llm_unavailable)."""
    if not await ollama_health():
        return [], True

    truncated = truncate_output(output)
    user_msg = USER_TEMPLATE.format(
        target=target,
        tool=phase.tool,
        command=phase.command,
        profile=profile,
        output=truncated,
    )
    messages: list[dict[str, str]] = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]

    for attempt in range(2):
        try:
            content = await _chat(messages, model)
            return _parse_findings_json(content, phase.scan_id, phase), False
        except json.JSONDecodeError:
            if attempt == 0:
                messages.append(
                    {
                        "role": "user",
                        "content": "Your previous response was not valid JSON — try again.",
                    }
                )
            else:
                log.warning("llm_json_parse_failed", phase_id=phase.id)
        except httpx.HTTPError as exc:
            log.warning("llm_request_failed", error=str(exc))
            return [], True

    return [], False