trustos/backend/app/services/ai_translator.py

"""
AI Risk Translator — calls OpenAI/Anthropic to generate plain-English
business-impact explanations for security findings.
"""
from datetime import datetime
from typing import Optional
from app.db.session import AsyncSessionLocal
from app.models.models import Finding, AttackPath
from sqlalchemy import select
import json
import logging

logger = logging.getLogger(__name__)

TRANSLATION_SYSTEM_PROMPT = """You are TrustOS, an AI cyber resilience advisor. 
Your role is to translate technical cybersecurity findings into clear, plain-English 
business impact statements for executive and non-technical audiences.

Rules:
- Never use CVE IDs, CVSS scores, or technical jargon in the executive summary
- Always frame risk in terms of business impact: customers, revenue, operations, reputation
- Be direct and calm — not alarmist, not dismissive
- Always provide a clear recommended action
- Output must be valid JSON matching the schema provided

Output JSON schema:
{
  "summary": "One sentence: what this is in plain English",
  "business_impact": "1-2 sentences: what could happen to the business if exploited",
  "impact_level": "Low|Medium|High|Critical",
  "remediation_steps": "3-5 concrete steps to fix this, numbered",
  "fix_priority": "urgent|soon|planned"
}"""


async def _call_llm(prompt: str) -> Optional[str]:
    """Call the configured LLM provider. Returns raw text response."""
    from app.core.config import settings
    try:
        if settings.AI_PROVIDER == "openai" and settings.OPENAI_API_KEY:
            from openai import AsyncOpenAI
            client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
            resp = await client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": TRANSLATION_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.3,
                response_format={"type": "json_object"},
            )
            return resp.choices[0].message.content
        elif settings.AI_PROVIDER == "anthropic" and settings.ANTHROPIC_API_KEY:
            from anthropic import AsyncAnthropic
            client = AsyncAnthropic(api_key=settings.ANTHROPIC_API_KEY)
            resp = await client.messages.create(
                model="claude-3-haiku-20240307",
                max_tokens=1024,
                system=TRANSLATION_SYSTEM_PROMPT,
                messages=[{"role": "user", "content": prompt}],
            )
            return resp.content[0].text
        else:
            logger.warning("No AI provider configured — skipping translation")
            return None
    except Exception as e:
        logger.error(f"LLM call failed: {e}")
        return None


async def translate_finding_async(finding_id: str):
    """Background task: generate AI translation for a finding and persist it."""
    async with AsyncSessionLocal() as db:
        result = await db.execute(select(Finding).where(Finding.id == finding_id))
        finding = result.scalar_one_or_none()
        if not finding:
            return

        prompt = f"""Translate this cybersecurity finding:

Title: {finding.title}
Severity: {finding.severity.value}
Category: {finding.category.value}
CVE ID: {finding.cve_id or 'N/A'}
CVSS Score: {finding.cvss_score or 'N/A'}
Technical Description: {finding.technical_description or 'Not provided'}
Affected Component: {finding.affected_component or 'Unknown'}

Provide the JSON output as specified."""

        raw = await _call_llm(prompt)
        if not raw:
            return

        try:
            data = json.loads(raw)
            finding.ai_summary = data.get("summary")
            finding.ai_business_impact = data.get("business_impact")
            finding.ai_impact_level = data.get("impact_level")
            finding.ai_remediation_steps = data.get("remediation_steps")
            finding.ai_fix_priority = data.get("fix_priority")
            finding.ai_generated_at = datetime.utcnow()
            await db.commit()
            logger.info(f"AI translation complete for finding {finding_id}")
        except (json.JSONDecodeError, KeyError) as e:
            logger.error(f"Failed to parse AI response for finding {finding_id}: {e}")


async def answer_finding_question(finding: Finding, question: str) -> str:
    """AI Security Coach: answer a specific question about a finding."""
    prompt = f"""A security professional is asking about this finding:

Title: {finding.title}
Summary: {finding.ai_summary or finding.technical_description}
Business Impact: {finding.ai_business_impact or 'See technical description'}
Category: {finding.category.value}

Their question: {question}

Answer in 2-4 sentences. Be specific to this finding. Use plain English."""

    system = "You are TrustOS AI Security Coach. Answer questions about specific security findings clearly and directly. Do not use CVE IDs or CVSS in your answers."

    from app.core.config import settings
    try:
        if settings.AI_PROVIDER == "openai" and settings.OPENAI_API_KEY:
            from openai import AsyncOpenAI
            client = AsyncOpenAI(api_key=settings.OPENAI_API_KEY)
            resp = await client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.5,
            )
            return resp.choices[0].message.content
    except Exception as e:
        logger.error(f"AI coach call failed: {e}")

    return "AI explanation is not available. Please review the technical description and remediation steps."


async def generate_attack_path_narrative(finding_id: str):
    """Generate an AI-written attack path narrative for a finding."""
    async with AsyncSessionLocal() as db:
        result = await db.execute(select(Finding).where(Finding.id == finding_id))
        finding = result.scalar_one_or_none()
        if not finding:
            return

        prompt = f"""Create an attack path for this vulnerability:

Title: {finding.title}
Summary: {finding.ai_summary or finding.technical_description}
Category: {finding.category.value}
Severity: {finding.severity.value}

Provide:
1. A plain-English narrative (2-3 sentences): how an attacker could exploit this path from the internet to sensitive data
2. A JSON list of nodes: [{{"id": "1", "label": "Internet", "type": "attacker", "risk_level": "none"}}, ...]
   - types: attacker, entry_point, pivot, target
   - risk_level: none, low, medium, high, critical
3. A JSON list of edges: [{{"source": "1", "target": "2"}}, ...]

Output JSON:
{{
  "narrative": "...",
  "nodes": [...],
  "edges": [...]
}}"""

        raw = await _call_llm(prompt)
        if not raw:
            return

        try:
            data = json.loads(raw)
            path = AttackPath(
                finding_id=finding_id,
                title=f"Attack path: {finding.title}",
                ai_narrative=data.get("narrative"),
                nodes_json=json.dumps(data.get("nodes", [])),
                edges_json=json.dumps(data.get("edges", [])),
            )
            db.add(path)
            await db.commit()
        except Exception as e:
            logger.error(f"Attack path generation failed for {finding_id}: {e}")