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>
170 lines
6.2 KiB
Python
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}
|