| 1 | import calendar |
| 2 | import csv |
| 3 | from datetime import datetime |
| 4 | from io import StringIO |
| 5 | from typing import Any |
| 6 | from typing import Dict |
| 7 | from typing import List |
| 8 | from typing import Optional |
| 9 | |
| 10 | from fastapi import APIRouter |
| 11 | from fastapi import Depends |
| 12 | from fastapi import HTTPException |
| 13 | from fastapi import Query |
| 14 | from fastapi import Security |
| 15 | from fastapi.responses import FileResponse |
| 16 | from fastapi.responses import StreamingResponse |
| 17 | from sqlalchemy.ext.asyncio import AsyncSession |
| 18 | from sqlalchemy.orm import selectinload |
| 19 | from sqlmodel import select |
| 20 | |
| 21 | from app.auth.routes.auth import AuthHandler |
| 22 | from app.customers.routes.customers import get_customer |
| 23 | from app.db.db_session import get_db |
| 24 | from app.incidents.models import Alert |
| 25 | from app.incidents.models import AlertToIoC |
| 26 | from app.incidents.models import AlertToTag |
| 27 | from app.incidents.models import Asset |
| 28 | from app.incidents.models import Case |
| 29 | from app.incidents.models import CaseAlertLink |
| 30 | from app.incidents.schema.db_operations import CaseDownloadDocxRequest |
| 31 | from app.incidents.services.reports import cleanup_temp_files |
| 32 | from app.incidents.services.reports import create_case_context |
| 33 | from app.incidents.services.reports import create_file_response |
| 34 | from app.incidents.services.reports import download_template |
| 35 | from app.incidents.services.reports import render_document_with_context |
| 36 | from app.incidents.services.reports import save_template_to_tempfile |
| 37 | from app.incidents.services.reports_pdf import convert_html_to_pdf |
| 38 | from app.incidents.services.reports_pdf import create_case_context_pdf |
| 39 | from app.incidents.services.reports_pdf import create_file_response_pdf |
| 40 | from app.incidents.services.reports_pdf import download_template_pdf |
| 41 | from app.incidents.services.reports_pdf import render_html_template |
| 42 | |
| 43 | incidents_report_router = APIRouter() |
| 44 | |
| 45 | # Constants |
| 46 | FIELDNAMES = [ |
| 47 | "Case ID", |
| 48 | "Case Name", |
| 49 | "Case Description", |
| 50 | "Case Creation Time", |
| 51 | "Case Status", |
| 52 | "Case Closed Time", |
| 53 | "Case Assigned To", |
| 54 | "Case Customer Code", |
| 55 | "Alert ID", |
| 56 | "Alert Name", |
| 57 | "Alert Description", |
| 58 | "Alert Status", |
| 59 | "Alert Creation Time", |
| 60 | "Alert Time Closed", |
| 61 | "Alert Source", |
| 62 | "Alert Assigned To", |
| 63 | "Assets", |
| 64 | "Tags", |
| 65 | "Comments", |
| 66 | "Customer Code", |
| 67 | ] |
| 68 | |
| 69 | |
| 70 | # Helper Functions |
| 71 | def _build_month_range(year: int, month: int) -> tuple[datetime, datetime]: |
| 72 | """Return (start, end) datetimes for a given year/month.""" |
| 73 | last_day = calendar.monthrange(year, month)[1] |
| 74 | start = datetime(year, month, 1) |
| 75 | end = datetime(year, month, last_day, 23, 59, 59) |
| 76 | return start, end |
| 77 | |
| 78 | |
| 79 | async def fetch_cases_with_related_data( |
| 80 | session: AsyncSession, |
| 81 | year: Optional[int] = None, |
| 82 | month: Optional[int] = None, |
| 83 | ) -> List[Case]: |
| 84 | """Fetch cases with related alerts, assets, tags, and comments.""" |
| 85 | query = select(Case).options( |
| 86 | selectinload(Case.alerts) |
| 87 | .selectinload(CaseAlertLink.alert) |
| 88 | .options(selectinload(Alert.assets), selectinload(Alert.tags).selectinload(AlertToTag.tag), selectinload(Alert.comments)), |
| 89 | ) |
| 90 | if year and month: |
| 91 | start, end = _build_month_range(year, month) |
| 92 | query = query.where(Case.case_creation_time >= start, Case.case_creation_time <= end) |
| 93 | result = await session.execute(query) |
| 94 | return result.scalars().all() |
| 95 | |
| 96 | |
| 97 | async def fetch_cases_by_customer( |
| 98 | session: AsyncSession, |
| 99 | customer_code: str, |
| 100 | year: Optional[int] = None, |
| 101 | month: Optional[int] = None, |
| 102 | ) -> List[Case]: |
| 103 | """Fetch cases for a specific customer with related data.""" |
| 104 | query = ( |
| 105 | select(Case) |
| 106 | .where(Case.customer_code == customer_code) |
| 107 | .options( |
| 108 | selectinload(Case.alerts) |
| 109 | .selectinload(CaseAlertLink.alert) |
| 110 | .options(selectinload(Alert.assets), selectinload(Alert.tags).selectinload(AlertToTag.tag), selectinload(Alert.comments)), |
| 111 | ) |
| 112 | ) |
| 113 | if year and month: |
| 114 | start, end = _build_month_range(year, month) |
| 115 | query = query.where(Case.case_creation_time >= start, Case.case_creation_time <= end) |
| 116 | result = await session.execute(query) |
| 117 | return result.scalars().all() |
| 118 | |
| 119 | |
| 120 | async def fetch_case_by_id(session: AsyncSession, case_id: int) -> Case: |
| 121 | """Fetch a case by its ID.""" |
| 122 | result = await session.execute( |
| 123 | select(Case) |
| 124 | .where(Case.id == case_id) |
| 125 | .options( |
| 126 | selectinload(Case.alerts) |
| 127 | .selectinload(CaseAlertLink.alert) |
| 128 | .options( |
| 129 | selectinload(Alert.assets).selectinload(Asset.alert_context), # Load alert_context |
| 130 | selectinload(Alert.tags).selectinload(AlertToTag.tag), # Load tags |
| 131 | selectinload(Alert.comments), # Load comments |
| 132 | selectinload(Alert.iocs).selectinload(AlertToIoC.ioc), # Load IoCs |
| 133 | ), |
| 134 | ), |
| 135 | ) |
| 136 | return result.scalars().first() |
| 137 | |
| 138 | |
| 139 | def serialize_case_alert_to_row(case: Case, alert: Alert) -> Dict[str, Any]: |
| 140 | """Serialize a case and its alert into a CSV row.""" |
| 141 | assets = ";".join(asset.asset_name for asset in alert.assets or []) |
| 142 | tags = ";".join(alert_tag.tag.tag for alert_tag in alert.tags or []) |
| 143 | comments = ";".join(comment.comment for comment in alert.comments or []) |
| 144 | |
| 145 | return { |
| 146 | "Case ID": case.id, |
| 147 | "Case Name": case.case_name, |
| 148 | "Case Description": case.case_description, |
| 149 | "Case Creation Time": case.case_creation_time.strftime("%Y-%m-%d %H:%M:%S"), |
| 150 | "Case Status": case.case_status, |
| 151 | "Case Closed Time": case.case_closed_time.strftime("%Y-%m-%d %H:%M:%S") if case.case_closed_time else "", |
| 152 | "Case Assigned To": case.assigned_to or "", |
| 153 | "Case Customer Code": case.customer_code or "", |
| 154 | "Alert ID": alert.id, |
| 155 | "Alert Name": alert.alert_name, |
| 156 | "Alert Description": alert.alert_description, |
| 157 | "Alert Status": alert.status, |
| 158 | "Alert Creation Time": alert.alert_creation_time.strftime("%Y-%m-%d %H:%M:%S"), |
| 159 | "Alert Time Closed": alert.time_closed.strftime("%Y-%m-%d %H:%M:%S") if alert.time_closed else "", |
| 160 | "Alert Source": alert.source, |
| 161 | "Alert Assigned To": alert.assigned_to or "", |
| 162 | "Assets": assets, |
| 163 | "Tags": tags, |
| 164 | "Comments": comments, |
| 165 | "Customer Code": alert.customer_code, |
| 166 | } |
| 167 | |
| 168 | |
| 169 | def generate_csv_content(rows: List[Dict[str, Any]]) -> StringIO: |
| 170 | """Generate CSV content from rows.""" |
| 171 | csv_stream = StringIO() |
| 172 | writer = csv.DictWriter(csv_stream, fieldnames=FIELDNAMES, lineterminator="\n") |
| 173 | writer.writeheader() |
| 174 | writer.writerows(rows) |
| 175 | csv_stream.seek(0) |
| 176 | return csv_stream |
| 177 | |
| 178 | |
| 179 | # Route Handler |
| 180 | @incidents_report_router.post( |
| 181 | "/generate-report-csv", |
| 182 | description="Generate a report for all cases. Optionally filter by year and month.", |
| 183 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 184 | ) |
| 185 | async def get_cases_export_all_route( |
| 186 | year: Optional[int] = Query(None, description="Filter by year (e.g. 2026)", ge=2000, le=2100), |
| 187 | month: Optional[int] = Query(None, description="Filter by month (1-12)", ge=1, le=12), |
| 188 | session: AsyncSession = Depends(get_db), |
| 189 | ) -> StreamingResponse: |
| 190 | if (year is None) != (month is None): |
| 191 | raise HTTPException(status_code=400, detail="Both year and month must be provided together") |
| 192 | cases = await fetch_cases_with_related_data(session, year=year, month=month) |
| 193 | rows = [serialize_case_alert_to_row(case, alert_link.alert) for case in cases for alert_link in case.alerts] |
| 194 | csv_stream = generate_csv_content(rows) |
| 195 | month_suffix = f"_{year}-{month:02d}" if year and month else "" |
| 196 | filename = f"cases_export{month_suffix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv" |
| 197 | response = StreamingResponse(csv_stream, media_type="text/csv") |
| 198 | response.headers["Content-Disposition"] = f"attachment; filename={filename}" |
| 199 | return response |
| 200 | |
| 201 | |
| 202 | @incidents_report_router.post( |
| 203 | "/generate-report-csv/{customer_code}", |
| 204 | description="Generate a report for a customer. Optionally filter by year and month.", |
| 205 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 206 | ) |
| 207 | async def get_cases_export_customer_route( |
| 208 | customer_code: str, |
| 209 | year: Optional[int] = Query(None, description="Filter by year (e.g. 2026)", ge=2000, le=2100), |
| 210 | month: Optional[int] = Query(None, description="Filter by month (1-12)", ge=1, le=12), |
| 211 | session: AsyncSession = Depends(get_db), |
| 212 | ) -> StreamingResponse: |
| 213 | if (year is None) != (month is None): |
| 214 | raise HTTPException(status_code=400, detail="Both year and month must be provided together") |
| 215 | await get_customer(customer_code=customer_code, session=session) |
| 216 | cases = await fetch_cases_by_customer(session, customer_code, year=year, month=month) |
| 217 | rows = [serialize_case_alert_to_row(case, alert_link.alert) for case in cases for alert_link in case.alerts] |
| 218 | csv_stream = generate_csv_content(rows) |
| 219 | month_suffix = f"_{year}-{month:02d}" if year and month else "" |
| 220 | filename = f"cases_export_{customer_code}{month_suffix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv" |
| 221 | response = StreamingResponse(csv_stream, media_type="text/csv") |
| 222 | response.headers["Content-Disposition"] = f"attachment; filename={filename}" |
| 223 | return response |
| 224 | |
| 225 | |
| 226 | @incidents_report_router.post( |
| 227 | "/generate-report-docx", |
| 228 | description="Generate a docx report for a case.", |
| 229 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 230 | ) |
| 231 | async def get_cases_export_docx_route( |
| 232 | request: CaseDownloadDocxRequest, |
| 233 | session: AsyncSession = Depends(get_db), |
| 234 | ) -> FileResponse: |
| 235 | case = await fetch_case_by_id(session, request.case_id) |
| 236 | if not case: |
| 237 | raise HTTPException(status_code=404, detail="No cases found") |
| 238 | |
| 239 | context = create_case_context(case) |
| 240 | |
| 241 | template_file_content = await download_template(request.template_name) |
| 242 | tmp_template_name = save_template_to_tempfile(template_file_content) |
| 243 | |
| 244 | rendered_file_name = render_document_with_context(tmp_template_name, context) |
| 245 | |
| 246 | response = create_file_response(file_path=rendered_file_name, file_name=request.file_name) |
| 247 | |
| 248 | # Clean up temporary files |
| 249 | cleanup_temp_files([tmp_template_name]) |
| 250 | |
| 251 | return response |
| 252 | |
| 253 | |
| 254 | @incidents_report_router.post( |
| 255 | "/generate-report-pdf", |
| 256 | description="Generate a PDF report for a case.", |
| 257 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 258 | ) |
| 259 | async def get_cases_export_pdf_route( |
| 260 | request: CaseDownloadDocxRequest, |
| 261 | session: AsyncSession = Depends(get_db), |
| 262 | ) -> FileResponse: |
| 263 | case = await fetch_case_by_id(session, request.case_id) |
| 264 | if not case: |
| 265 | raise HTTPException(status_code=404, detail="No cases found") |
| 266 | |
| 267 | context = create_case_context_pdf(case) |
| 268 | |
| 269 | # Download and save the template |
| 270 | tmp_template_name = await download_template_pdf(request.template_name) |
| 271 | |
| 272 | # Render the HTML template with the context |
| 273 | rendered_html_file_name = render_html_template(tmp_template_name, context) |
| 274 | |
| 275 | # Convert HTML to PDF using WeasyPrint |
| 276 | rendered_pdf_file_name = convert_html_to_pdf(rendered_html_file_name) |
| 277 | |
| 278 | # Create the FileResponse for PDF |
| 279 | response = create_file_response_pdf(file_path=rendered_pdf_file_name, file_name=request.file_name.replace(".docx", ".pdf")) |
| 280 | |
| 281 | # Clean up temporary files |
| 282 | cleanup_temp_files([tmp_template_name, rendered_html_file_name]) |
| 283 | |
| 284 | return response |
| 285 | |
| 286 | |
| 287 | # ! TODO: ROUTE FOR MARKDOWN TEMPLATE ! # |