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>
This commit is contained in:
drjones
2026-07-07 05:15:52 +00:00
parent 5e22c83919
commit 989c00e5fb
5 changed files with 418 additions and 144 deletions

View File

@@ -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}