"""Bounded parser and shared upload contract for asset bulk mutations.""" from __future__ import annotations import csv import io import logging import zipfile from collections.abc import Iterable from pathlib import Path from typing import Any from fastapi import UploadFile from openpyxl import load_workbook from app.core.exceptions import AppException logger = logging.getLogger(__name__) MAX_UPLOAD_BYTES = 5 * 1024 * 1024 MAX_BULK_ROWS = 2_000 MAX_BULK_COLUMNS = 100 MAX_CELL_STRING_LENGTH = 32_767 # XLSX is a ZIP container. The compressed upload cap alone does not bound # what openpyxl has to allocate while it loads workbook metadata/shared # strings, so reject suspicious archives before handing them to openpyxl. MAX_XLSX_ZIP_MEMBERS = 1_000 MAX_XLSX_UNCOMPRESSED_BYTES = 32 * 1024 * 1024 MAX_XLSX_MEMBER_BYTES = 16 * 1024 * 1024 MAX_XLSX_COMPRESSION_RATIO = 200.0 ASSET_BULK_ALLOWED_EXTENSIONS = frozenset({".csv", ".xlsx"}) _ALLOWED_MIMES_BY_EXTENSION = { ".csv": frozenset( { "text/csv", "application/csv", # Excel commonly labels an exported CSV with its legacy MIME. "application/vnd.ms-excel", "application/octet-stream", } ), ".xlsx": frozenset( { "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/octet-stream", } ), } def _invalid_file(message: str) -> AppException: return AppException( code="INVALID_ASSET_BULK_FILE", message=message, status_code=400, ) def _limit_exceeded(message: str) -> AppException: return AppException( code="ASSET_BULK_IMPORT_LIMIT_EXCEEDED", message=message, status_code=413, ) def validate_bulk_upload_contract(upload: UploadFile) -> str: """Validate the format contract shared by intake and bulk update.""" extension = Path(upload.filename or "").suffix.lower() content_type = (upload.content_type or "").lower() allowed_mimes = _ALLOWED_MIMES_BY_EXTENSION.get(extension) if allowed_mimes is None or content_type not in allowed_mimes: raise AppException( code="UNSUPPORTED_ASSET_BULK_FILE", message="Unsupported file type; only .csv and .xlsx are accepted", status_code=415, ) return extension def _collect_bounded_rows[RowT](rows: Iterable[RowT]) -> list[RowT]: """Collect no more than 2,000 rows, consuming one sentinel row at most.""" collected: list[RowT] = [] for row_number, row in enumerate(rows, start=1): if row_number > MAX_BULK_ROWS: raise _limit_exceeded(f"Too many rows; maximum {MAX_BULK_ROWS} per upload") collected.append(row) return collected def _validate_cell(value: Any) -> None: if isinstance(value, str) and len(value) > MAX_CELL_STRING_LENGTH: raise _limit_exceeded( f"Cell text exceeds the maximum length of {MAX_CELL_STRING_LENGTH} characters" ) def _validate_row_width(row: list[Any] | tuple[Any, ...]) -> None: if len(row) > MAX_BULK_COLUMNS: raise _limit_exceeded(f"Too many columns; maximum {MAX_BULK_COLUMNS} per upload") for value in row: _validate_cell(value) def _preflight_xlsx_archive(payload: bytes) -> None: try: with zipfile.ZipFile(io.BytesIO(payload)) as archive: members = [info for info in archive.infolist() if not info.is_dir()] except (zipfile.BadZipFile, OSError) as exc: raise _invalid_file("Invalid XLSX archive") from exc if len(members) > MAX_XLSX_ZIP_MEMBERS: raise _limit_exceeded(f"XLSX archive has too many members; maximum {MAX_XLSX_ZIP_MEMBERS}") total_uncompressed = 0 for member in members: if member.flag_bits & 0x1: raise _invalid_file("Encrypted XLSX archives are not supported") if member.file_size > MAX_XLSX_MEMBER_BYTES: raise _limit_exceeded("XLSX archive member is too large") total_uncompressed += member.file_size if total_uncompressed > MAX_XLSX_UNCOMPRESSED_BYTES: raise _limit_exceeded("XLSX uncompressed size is too large") if member.file_size: ratio = member.file_size / max(member.compress_size, 1) if ratio > MAX_XLSX_COMPRESSION_RATIO: raise _limit_exceeded("XLSX compression ratio is suspiciously high") def _parse_csv(payload: bytes) -> list[dict[str, Any]]: try: text = payload.decode("utf-8-sig") except UnicodeDecodeError as exc: raise _invalid_file("CSV must be UTF-8 encoded") from exc try: reader = csv.reader(io.StringIO(text)) headers = next(reader, None) if headers is None: return [] _validate_row_width(headers) def data_rows() -> Iterable[dict[str, Any]]: for row in reader: # Match ``csv.DictReader``'s existing blank-line behaviour. if not row: continue _validate_row_width(row) yield { headers[index]: row[index] for index in range(min(len(headers), len(row))) if headers[index] } return _collect_bounded_rows(data_rows()) except csv.Error as exc: raise _invalid_file("Malformed CSV file") from exc def _trim_trailing_empty_cells(row: tuple[Any, ...]) -> tuple[Any, ...]: last = len(row) while last and row[last - 1] is None: last -= 1 return row[:last] def _parse_xlsx(payload: bytes) -> list[dict[str, Any]]: _preflight_xlsx_archive(payload) workbook = None try: workbook = load_workbook( io.BytesIO(payload), read_only=True, data_only=True, ) worksheet = workbook.active if worksheet.max_column and worksheet.max_column > MAX_BULK_COLUMNS: raise _limit_exceeded(f"Too many columns; maximum {MAX_BULK_COLUMNS} per upload") # One sentinel column catches a populated 101st cell without asking # openpyxl to materialise an attacker-controlled XFD-wide row. row_iterator = worksheet.iter_rows( values_only=True, max_col=MAX_BULK_COLUMNS + 1, ) raw_headers = next(row_iterator, None) if raw_headers is None: return [] header_values = _trim_trailing_empty_cells(raw_headers) _validate_row_width(header_values) headers = [str(cell).strip() if cell is not None else "" for cell in header_values] def data_rows() -> Iterable[dict[str, Any]]: for raw_row in row_iterator: row = _trim_trailing_empty_cells(raw_row) _validate_row_width(row) yield { headers[index]: row[index] for index in range(min(len(headers), len(row))) if headers[index] } return _collect_bounded_rows(data_rows()) except AppException: raise except Exception as exc: # openpyxl exposes several parser exception types logger.info("asset bulk XLSX parse rejected: %s", type(exc).__name__) raise _invalid_file("Malformed XLSX file") from exc finally: if workbook is not None: workbook.close() async def read_bulk_rows(upload: UploadFile) -> list[dict[str, Any]]: """Read a CSV/XLSX upload with bounded compressed and expanded work.""" try: extension = validate_bulk_upload_contract(upload) payload = await upload.read(MAX_UPLOAD_BYTES + 1) finally: await upload.close() if len(payload) > MAX_UPLOAD_BYTES: raise _limit_exceeded(f"File exceeds maximum size of {MAX_UPLOAD_BYTES // (1024 * 1024)}MB") if extension == ".csv": return _parse_csv(payload) return _parse_xlsx(payload)