diff --git a/backend/app/api/routes/findings.py b/backend/app/api/routes/findings.py index 721d621..a30e026 100644 --- a/backend/app/api/routes/findings.py +++ b/backend/app/api/routes/findings.py @@ -3,11 +3,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, desc from typing import List, Optional from datetime import datetime +from pydantic import BaseModel from app.db.session import get_db from app.models.models import Finding, FindingStatus, FindingSeverity from app.schemas.schemas import FindingOut, FindingCreate, FindingStatusUpdate from app.core.security import require_executive_or_above, require_it_or_above +from app.services.ai_translator import translate_finding_async, answer_finding_question router = APIRouter(prefix="/findings", tags=["findings"]) @@ -125,3 +127,43 @@ async def toggle_top_risk( await db.commit() await db.refresh(finding) return finding + + +@router.post("/{finding_id}/ai-translate") +async def translate_finding( + finding_id: str, + payload: dict = Depends(require_it_or_above), + db: AsyncSession = Depends(get_db), +): + result = await db.execute(select(Finding).where(Finding.id == finding_id)) + finding = result.scalar_one_or_none() + if not finding: + raise HTTPException(status_code=404, detail="Finding not found") + if payload.get("role") != "trustos_admin" and payload.get("tenant_id") != finding.tenant_id: + raise HTTPException(status_code=403, detail="Access denied") + + import asyncio + asyncio.create_task(translate_finding_async(finding_id)) + return {"status": "Translation requested"} + + +class AIQuestionRequest(BaseModel): + question: str + + +@router.post("/{finding_id}/ai-question") +async def ask_ai_about_finding( + finding_id: str, + request: AIQuestionRequest, + payload: dict = Depends(require_executive_or_above), + db: AsyncSession = Depends(get_db), +): + result = await db.execute(select(Finding).where(Finding.id == finding_id)) + finding = result.scalar_one_or_none() + if not finding: + raise HTTPException(status_code=404, detail="Finding not found") + if payload.get("role") != "trustos_admin" and payload.get("tenant_id") != finding.tenant_id: + raise HTTPException(status_code=403, detail="Access denied") + + answer = await answer_finding_question(finding, request.question) + return {"answer": answer} diff --git a/backend/app/api/routes/reports.py b/backend/app/api/routes/reports.py index 9b27c89..b58b898 100644 --- a/backend/app/api/routes/reports.py +++ b/backend/app/api/routes/reports.py @@ -1,4 +1,5 @@ from fastapi import APIRouter, Depends, HTTPException, Query +from fastapi.responses import StreamingResponse from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, desc from typing import List @@ -6,7 +7,7 @@ from datetime import datetime import json from app.db.session import get_db -from app.models.models import AuditReport, Finding, RiskScore, Executive, AuthorizedAsset, FindingStatus +from app.models.models import AuditReport, Finding, RiskScore, Executive, AuthorizedAsset, FindingStatus, Tenant from app.schemas.schemas import AuditReportOut, AuditReportCreate from app.core.security import require_admin @@ -82,11 +83,6 @@ async def generate_audit_report( await db.commit() await db.refresh(report) - # Kick off PDF generation in background - from app.services.report_generator import generate_pdf_for_report - import asyncio - asyncio.create_task(generate_pdf_for_report(report.id)) - return report @@ -101,3 +97,81 @@ async def get_report( if not report: raise HTTPException(status_code=404, detail="Report not found") return report + + +@router.get("/{report_id}/pdf") +async def download_report_pdf( + report_id: str, + payload: dict = Depends(require_admin), + db: AsyncSession = Depends(get_db), +): + result = await db.execute(select(AuditReport).where(AuditReport.id == report_id)) + report = result.scalar_one_or_none() + if not report: + raise HTTPException(status_code=404, detail="Report not found") + + tenant_result = await db.execute(select(Tenant).where(Tenant.id == report.tenant_id)) + tenant = tenant_result.scalar_one_or_none() + + findings_result = await db.execute( + select(Finding).where(Finding.tenant_id == report.tenant_id).order_by(desc(Finding.created_at)) + ) + findings = findings_result.scalars().all() + + score_result = await db.execute( + select(RiskScore).where(RiskScore.tenant_id == report.tenant_id).order_by(desc(RiskScore.score_date)) + ) + scores = score_result.scalars().all() + + from app.services.report_generator import generate_findings_pdf + latest_score = scores[0].overall_score if scores else 0 + pdf_io = await generate_findings_pdf( + tenant_name=tenant.name if tenant else "Unknown", + cyber_score=latest_score, + findings=findings, + risk_scores=scores, + ) + + return StreamingResponse( + iter([pdf_io.getvalue()]), + media_type="application/pdf", + headers={"Content-Disposition": f"attachment; filename=report_{report_id}.pdf"}, + ) + + +@router.post("/{tenant_id}/pdf-snapshot") +async def generate_pdf_snapshot( + tenant_id: str, + payload: dict = Depends(require_admin), + db: AsyncSession = Depends(get_db), +): + """Generate a one-off PDF report for a tenant (not stored as a record).""" + tenant_result = await db.execute(select(Tenant).where(Tenant.id == tenant_id)) + tenant = tenant_result.scalar_one_or_none() + if not tenant: + raise HTTPException(status_code=404, detail="Tenant not found") + + findings_result = await db.execute( + select(Finding).where(Finding.tenant_id == tenant_id).order_by(desc(Finding.created_at)) + ) + findings = findings_result.scalars().all() + + score_result = await db.execute( + select(RiskScore).where(RiskScore.tenant_id == tenant_id).order_by(desc(RiskScore.score_date)) + ) + scores = score_result.scalars().all() + + from app.services.report_generator import generate_findings_pdf + latest_score = scores[0].overall_score if scores else 0 + pdf_io = await generate_findings_pdf( + tenant_name=tenant.name, + cyber_score=latest_score, + findings=findings, + risk_scores=scores, + ) + + return StreamingResponse( + iter([pdf_io.getvalue()]), + media_type="application/pdf", + headers={"Content-Disposition": f"attachment; filename=trustos_report_{tenant_id}.pdf"}, + ) diff --git a/backend/app/services/ai_translator.py b/backend/app/services/ai_translator.py index 1aec079..bc9a3d1 100644 --- a/backend/app/services/ai_translator.py +++ b/backend/app/services/ai_translator.py @@ -37,7 +37,7 @@ 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: + 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( @@ -50,7 +50,7 @@ async def _call_llm(prompt: str) -> Optional[str]: response_format={"type": "json_object"}, ) return resp.choices[0].message.content - elif settings.AI_PROVIDER == "anthropic" and settings.ANTHROPIC_API_KEY: + 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( @@ -61,11 +61,22 @@ async def _call_llm(prompt: str) -> Optional[str]: ) return resp.content[0].text else: - logger.warning("No AI provider configured — skipping translation") - return None + 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}") - return None + 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): @@ -94,14 +105,17 @@ Provide the JSON output as specified.""" 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}") + 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}") @@ -123,7 +137,7 @@ Answer in 2-4 sentences. Be specific to this finding. Use plain English.""" from app.core.config import settings try: - if settings.AI_PROVIDER == "openai" and settings.OPENAI_API_KEY: + 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( @@ -135,10 +149,32 @@ Answer in 2-4 sentences. Be specific to this finding. Use plain English.""" 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 "AI explanation is not available. Please review the technical description and remediation steps." + 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): @@ -149,7 +185,16 @@ async def generate_attack_path_narrative(finding_id: str): if not finding: return - prompt = f"""Create an attack path for this vulnerability: + 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} @@ -169,10 +214,9 @@ Output JSON: "nodes": [...], "edges": [...] }}""" - - raw = await _call_llm(prompt) - if not raw: - return + raw = await _call_llm(prompt) + if not raw: + raw = _generate_mock_attack_path(finding) try: data = json.loads(raw) @@ -185,5 +229,29 @@ Output JSON: ) 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, + }) diff --git a/backend/app/services/report_generator.py b/backend/app/services/report_generator.py index 52889a3..4fa9ccd 100644 --- a/backend/app/services/report_generator.py +++ b/backend/app/services/report_generator.py @@ -1,139 +1,215 @@ -""" -PDF report generator for Vault Audit Reports. -Uses Jinja2 + WeasyPrint to produce branded PDFs. -""" -import os -import json -import logging +"""PDF Report Generator — creates professional security reports.""" from datetime import datetime -from pathlib import Path - -from jinja2 import Environment, PackageLoader, select_autoescape, DictLoader -from app.db.session import AsyncSessionLocal -from app.models.models import AuditReport, Tenant, Finding, FindingStatus -from app.core.config import settings -from sqlalchemy import select, desc +from typing import List, Optional +from jinja2 import Template +from weasyprint import HTML, CSS +from io import BytesIO +from app.models.models import Finding, RiskScore +import logging logger = logging.getLogger(__name__) -REPORT_HTML_TEMPLATE = """ - - +HTML_TEMPLATE = """ + - - + + -
⚠ CONFIDENTIAL — This report contains sensitive security information. Do not distribute without authorization.
-
- -
{{ report.title }}
-
Vault Audit Report · {{ tenant.name }} · Generated {{ report.report_date.strftime('%B %d, %Y') }}
-
+
+

{{ tenant_name }} — Cyber Risk Report

+
+

Report generated on {{ report_date }}

+
+
-
-
{{ report.baseline_score | int }}
-
Cyber Health Score at Audit Date
100 = Optimal · 0 = Critical Risk
-
+
+
{{ cyber_score }}
+
Cyber Health Score
+
-{% if report.executive_summary %} -

Executive Summary

-
{{ report.executive_summary }}
-{% endif %} +
+

Risk Summary

+
+
+
{{ critical_count }}
+
Critical
+
+
+
{{ high_count }}
+
High
+
+
+
{{ medium_count }}
+
Medium
+
+
+
{{ low_count }}
+
Low
+
+
+
-{% if report.scope_description %} -

Scope

-

{{ report.scope_description }}

-{% endif %} +
+

Executive Summary

+

{{ summary }}

+
-

Key Findings

-{% for f in findings %} -
-
{{ f.title }} {{ f.severity | upper }}
- {% if f.ai_summary %} -
{{ f.ai_summary }}
- {% endif %} - {% if f.ai_business_impact %} -
Business Impact: {{ f.ai_business_impact }}
- {% endif %} -
-{% endfor %} +
+

Findings ({{ findings_count }})

+ {% for finding in findings %} +
+
{{ loop.index }}. {{ finding.title }}
+ {{ finding.severity | upper }} +
Category: {{ finding.category }}
+ {% if finding.ai_summary %} +
{{ finding.ai_summary }}
+ {% endif %} +
+ {% endfor %} +
- + """ -async def generate_pdf_for_report(report_id: str): - """Generate a branded PDF for a Vault Audit Report and store the path.""" - async with AsyncSessionLocal() as db: - try: - r_result = await db.execute(select(AuditReport).where(AuditReport.id == report_id)) - report = r_result.scalar_one_or_none() - if not report: - return +async def generate_findings_pdf( + tenant_name: str, + cyber_score: float, + findings: List[Finding], + risk_scores: Optional[List[RiskScore]] = None, +) -> BytesIO: + """Generate a professional PDF report of security findings.""" - t_result = await db.execute(select(Tenant).where(Tenant.id == report.tenant_id)) - tenant = t_result.scalar_one_or_none() + critical = sum(1 for f in findings if f.severity.value == "critical") + high = sum(1 for f in findings if f.severity.value == "high") + medium = sum(1 for f in findings if f.severity.value == "medium") + low = sum(1 for f in findings if f.severity.value == "low") - # Get findings snapshot - f_result = await db.execute( - select(Finding) - .where( - Finding.tenant_id == report.tenant_id, - Finding.status.in_([FindingStatus.open, FindingStatus.in_progress]) - ) - .order_by(Finding.created_at) - .limit(20) - ) - findings = f_result.scalars().all() + context = { + "tenant_name": tenant_name, + "cyber_score": round(cyber_score, 1), + "report_date": datetime.utcnow().strftime("%B %d, %Y"), + "critical_count": critical, + "high_count": high, + "medium_count": medium, + "low_count": low, + "findings_count": len(findings), + "findings": [ + { + "title": f.title, + "severity": f.severity.value, + "category": f.category.value, + "ai_summary": f.ai_summary, + } + for f in findings + ], + "summary": f"This report contains {len(findings)} security findings affecting {tenant_name}, with {critical} critical issues requiring immediate attention.", + } - # Render HTML - env = Environment(loader=DictLoader({"report.html": REPORT_HTML_TEMPLATE})) - template = env.get_template("report.html") - html = template.render(report=report, tenant=tenant, findings=findings) + template = Template(HTML_TEMPLATE) + html_string = template.render(**context) - # Write PDF - storage = Path(settings.STORAGE_PATH) / "reports" - storage.mkdir(parents=True, exist_ok=True) - pdf_path = storage / f"vault-audit-{report_id}.pdf" + html = HTML(string=html_string, base_url=".") + pdf_bytes = html.write_pdf() - from weasyprint import HTML as WH - WH(string=html).write_pdf(str(pdf_path)) - - report.pdf_path = str(pdf_path) - await db.commit() - logger.info(f"PDF generated: {pdf_path}") - - except Exception as e: - logger.error(f"PDF generation failed for report {report_id}: {e}") + pdf_io = BytesIO(pdf_bytes) + pdf_io.seek(0) + return pdf_io diff --git a/backend/test_db.py b/backend/test_db.py new file mode 100644 index 0000000..8dbe33b --- /dev/null +++ b/backend/test_db.py @@ -0,0 +1,14 @@ +import asyncio +from app.db.session import AsyncSessionLocal +from app.models.models import User +from sqlalchemy import select + +async def test(): + db = AsyncSessionLocal() + result = await db.execute(select(User)) + users = result.scalars().all() + print(f'Found {len(users)} users') + for u in users: + print(f' - {u.email}: {u.role.value}') + +asyncio.run(test())