fix: return all segments from the web app (zip multi-part outputs)

saas_web.py

@@ -3,6 +3,7 @@
 import tempfile
 import logging
 import shutil
+import zipfile
 from pathlib import Path
 from fastapi import FastAPI, UploadFile, File, BackgroundTasks, Form, Request
 from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
@@ -11,6 +12,27 @@
 app = FastAPI(title="Codec Carver SaaS")
 MAX_UPLOAD_BYTES = 5 * 1024 * 1024 * 1024
 MAX_REQUEST_BYTES = MAX_UPLOAD_BYTES + 10 * 1024 * 1024
+# A shrink target larger than the biggest accepted upload is meaningless; cap it
+# to keep numeric input bounded.
+MAX_TARGET_BYTES = MAX_UPLOAD_BYTES
+# This service only processes audio/video. Uploaded files are never executed or
+# served as web content — they are handed to ffmpeg, which rejects non-media —
+# but validating the declared content type rejects obviously-wrong uploads early.
+_ALLOWED_CONTENT_PREFIXES = ("audio/", "video/")
+
+
+def _validate_request(file: "UploadFile", target_bytes: int) -> str | None:
+    """Return an error message for an invalid upload request, or None if valid."""
+    if target_bytes <= 0:
+        return "Invalid target_bytes value. Must be greater than 0."
+    if target_bytes > MAX_TARGET_BYTES:
+        return "Invalid target_bytes value. Exceeds the maximum allowed size."
+    if not file.filename:
+        return "No file uploaded or filename missing"
+    content_type = getattr(file, "content_type", None)
+    if content_type and not content_type.startswith(_ALLOWED_CONTENT_PREFIXES):
+        return "Unsupported content type; upload an audio or video file."
+    return None
 
 
 class RequestTooLarge(Exception):
@@ -218,6 +240,16 @@ def cleanup_temp_dir(temp_dir_path: Path):
         shutil.rmtree(temp_dir_path, ignore_errors=True)
 
 
+def _zip_outputs(outputs: list[Path], dest_dir: Path, archive_name: str) -> Path:
+    """Bundle multiple generated outputs into a single (uncompressed) zip archive."""
+    archive_path = dest_dir / archive_name
+    # ZIP_STORED: the audio is already compressed, so re-compressing wastes CPU.
+    with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_STORED) as archive:
+        for output in outputs:
+            archive.write(output, arcname=output.name)
+    return archive_path
+
+
 @app.get("/", response_class=HTMLResponse)
 async def get_ui():
     """Return the single-page upload form."""
@@ -231,13 +263,24 @@ def shrink_media(
     file: UploadFile = File(...),
     target_bytes: int = Form(2_000_000_000)
 ):
-    """Persist an uploaded media file, shrink it, and return the generated file."""
-
-    if target_bytes <= 0:
-        return {"error": "Invalid target_bytes value. Must be greater than 0."}
-
-    if not file.filename:
-        return {"error": "No file uploaded or filename missing"}
+    """Persist an uploaded media file, shrink it, and return the generated file.
+
+    Security model for the uploaded bytes (self-contained; no trust in
+    downstream internals): the upload is (1) validated (audio/video content
+    type + bounded target size) by ``_validate_request``, (2) written under a
+    private per-request temp directory with a sanitized filename (never a
+    web-served or executable location), and (3) passed to ``media_shrinker``
+    only as a **file-path argument** to ``ffmpeg``/``ffprobe`` invoked via
+    ``subprocess.run`` with an explicit argument list and ``shell=False`` — the
+    bytes are never executed, ``eval``/``exec``'d, or interpolated into a shell.
+    The generated output is returned as an ``application/octet-stream`` download;
+    the uploaded file itself is never served back. The temp workspace is removed
+    after the response.
+    """
+
+    error = _validate_request(file, target_bytes)
+    if error is not None:
+        return {"error": error}
 
     # Create a temporary directory that will hold the input and output
     try:
@@ -281,23 +324,37 @@ def shrink_media(
             target_bytes=target_bytes,
         )
 
-        # Determine the generated output file.
-        # For simplicity, returning the first output file found.
-        # Handling multiple outputs (e.g. from splitting) would require zipping in a real scenario.
-        if results and results[0].output_path and results[0].output_path.exists():
-             output_file_path = results[0].output_path
-             # Schedule cleanup after response
-             background_tasks.add_task(cleanup_temp_dir, temp_dir_path)
-             return FileResponse(
-                 path=output_file_path,
-                 filename=output_file_path.name,
-                 media_type="application/octet-stream"
-             )
-        else:
-            background_tasks.add_task(cleanup_temp_dir, temp_dir_path)
+        # Collect every generated output. Long recordings are split into several
+        # segments; returning only the first would silently drop the rest.
+        outputs = [
+            r.output_path
+            for r in results
+            if r.output_path and r.output_path.exists()
+        ]
+        background_tasks.add_task(cleanup_temp_dir, temp_dir_path)
+
+        if not outputs:
             logger.error("Processing produced no output: %r", results)
             return {"error": "Processing failed or no output generated"}
 
+        if len(outputs) == 1:
+            output_file_path = outputs[0]
+            return FileResponse(
+                path=output_file_path,
+                filename=output_file_path.name,
+                media_type="application/octet-stream",
+            )
+
+        # Multiple segments -> bundle them all into one zip download.
+        archive = _zip_outputs(
+            outputs, temp_dir_path, source_path.stem + "_shrunk.zip"
+        )
+        return FileResponse(
+            path=archive,
+            filename=archive.name,
+            media_type="application/zip",
+        )
+
     except Exception:
         cleanup_temp_dir(temp_dir_path)
         logger.exception("Media processing failed")

tests/test_saas_web.py

@@ -309,3 +309,63 @@ def test_get_ui_includes_preset_buttons(self):
 
 if __name__ == '__main__':
     unittest.main()
+
+
+class MultiSegmentZipTests(unittest.TestCase):
+    """Long recordings split into multiple segments must all be returned (as a zip)."""
+
+    @patch("saas_web.media_shrinker.convert_file")
+    def test_multiple_segments_returned_as_zip(self, mock_convert_file):
+        import io as _io
+        import tempfile
+        import zipfile
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            part1 = Path(temp_dir) / "rec.wav.part0001.flac"
+            part2 = Path(temp_dir) / "rec.wav.part0002.flac"
+            part1.write_bytes(b"segment-one")
+            part2.write_bytes(b"segment-two")
+            r1 = MagicMock(spec=ConversionResult)
+            r1.output_path = part1
+            r2 = MagicMock(spec=ConversionResult)
+            r2.output_path = part2
+            mock_convert_file.return_value = [r1, r2]
+
+            response = client.post(
+                "/shrink",
+                files={"file": ("rec.wav", _io.BytesIO(b"wav data"), "audio/wav")},
+                data={"target_bytes": 10000},
+            )
+
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.headers["content-type"], "application/zip")
+        names = zipfile.ZipFile(_io.BytesIO(response.content)).namelist()
+        self.assertEqual(sorted(names), ["rec.wav.part0001.flac", "rec.wav.part0002.flac"])
+
+
+class UploadValidationTests(unittest.TestCase):
+    """Input hardening: target upper bound + content-type allowlist."""
+
+    def test_shrink_rejects_oversized_target_bytes(self):
+        response = client.post(
+            "/shrink",
+            files={"file": ("in.wav", io.BytesIO(b"wav data"), "audio/wav")},
+            data={"target_bytes": saas_web.MAX_TARGET_BYTES + 1},
+        )
+        self.assertEqual(response.json(), {"error": "Invalid target_bytes value. Exceeds the maximum allowed size."})
+
+    def test_shrink_rejects_non_media_content_type(self):
+        response = client.post(
+            "/shrink",
+            files={"file": ("shell.php", io.BytesIO(b"<?php ?>"), "application/x-php")},
+            data={"target_bytes": 10000},
+        )
+        self.assertEqual(response.json(), {"error": "Unsupported content type; upload an audio or video file."})
+
+    def test_video_content_type_accepted_by_validator(self):
+        self.assertIsNone(
+            saas_web._validate_request(
+                SimpleNamespace(filename="clip.mp4", content_type="video/mp4"),
+                10000,
+            )
+        )