Files
trustos/backend/app/services/report_generator.py
drjones 989c00e5fb Implement AI translation, attack paths, and PDF reports - Advanced features phase
Features added:
- AI Finding Translation endpoints (POST /findings/{id}/ai-translate)
- AI Security Coach endpoint (POST /findings/{id}/ai-question)
- Attack Path visualization generation (POST /attack-paths/{id}/generate, GET /attack-paths/{id})
- Mock AI implementations for demo mode (no API keys required)
- PDF Report generation and download endpoints
- Report snapshot feature for on-demand PDF generation

Technical improvements:
- Mock translation system for findings and attack paths
- Async task-based AI processing
- Graph-based attack path with nodes and edges
- Professional HTML-to-PDF conversion with WeasyPrint
- Jinja2 templating for report generation

Database updates:
- AttackPath table integrated with mock narrative generation
- AI fields populated via async tasks

Testing:
- All E2E tests verified passing (login, dashboard, findings, all roles)
- AI endpoints tested and working with mock data
- PDF report generation produces valid 18KB+ documents
- Attack path generation creates proper graph structures

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
2026-07-07 05:15:52 +00:00

216 lines
6.7 KiB
Python

"""PDF Report Generator — creates professional security reports."""
from datetime import datetime
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__)
HTML_TEMPLATE = """
<html>
<head>
<meta charset="utf-8">
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
color: #1f2937;
line-height: 1.6;
background: white;
padding: 40px;
}
.header {
border-bottom: 3px solid #3b82d4;
padding-bottom: 20px;
margin-bottom: 30px;
}
.header h1 { font-size: 28px; color: #0f172a; }
.header .meta {
margin-top: 10px;
font-size: 12px;
color: #6b7280;
}
.score-box {
background: linear-gradient(135deg, #3b82d4 0%, #1e40af 100%);
color: white;
padding: 30px;
border-radius: 8px;
margin: 30px 0;
text-align: center;
}
.score-box .number { font-size: 48px; font-weight: bold; }
.score-box .label { font-size: 14px; opacity: 0.9; margin-top: 10px; }
.section {
margin: 40px 0;
page-break-inside: avoid;
}
.section h2 {
font-size: 20px;
color: #0f172a;
border-left: 4px solid #3b82d4;
padding-left: 15px;
margin-bottom: 20px;
}
.finding-card {
border: 1px solid #e5e7eb;
border-radius: 6px;
padding: 20px;
margin-bottom: 15px;
page-break-inside: avoid;
}
.finding-title {
font-size: 16px;
font-weight: 600;
color: #0f172a;
margin-bottom: 10px;
}
.severity-critical { color: #dc2626; background: #fee2e2; }
.severity-high { color: #ea580c; background: #fef3c7; }
.severity-medium { color: #d97706; background: #fef3c7; }
.severity-low { color: #16a34a; background: #dcfce7; }
.severity-badge {
display: inline-block;
padding: 4px 12px;
border-radius: 4px;
font-size: 12px;
font-weight: 600;
margin-bottom: 10px;
}
.finding-desc {
font-size: 13px;
color: #4b5563;
margin: 10px 0;
}
.stats {
display: grid;
grid-template-columns: repeat(4, 1fr);
gap: 20px;
}
.stat-box {
text-align: center;
padding: 15px;
border: 1px solid #e5e7eb;
border-radius: 6px;
}
.stat-number { font-size: 24px; font-weight: bold; color: #3b82d4; }
.stat-label { font-size: 12px; color: #6b7280; margin-top: 5px; }
.footer {
margin-top: 50px;
padding-top: 20px;
border-top: 1px solid #e5e7eb;
font-size: 11px;
color: #9ca3af;
text-align: center;
}
</style>
</head>
<body>
<div class="header">
<h1>{{ tenant_name }} — Cyber Risk Report</h1>
<div class="meta">
<p>Report generated on {{ report_date }}</p>
</div>
</div>
<div class="score-box">
<div class="number">{{ cyber_score }}</div>
<div class="label">Cyber Health Score</div>
</div>
<div class="section">
<h2>Risk Summary</h2>
<div class="stats">
<div class="stat-box">
<div class="stat-number">{{ critical_count }}</div>
<div class="stat-label">Critical</div>
</div>
<div class="stat-box">
<div class="stat-number">{{ high_count }}</div>
<div class="stat-label">High</div>
</div>
<div class="stat-box">
<div class="stat-number">{{ medium_count }}</div>
<div class="stat-label">Medium</div>
</div>
<div class="stat-box">
<div class="stat-number">{{ low_count }}</div>
<div class="stat-label">Low</div>
</div>
</div>
</div>
<div class="section">
<h2>Executive Summary</h2>
<p>{{ summary }}</p>
</div>
<div class="section">
<h2>Findings ({{ findings_count }})</h2>
{% for finding in findings %}
<div class="finding-card">
<div class="finding-title">{{ loop.index }}. {{ finding.title }}</div>
<span class="severity-badge severity-{{ finding.severity }}">{{ finding.severity | upper }}</span>
<div class="finding-desc"><strong>Category:</strong> {{ finding.category }}</div>
{% if finding.ai_summary %}
<div class="finding-desc">{{ finding.ai_summary }}</div>
{% endif %}
</div>
{% endfor %}
</div>
<div class="footer">
<p>This report is confidential and for authorized recipients only.</p>
<p>TrustOS — The AI Operating System for Cyber Resilience</p>
</div>
</body>
</html>
"""
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."""
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")
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.",
}
template = Template(HTML_TEMPLATE)
html_string = template.render(**context)
html = HTML(string=html_string, base_url=".")
pdf_bytes = html.write_pdf()
pdf_io = BytesIO(pdf_bytes)
pdf_io.seek(0)
return pdf_io