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 and not settings.OPENAI_API_KEY.startswith("sk-..."):
            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 and not settings.ANTHROPIC_API_KEY.startswith("sk-ant-"):
            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.info("No valid AI provider configured — using mock translation")
            return _generate_mock_translation(prompt)
    except Exception as e:
        logger.error(f"LLM call failed: {e}, using mock translation")
        return _generate_mock_translation(prompt)


def _generate_mock_translation(prompt: str) -> str:
    """Generate a mock AI translation for demo purposes."""
    return json.dumps({
        "summary": "Security vulnerability detected in system component",
        "business_impact": "Unauthorized access or data breach potential if exploited by attackers",
        "impact_level": "High",
        "remediation_steps": "1. Patch the affected component to latest version 2. Deploy patch during maintenance window 3. Verify patch application 4. Monitor logs for suspicious activity 5. Conduct security scan to confirm fix",
        "fix_priority": "soon"
    })


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)
            if "summary" in data:
                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}")
            else:
                logger.warning(f"Invalid AI response format 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 and not settings.OPENAI_API_KEY.startswith("sk-..."):
            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
        elif settings.AI_PROVIDER == "anthropic" and settings.ANTHROPIC_API_KEY and not settings.ANTHROPIC_API_KEY.startswith("sk-ant-"):
            from anthropic import AsyncAnthropic
            client = AsyncAnthropic(api_key=settings.ANTHROPIC_API_KEY)
            resp = await client.messages.create(
                model="claude-3-haiku-20240307",
                max_tokens=256,
                system=system,
                messages=[{"role": "user", "content": prompt}],
            )
            return resp.content[0].text
    except Exception as e:
        logger.error(f"AI coach call failed: {e}")

    return f"Based on this {finding.category.value} issue, {_generate_mock_question_answer(finding, question)}"


def _generate_mock_question_answer(finding: Finding, question: str) -> str:
    """Generate mock AI response to questions about findings."""
    if "risk" in question.lower() or "impact" in question.lower():
        return finding.ai_business_impact or "This finding could allow attackers to compromise system integrity."
    elif "fix" in question.lower() or "remediate" in question.lower() or "resolve" in question.lower():
        return finding.ai_remediation_steps or "Follow the listed remediation steps to address this issue."
    elif "timeline" in question.lower() or "urgent" in question.lower() or "priority" in question.lower():
        return f"This {finding.severity.value}-severity issue should be addressed as soon as possible."
    else:
        return "Review the finding details above for comprehensive information about this security issue."


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

        from app.core.config import settings
        use_mock = not (
            (settings.AI_PROVIDER == "openai" and settings.OPENAI_API_KEY and not settings.OPENAI_API_KEY.startswith("sk-...")) or
            (settings.AI_PROVIDER == "anthropic" and settings.ANTHROPIC_API_KEY and not settings.ANTHROPIC_API_KEY.startswith("sk-ant-"))
        )

        if use_mock:
            raw = _generate_mock_attack_path(finding)
        else:
            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:
                raw = _generate_mock_attack_path(finding)

        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()
            logger.info(f"Attack path generated for finding {finding_id}")
        except Exception as e:
            logger.error(f"Attack path generation failed for {finding_id}: {e}")


def _generate_mock_attack_path(finding: Finding) -> str:
    """Generate a mock attack path for demo purposes."""
    nodes = [
        {"id": "1", "label": "Internet", "type": "attacker", "risk_level": "none"},
        {"id": "2", "label": "Public Endpoint", "type": "entry_point", "risk_level": "critical"},
        {"id": "3", "label": "Web Server", "type": "pivot", "risk_level": "high"},
        {"id": "4", "label": "Database", "type": "target", "risk_level": "critical"},
    ]
    edges = [
        {"source": "1", "target": "2"},
        {"source": "2", "target": "3"},
        {"source": "3", "target": "4"},
    ]

    narrative = f"An attacker from the internet discovers the exposed entry point in your {finding.category.value} infrastructure. They exploit the vulnerability to pivot through your web tier and ultimately access sensitive data in your backend database."

    return json.dumps({
        "narrative": narrative,
        "nodes": nodes,
        "edges": edges,
    })