Files
trustos/backend/app/api/routes/findings.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

170 lines
6.2 KiB
Python

from fastapi import APIRouter, Depends, HTTPException, Query
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"])
@router.get("", response_model=List[FindingOut])
async def list_findings(
tenant_id: str = Query(...),
severity: Optional[str] = None,
status: Optional[str] = None,
category: Optional[str] = None,
limit: int = 50,
offset: int = 0,
payload: dict = Depends(require_executive_or_above),
db: AsyncSession = Depends(get_db),
):
if payload.get("role") != "trustos_admin" and payload.get("tenant_id") != tenant_id:
raise HTTPException(status_code=403, detail="Access denied")
query = select(Finding).where(Finding.tenant_id == tenant_id)
if severity:
query = query.where(Finding.severity == severity)
if status:
query = query.where(Finding.status == status)
if category:
query = query.where(Finding.category == category)
query = query.order_by(desc(Finding.created_at)).limit(limit).offset(offset)
result = await db.execute(query)
return result.scalars().all()
@router.get("/{finding_id}", response_model=FindingOut)
async def get_finding(
finding_id: str,
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")
return finding
@router.post("", response_model=FindingOut, status_code=201)
async def create_finding(
tenant_id: str = Query(...),
body: FindingCreate = ...,
payload: dict = Depends(require_it_or_above),
db: AsyncSession = Depends(get_db),
):
if payload.get("role") != "trustos_admin" and payload.get("tenant_id") != tenant_id:
raise HTTPException(status_code=403, detail="Access denied")
finding = Finding(tenant_id=tenant_id, **body.model_dump())
db.add(finding)
await db.commit()
await db.refresh(finding)
return finding
@router.patch("/{finding_id}/status", response_model=FindingOut)
async def update_finding_status(
finding_id: str,
body: FindingStatusUpdate,
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")
finding.status = body.status
if body.resolution_note:
finding.resolution_note = body.resolution_note
if body.assignee_email:
finding.assignee_email = body.assignee_email
if body.due_date:
finding.due_date = body.due_date
if body.status == FindingStatus.resolved:
finding.resolved_at = datetime.utcnow()
if body.status == FindingStatus.verified:
finding.verified_at = datetime.utcnow()
if not body.resolution_note:
raise HTTPException(status_code=400, detail="A resolution note is required to verify a finding")
finding.updated_at = datetime.utcnow()
await db.commit()
await db.refresh(finding)
# Trigger async risk score recalculation (fire and forget)
from app.services.risk_calculator import recalculate_risk_score
import asyncio
asyncio.create_task(recalculate_risk_score(finding.tenant_id))
return finding
@router.patch("/{finding_id}/top-risk", response_model=FindingOut)
async def toggle_top_risk(
finding_id: str,
is_top_risk: bool = Query(...),
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="Not found")
finding.is_top_risk = is_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}