diff --git a/saas_web.py b/saas_web.py index 2f177f8..7019f48 100644 --- a/saas_web.py +++ b/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") diff --git a/tests/test_saas_web.py b/tests/test_saas_web.py index 06852fa..f76135b 100644 --- a/tests/test_saas_web.py +++ b/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""), "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, + ) + )