trustos/backend/app/api/routes/reports.py

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
from datetime import datetime
import json

from app.db.session import get_db
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, require_executive_or_above, require_it_or_above


def _check_tenant_access(payload: dict, tenant_id: str):
    if payload.get("role") != "trustos_admin" and payload.get("tenant_id") != tenant_id:
        raise HTTPException(status_code=403, detail="Access denied")

router = APIRouter(prefix="/audit-reports", tags=["audit-reports"])


@router.get("", response_model=List[AuditReportOut])
async def list_reports(
    tenant_id: str = Query(...),
    payload: dict = Depends(require_executive_or_above),
    db: AsyncSession = Depends(get_db),
):
    _check_tenant_access(payload, tenant_id)
    result = await db.execute(
        select(AuditReport)
        .where(AuditReport.tenant_id == tenant_id)
        .order_by(desc(AuditReport.report_date))
    )
    return result.scalars().all()


@router.post("/generate", response_model=AuditReportOut, status_code=201)
async def generate_audit_report(
    tenant_id: str = Query(...),
    body: AuditReportCreate = ...,
    payload: dict = Depends(require_admin),
    db: AsyncSession = Depends(get_db),
):
    """Generate a Vault Audit Report snapshot for a tenant."""
    # Get current score
    score_result = await db.execute(
        select(RiskScore)
        .where(RiskScore.tenant_id == tenant_id)
        .order_by(desc(RiskScore.score_date))
        .limit(1)
    )
    latest_score = score_result.scalar_one_or_none()

    # Snapshot top findings
    findings_result = await db.execute(
        select(Finding)
        .where(
            Finding.tenant_id == tenant_id,
            Finding.status.in_([FindingStatus.open, FindingStatus.in_progress])
        )
        .order_by(Finding.created_at)
        .limit(20)
    )
    top_findings = findings_result.scalars().all()
    findings_snapshot = [
        {
            "id": f.id,
            "title": f.title,
            "severity": f.severity.value,
            "category": f.category.value,
            "ai_summary": f.ai_summary,
            "ai_impact_level": f.ai_impact_level,
        }
        for f in top_findings
    ]

    report = AuditReport(
        tenant_id=tenant_id,
        title=body.title,
        report_date=datetime.utcnow(),
        baseline_score=latest_score.overall_score if latest_score else None,
        executive_summary=body.executive_summary,
        scope_description=body.scope_description,
        key_findings_json=json.dumps(findings_snapshot),
        is_baseline=True,
        generated_by=payload.get("sub"),
    )
    db.add(report)
    await db.commit()
    await db.refresh(report)

    return report


@router.get("/{report_id}", response_model=AuditReportOut)
async def get_report(
    report_id: str,
    payload: dict = Depends(require_executive_or_above),
    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")
    _check_tenant_access(payload, report.tenant_id)
    return report


@router.get("/{report_id}/pdf")
async def download_report_pdf(
    report_id: str,
    payload: dict = Depends(require_executive_or_above),
    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")
    _check_tenant_access(payload, report.tenant_id)

    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_it_or_above),
    db: AsyncSession = Depends(get_db),
):
    """Generate a one-off PDF report for a tenant (not stored as a record)."""
    _check_tenant_access(payload, tenant_id)
    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"},
    )