From 55745531cc265f65b07328065f057bd6b6eb6c54 Mon Sep 17 00:00:00 2001 From: HC-ONLINE Date: Fri, 16 Oct 2026 14:03:19 -0500 Subject: [PATCH 1/2] refactor(export): centralize supported export formats --- src/ciberwebscan/api/models/requests.py | 12 +++++++----- src/ciberwebscan/api/models/responses.py | 2 +- src/ciberwebscan/export/__init__.py | 5 +++++ src/ciberwebscan/export/base.py | 12 +++++++++--- 4 files changed, 22 insertions(+), 9 deletions(-) diff --git a/src/ciberwebscan/api/models/requests.py b/src/ciberwebscan/api/models/requests.py index 8a42e76..9a34aca 100644 --- a/src/ciberwebscan/api/models/requests.py +++ b/src/ciberwebscan/api/models/requests.py @@ -11,6 +11,8 @@ from pydantic import BaseModel, Field, HttpUrl, field_validator +from ciberwebscan.export import ExportFormat + # ============================================================================= # Scrape Requests # ============================================================================= @@ -77,7 +79,7 @@ class ScrapeRequest(BaseModel): default=None, description="Optional output file path for exported results", ) - export_format: Literal["json", "jsonl", "csv", "html"] = Field( + export_format: ExportFormat = Field( default="json", description="Export format when export path is provided", ) @@ -112,7 +114,7 @@ class ScrapeBatchRequest(BaseModel): default=None, description="Optional output file path for exported batch results", ) - export_format: Literal["json", "jsonl", "csv", "html"] = Field( + export_format: ExportFormat = Field( default="jsonl", description="Export format when export path is provided", ) @@ -176,7 +178,7 @@ class AnalyzeRequest(BaseModel): default=None, description="Optional output file path for exported results", ) - export_format: Literal["json", "jsonl", "csv", "html"] = Field( + export_format: ExportFormat = Field( default="json", description="Export format when export path is provided", ) @@ -265,7 +267,7 @@ class AttackRequest(BaseModel): default=None, description="Optional output file path for exported results", ) - export_format: Literal["json", "jsonl", "csv", "html"] = Field( + export_format: ExportFormat = Field( default="json", description="Export format when export path is provided", ) @@ -319,7 +321,7 @@ class ExportRequest(BaseModel): """Request for exporting results.""" job_id: str = Field(..., description="Job ID from a previous operation") - format: Literal["jsonl", "json", "csv", "html"] = "jsonl" + format: ExportFormat = "jsonl" include_raw_html: bool = False include_screenshots: bool = False diff --git a/src/ciberwebscan/api/models/responses.py b/src/ciberwebscan/api/models/responses.py index 54171fe..732bb35 100644 --- a/src/ciberwebscan/api/models/responses.py +++ b/src/ciberwebscan/api/models/responses.py @@ -267,7 +267,7 @@ class DownloadInfo(BaseModel): expires_at: datetime attempts_remaining: int file_format: str = Field( - default="json", description="Export format (json/jsonl/csv)" + default="json", description="Export format (json/jsonl/csv/html)" ) diff --git a/src/ciberwebscan/export/__init__.py b/src/ciberwebscan/export/__init__.py index 63a2376..3a2e958 100644 --- a/src/ciberwebscan/export/__init__.py +++ b/src/ciberwebscan/export/__init__.py @@ -31,8 +31,10 @@ # Base classes and utilities from ciberwebscan.export.base import ( + EXPORT_FORMATS, BaseExporter, ExportError, + ExportFormat, ExportValidationError, ExportWriteError, StreamingExporter, @@ -102,6 +104,9 @@ "ExportError", "ExportWriteError", "ExportValidationError", + # Export formats + "ExportFormat", + "EXPORT_FORMATS", # Factory functions "export_to_file", "get_exporter", diff --git a/src/ciberwebscan/export/base.py b/src/ciberwebscan/export/base.py index 4fa48be..ad5c998 100644 --- a/src/ciberwebscan/export/base.py +++ b/src/ciberwebscan/export/base.py @@ -13,13 +13,19 @@ from contextlib import contextmanager from datetime import datetime from pathlib import Path -from typing import TYPE_CHECKING, Any, TextIO +from typing import TYPE_CHECKING, Any, Literal, TextIO, get_args if TYPE_CHECKING: from ciberwebscan.export.models import AnalysisReport logger = logging.getLogger(__name__) +# Canonical export formats. Single source of truth: the Literal drives +# validation across API models and download metadata, and EXPORT_FORMATS +# is derived from it for runtime membership checks. +ExportFormat = Literal["json", "jsonl", "csv", "html"] +EXPORT_FORMATS: frozenset[str] = frozenset(get_args(ExportFormat)) + class ExportError(Exception): """Base exception for export errors.""" @@ -320,9 +326,9 @@ def get_exporter( } format_lower = format.lower() - if format_lower not in exporters: + if format_lower not in EXPORT_FORMATS: raise ExportError( - f"Unknown export format: {format}. Use: {list(exporters.keys())}" + f"Unknown export format: {format}. Use: {sorted(EXPORT_FORMATS)}" ) return exporters[format_lower](output_path=output_path, **kwargs) From 4b35c768cb869d781ec07e7f8758b59a71d19161 Mon Sep 17 00:00:00 2001 From: HC-ONLINE Date: Fri, 16 Oct 2026 14:04:44 -0500 Subject: [PATCH 2/2] fix(api): generate download filenames from the exported format --- docs/API.md | 2 +- .../api/helpers/download_helper.py | 29 ++++- src/ciberwebscan/api/routes/download.py | 31 ++++- src/ciberwebscan/services/download_service.py | 26 +++- .../api/test_api_export_download.py | 36 ++++++ tests/unit/api/helpers/__init__.py | 0 .../unit/api/helpers/test_download_helper.py | 120 ++++++++++++++++++ tests/unit/api/routes/test_download.py | 37 +++++- tests/unit/services/test_download_service.py | 27 ++++ 9 files changed, 302 insertions(+), 6 deletions(-) create mode 100644 tests/unit/api/helpers/__init__.py create mode 100644 tests/unit/api/helpers/test_download_helper.py diff --git a/docs/API.md b/docs/API.md index 9b26c02..2881ced 100644 --- a/docs/API.md +++ b/docs/API.md @@ -611,7 +611,7 @@ Download previously exported results using a time-limited token. X-API-Key: your-api-key-here ``` -**Response:** File stream (application/octet-stream) +**Response:** File stream (application/octet-stream). The `Content-Disposition` header carries a filename matching the export format: `attachment; filename="export.json"` (also `.jsonl`, `.csv` or `.html`). **Status Codes:** diff --git a/src/ciberwebscan/api/helpers/download_helper.py b/src/ciberwebscan/api/helpers/download_helper.py index 73a2aef..cf17ad4 100644 --- a/src/ciberwebscan/api/helpers/download_helper.py +++ b/src/ciberwebscan/api/helpers/download_helper.py @@ -7,17 +7,41 @@ from __future__ import annotations +from pathlib import Path from typing import TypeVar from urllib.parse import urlparse from fastapi import Request +from ciberwebscan.export import EXPORT_FORMATS from ciberwebscan.services.base import ServiceResult from ciberwebscan.services.download_service import DownloadService T = TypeVar("T") +def _detect_export_format(export_format: str | None, export_path: Path) -> str: + """Resolve the export format for download metadata. + + Preference order: the format reported by the service, then the export + file extension, then ``json``. Unknown values are discarded so they + can never surface in download metadata. + + Args: + export_format: Format reported by the service result, if any. + export_path: Path of the exported file. + + Returns: + One of ``EXPORT_FORMATS``. + """ + if export_format is not None and export_format in EXPORT_FORMATS: + return export_format + suffix = export_path.suffix.removeprefix(".").lower() + if suffix in EXPORT_FORMATS: + return suffix + return "json" + + def enrich_response_with_token( result: ServiceResult[T], user_id: str, @@ -46,11 +70,12 @@ def enrich_response_with_token( if result.export_path is None: return result.data, None, None - # Generate token from the exported file + # Generate token from the exported file with the real export format + # (previously hardcoded to "json", which mislabeled every download). token_result = download_service.generate_download_token( file_path=result.export_path, user_id=user_id, - file_format="json", # Default format, could be detected from path extension + file_format=_detect_export_format(result.export_format, result.export_path), ) # If token generation failed, log but don't error diff --git a/src/ciberwebscan/api/routes/download.py b/src/ciberwebscan/api/routes/download.py index 31206ad..621fb71 100644 --- a/src/ciberwebscan/api/routes/download.py +++ b/src/ciberwebscan/api/routes/download.py @@ -16,6 +16,7 @@ from ciberwebscan.api.auth import AuthenticatedUser, get_current_user from ciberwebscan.config.loader import get_config +from ciberwebscan.export import EXPORT_FORMATS from ciberwebscan.services.download_service import DownloadService from ciberwebscan.utils.logging import mask_identifier @@ -23,6 +24,24 @@ router = APIRouter(prefix="/download", tags=["download"]) + +def _content_disposition(file_format: str | None) -> str: + """Build the Content-Disposition header for an export format. + + Only formats in ``EXPORT_FORMATS`` reach the header, so a + tampered registry value can never inject header content; anything + else falls back to ``json``. + + Args: + file_format: Format stored in the token metadata, if any. + + Returns: + The Content-Disposition value with a format-specific filename. + """ + extension = file_format if file_format in EXPORT_FORMATS else "json" + return f'attachment; filename="export.{extension}"' + + # Service instance _download_service = DownloadService() @@ -228,9 +247,19 @@ async def download_file( detail="Failed to retrieve file stream", ) + # Token metadata: the stored export format drives the download filename + info_result = download_service.get_info(token=token) + if not info_result.success or info_result.data is None: + reservation.refund_once() + logger.error(f"Failed to get download metadata: {info_result.error}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to retrieve file metadata", + ) + # Prepare headers headers = { - "Content-Disposition": "attachment; filename=export.json", + "Content-Disposition": _content_disposition(info_result.data.file_format), "X-Content-Type-Options": "nosniff", "Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", diff --git a/src/ciberwebscan/services/download_service.py b/src/ciberwebscan/services/download_service.py index 2c9131c..d774792 100644 --- a/src/ciberwebscan/services/download_service.py +++ b/src/ciberwebscan/services/download_service.py @@ -141,7 +141,7 @@ def generate_download_token( Args: file_path: Path to file to download. Must be within allowed directories. user_id: ID of user requesting download - file_format: Format of the exported file (json/jsonl/csv) + file_format: Format of the exported file (json/jsonl/csv/html) Returns: ServiceResult with DownloadTokenResponse containing token and expiration @@ -329,6 +329,30 @@ def restore_attempt(self, token: str) -> ServiceResult[bool]: self.logger.error(f"Error restoring attempt for token {token}: {e}") return ServiceResult(success=False, error=str(e)) + def get_info(self, token: str) -> ServiceResult[DownloadInfo]: + """ + Retrieve metadata (including file_format) for a download token. + + Read-only: expiration, ownership and attempt checks are the + caller's responsibility (``validate_download_request``), which + runs before any use of this metadata. + + Args: + token: Download token + + Returns: + ServiceResult with the token's DownloadInfo, or a failure + result when the token no longer exists. + """ + try: + info = run_async(_registry.get_info(token)) + if info is None: + return ServiceResult(success=False, error="Token not found") + return ServiceResult(success=True, data=info) + except Exception as e: + self.logger.error(f"Error getting token info: {e}") + return ServiceResult(success=False, error=str(e)) + def get_file_stream( self, token: str, diff --git a/tests/integration/api/test_api_export_download.py b/tests/integration/api/test_api_export_download.py index e1fa20a..7c5a9eb 100644 --- a/tests/integration/api/test_api_export_download.py +++ b/tests/integration/api/test_api_export_download.py @@ -265,6 +265,42 @@ def test_downloaded_html_has_valid_content( assert "

Report

" in text +class TestDownloadContentDisposition: + """The Content-Disposition filename must reflect the export format.""" + + @pytest.mark.parametrize( + ("file_format", "file_name"), + [ + ("json", "export.json"), + ("jsonl", "export.jsonl"), + ("csv", "export.csv"), + ("html", "export.html"), + ], + ) + def test_header_filename_matches_export_format( + self, in_process_client: tuple, tmp_path: Path, file_format: str, file_name: str + ): + client, _ = in_process_client + + test_file = tmp_path / f"data.{file_format}" + test_file.write_text("payload") + + service = DownloadService() + result = service.generate_download_token( + file_path=test_file, + user_id="test_user", + file_format=file_format, + ) + token = result.data.token + + response = client.get(f"/api/download/{token}") + assert response.status_code == 200 + assert ( + response.headers["content-disposition"] + == f'attachment; filename="{file_name}"' + ) + + class TestTokenLifecycle: """End-to-end token lifecycle test.""" diff --git a/tests/unit/api/helpers/__init__.py b/tests/unit/api/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/api/helpers/test_download_helper.py b/tests/unit/api/helpers/test_download_helper.py new file mode 100644 index 0000000..2e200a4 --- /dev/null +++ b/tests/unit/api/helpers/test_download_helper.py @@ -0,0 +1,120 @@ +""" +Unit tests for download token enrichment in routes. + +Covers export-format detection when generating download metadata: the +format reported by the service wins, the export file extension is the +fallback, and unknown values resolve to "json" instead of leaking into +download metadata. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from ciberwebscan.api.helpers.download_helper import ( + _detect_export_format, + enrich_response_with_token, +) +from ciberwebscan.api.models.responses import DownloadTokenResponse +from ciberwebscan.services.base import ServiceResult + +pytestmark = pytest.mark.unit + + +# ============================================================================= +# Test: Format detection +# ============================================================================= + + +class TestDetectExportFormat: + """Preference order: service format → file extension → json.""" + + @pytest.mark.parametrize( + ("export_format", "path", "expected"), + [ + ("csv", "out.json", "csv"), + ("jsonl", "out.csv", "jsonl"), + ("html", "out.txt", "html"), + (None, "out.csv", "csv"), + (None, "out.jsonl", "jsonl"), + (None, "out.HTML", "html"), + ("xml", "out.csv", "csv"), + ('csv"\r\nX: 1', "out.csv", "csv"), + ("xml", "out.txt", "json"), + (None, "out", "json"), + (None, "out.xml", "json"), + ("", "out", "json"), + ], + ) + def test_detection(self, export_format: str | None, path: str, expected: str): + assert _detect_export_format(export_format, Path(path)) == expected + + +# ============================================================================= +# Test: Token enrichment +# ============================================================================= + + +class TestEnrichResponseWithToken: + """The generated token stores the detected export format.""" + + def _make_request(self) -> Mock: + request = Mock() + request.url_for.return_value = "http://testserver/api/download/tok-1" + return request + + def _make_service(self) -> Mock: + service = Mock() + service.generate_download_token.return_value = ServiceResult.ok( + DownloadTokenResponse(token="tok-1", expires_at=datetime.now(timezone.utc)) + ) + return service + + def test_token_stores_service_reported_format(self): + result: ServiceResult[dict] = ServiceResult.ok( + data={"ok": True}, + export_path=Path("exports/report.csv"), + export_format="csv", + ) + + service = self._make_service() + _, token, url = enrich_response_with_token( + result, "user-1", service, self._make_request() + ) + + assert token == "tok-1" + assert url == "/api/download/tok-1" + service.generate_download_token.assert_called_once_with( + file_path=Path("exports/report.csv"), + user_id="user-1", + file_format="csv", + ) + + def test_token_falls_back_to_extension_without_service_format(self): + result: ServiceResult[dict] = ServiceResult.ok( + data={"ok": True}, export_path=Path("exports/report.jsonl") + ) + + service = self._make_service() + enrich_response_with_token(result, "user-1", service, self._make_request()) + + assert ( + service.generate_download_token.call_args.kwargs["file_format"] == "jsonl" + ) + + def test_no_export_path_returns_no_token(self): + result: ServiceResult[dict] = ServiceResult.ok(data={"ok": True}) + + service = self._make_service() + data, token, url = enrich_response_with_token( + result, "user-1", service, self._make_request() + ) + + assert data == {"ok": True} + assert token is None + assert url is None + service.generate_download_token.assert_not_called() diff --git a/tests/unit/api/routes/test_download.py b/tests/unit/api/routes/test_download.py index 00e6e62..1f92888 100644 --- a/tests/unit/api/routes/test_download.py +++ b/tests/unit/api/routes/test_download.py @@ -16,7 +16,11 @@ import pytest -from ciberwebscan.api.routes.download import _AttemptReservation, _stream_and_cleanup +from ciberwebscan.api.routes.download import ( + _AttemptReservation, + _content_disposition, + _stream_and_cleanup, +) from ciberwebscan.config.loader import get_config from ciberwebscan.services.download_service import DownloadService, _registry @@ -128,6 +132,37 @@ def test_refund_once_without_consumed_reservation_is_refused( assert _attempts(token) == max_retries +# ============================================================================= +# Test: Content-Disposition header +# ============================================================================= + + +class TestContentDisposition: + """The download filename always reflects a supported export format.""" + + @pytest.mark.parametrize( + ("file_format", "expected"), + [ + ("json", 'attachment; filename="export.json"'), + ("jsonl", 'attachment; filename="export.jsonl"'), + ("csv", 'attachment; filename="export.csv"'), + ("html", 'attachment; filename="export.html"'), + ], + ) + def test_supported_formats_produce_matching_filename( + self, file_format: str, expected: str + ): + assert _content_disposition(file_format) == expected + + @pytest.mark.parametrize( + "file_format", + [None, "", "xml", "js", "json\r\nX-Injected: 1", 'csv"; evil=y'], + ) + def test_unsupported_values_fall_back_to_json(self, file_format: str | None): + """Registry values outside the allowlist never reach the header.""" + assert _content_disposition(file_format) == 'attachment; filename="export.json"' + + # ============================================================================= # Test: Streaming finalization # ============================================================================= diff --git a/tests/unit/services/test_download_service.py b/tests/unit/services/test_download_service.py index d29566a..d07daf5 100644 --- a/tests/unit/services/test_download_service.py +++ b/tests/unit/services/test_download_service.py @@ -128,6 +128,33 @@ def test_generate_token_file_too_large(self, service: DownloadService): # ============================================================================= +class TestGetInfo: + """get_info exposes stored metadata (notably file_format) to callers.""" + + def test_get_info_returns_stored_file_format( + self, service: DownloadService, test_file: Path + ): + gen_result = service.generate_download_token( + file_path=test_file, + user_id="test_user", + file_format="csv", + ) + assert gen_result.success is True + + info_result = service.get_info(gen_result.data.token) + + assert info_result.success is True + assert info_result.data is not None + assert info_result.data.file_format == "csv" + assert info_result.data.user_id == "test_user" + + def test_get_info_unknown_token_fails(self, service: DownloadService): + info_result = service.get_info("missing-token") + + assert info_result.success is False + assert info_result.error == "Token not found" + + class TestValidateDownloadRequest: """Tests for request validation."""