Merge pull request #27 from CompassSecurity/fix/file-upload-handling

fix(backend):improved upload filter against polyglot files

Author: fxai

backend/app/services/file/file.py

@@ -110,6 +110,25 @@ def validate_file_content(file: UploadFile) -> FileType:
     )
 
 
+_CANONICAL_EXTENSION = {
+    FileType.PNG: ".png",
+    FileType.JPEG: ".jpg",
+    FileType.JPG: ".jpg",
+    FileType.TXT: ".txt",
+}
+
+
+def _normalize_filename(filename: str, file_type: FileType) -> str:
+    """
+    Force the filename extension to match the detected content type so that a
+    dangerous extension (e.g. .php) cannot be preserved by uploading a polyglot
+    file with valid image magic bytes.
+    """
+    stem = filename.rsplit(".", 1)[0] if "." in filename else filename
+    stem = stem or "unnamed"
+    return f"{stem}{_CANONICAL_EXTENSION[file_type]}"
+
+
 def upload_file_service(
     activity_id: uuid.UUID,
     file: UploadFile,
@@ -126,9 +145,9 @@ def upload_file_service(
     detected_file_type = validate_file_content(file)
     file_content = file.file.read()
 
-    final_filename = sanitize_filename(file.filename) or "unnamed"
-    if detected_file_type == FileType.TXT and not final_filename.endswith(".txt"):
-        final_filename += ".txt"
+    final_filename = _normalize_filename(
+        sanitize_filename(file.filename) or "unnamed", detected_file_type
+    )
 
     category = (
         FileCategory.RED

backend/tests/api/v1/assessments/test_activity_acl.py

@@ -301,6 +301,73 @@ def test_blue_user_cannot_upload_when_not_waiting(
     assert response.status_code == 403
 
 
+_PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
+
+
+def test_upload_polyglot_keeps_safe_extension(
+    client: TestClient,
+    session: Session,
+    test_assessment: Assessment,
+    test_activity: Activity,
+    auth_headers_regular: dict[str, str],
+    test_acl_red: Acl,
+):
+    """A PNG-magic file named shell.php is stored with a .png extension, not .php."""
+    response = client.post(
+        f"/api/v1/assessments/{test_assessment.id}/activity/{test_activity.id}/upload",
+        files={"file": ("shell.php", _PNG_BYTES, "image/png")},
+        headers=auth_headers_regular,
+    )
+    assert response.status_code == 200
+
+    file_id = uuid.UUID(response.json()["file_id"])
+    stored = session.get(File, file_id)
+    assert stored.filename == "shell.png"
+    assert ".php" not in stored.filename
+    assert stored.content_type == FileType.PNG
+
+
+def test_upload_text_normalized_to_txt(
+    client: TestClient,
+    session: Session,
+    test_assessment: Assessment,
+    test_activity: Activity,
+    auth_headers_regular: dict[str, str],
+    test_acl_red: Acl,
+):
+    """UTF-8 content with a non-txt extension is stored as .txt."""
+    response = client.post(
+        f"/api/v1/assessments/{test_assessment.id}/activity/{test_activity.id}/upload",
+        files={"file": ("notes.md", b"just some notes", "text/plain")},
+        headers=auth_headers_regular,
+    )
+    assert response.status_code == 200
+
+    stored = session.get(File, uuid.UUID(response.json()["file_id"]))
+    assert stored.filename == "notes.txt"
+    assert stored.content_type == FileType.TXT
+
+
+def test_upload_matching_extension_unchanged(
+    client: TestClient,
+    session: Session,
+    test_assessment: Assessment,
+    test_activity: Activity,
+    auth_headers_regular: dict[str, str],
+    test_acl_red: Acl,
+):
+    """A real png keeps its name (no double extension)."""
+    response = client.post(
+        f"/api/v1/assessments/{test_assessment.id}/activity/{test_activity.id}/upload",
+        files={"file": ("photo.png", _PNG_BYTES, "image/png")},
+        headers=auth_headers_regular,
+    )
+    assert response.status_code == 200
+
+    stored = session.get(File, uuid.UUID(response.json()["file_id"]))
+    assert stored.filename == "photo.png"
+
+
 def test_blue_user_fail_hidden(
     client: TestClient,
     session: Session,