diff --git a/config.py b/config.py index cf50904..d40e600 100644 --- a/config.py +++ b/config.py @@ -3,11 +3,37 @@ import os from pathlib import Path -from typing import Set +from typing import Set, Tuple OUTPUTS_DIR = Path("./outputs").resolve() TEMP_UPLOADS_DIR = Path("./temp_uploads").resolve() + +def _positive_int_env(name: str, default: int) -> int: + """Read a positive integer deployment limit without silently disabling it.""" + value = int(os.getenv(name, str(default))) + if value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _csv_env(name: str) -> Tuple[str, ...]: + return tuple(item.strip() for item in os.getenv(name, "").split(",") if item.strip()) + + +# Public upload guardrails. Railway can override these per environment without a code change. +MAX_UPLOAD_FILES = _positive_int_env("MAX_UPLOAD_FILES", 50) +MAX_UPLOAD_FILE_BYTES = _positive_int_env("MAX_UPLOAD_FILE_BYTES", 100 * 1024 * 1024) +MAX_UPLOAD_TOTAL_BYTES = _positive_int_env("MAX_UPLOAD_TOTAL_BYTES", 500 * 1024 * 1024) +UPLOAD_CHUNK_BYTES = 1024 * 1024 + +# MCP is disabled unless an operator deliberately enables it in a controlled environment. +MCP_ENABLED = os.getenv("MCP_ENABLED", "false").strip().lower() in {"1", "true", "yes", "on"} +MCP_SOURCE_ROOTS = tuple(Path(item).expanduser().resolve() for item in _csv_env("MCP_SOURCE_ROOTS")) + +# Set this to the exact production domain(s) to reject forged Host headers. +TRUSTED_HOSTS = _csv_env("TRUSTED_HOSTS") + TOOL_DISPLAY_NAME: str = "PhotoPackager" ORIGINAL_AUTHOR: str = "Steven Seagondollar, DropShock Digital LLC" diff --git a/main.py b/main.py index cd559b7..e63b3dc 100644 --- a/main.py +++ b/main.py @@ -4,12 +4,22 @@ from typing import List from fastapi import FastAPI, UploadFile, File, Form, HTTPException +from fastapi.middleware.trustedhost import TrustedHostMiddleware from fastapi.responses import FileResponse, JSONResponse from fastapi.staticfiles import StaticFiles # Import shared components from worker import celery_app -from config import OUTPUTS_DIR, TEMP_UPLOADS_DIR +from config import ( + MCP_ENABLED, + MAX_UPLOAD_FILES, + MAX_UPLOAD_FILE_BYTES, + MAX_UPLOAD_TOTAL_BYTES, + OUTPUTS_DIR, + TEMP_UPLOADS_DIR, + TRUSTED_HOSTS, + UPLOAD_CHUNK_BYTES, +) from schemas import JobSettings, JobResponse # Import MCP Server components @@ -18,9 +28,12 @@ app = FastAPI() +if TRUSTED_HOSTS: + app.add_middleware(TrustedHostMiddleware, allowed_hosts=list(TRUSTED_HOSTS)) + STATIC_DIR = Path(__file__).parent / "frontend" / "dist" -mcp_server = FastMCP(tools=get_tools()) +mcp_server = FastMCP(tools=get_tools()) if MCP_ENABLED else None # Ensure base directories exist TEMP_UPLOADS_DIR.mkdir(exist_ok=True) @@ -43,6 +56,41 @@ def _dump_job_settings(job_settings: JobSettings) -> dict: return dumper() return job_settings.dict() + +async def _save_uploads(files: List[UploadFile], job_dir: Path) -> None: + """Write uploads with bounded file and aggregate sizes to protect disk and workers.""" + if len(files) > MAX_UPLOAD_FILES: + raise HTTPException(status_code=413, detail=f"Too many files; maximum is {MAX_UPLOAD_FILES}.") + + total_bytes = 0 + for file in files: + if not file.filename: + await file.close() + continue + + clean_filename = Path(file.filename).name + if not clean_filename or clean_filename in {".", ".."}: + await file.close() + continue + + file_path = job_dir / clean_filename + file_bytes = 0 + try: + with open(file_path, "wb") as buffer: + while chunk := await file.read(UPLOAD_CHUNK_BYTES): + file_bytes += len(chunk) + total_bytes += len(chunk) + if file_bytes > MAX_UPLOAD_FILE_BYTES: + raise HTTPException(status_code=413, detail="An uploaded file is too large.") + if total_bytes > MAX_UPLOAD_TOTAL_BYTES: + raise HTTPException(status_code=413, detail="The upload is too large.") + buffer.write(chunk) + except Exception: + file_path.unlink(missing_ok=True) + raise + finally: + await file.close() + def get_job_status(job_id: str): """Helper to get the status of a Celery task.""" task_result = celery_app.AsyncResult(job_id) @@ -85,17 +133,11 @@ async def create_packaging_job( job_dir = TEMP_UPLOADS_DIR / job_id job_dir.mkdir(parents=True, exist_ok=True) - # Save uploaded files - for file in files: - if not file.filename: - continue - # Strict sanitization: remove path components, get purely the filename. - clean_filename = Path(file.filename).name - if not clean_filename or clean_filename == '.' or clean_filename == '..': - continue - file_path = job_dir / clean_filename - with open(file_path, "wb") as buffer: - shutil.copyfileobj(file.file, buffer) + try: + await _save_uploads(files, job_dir) + except Exception: + shutil.rmtree(job_dir, ignore_errors=True) + raise # Launch background task with Celery celery_app.send_task( @@ -142,7 +184,8 @@ async def download_zip_package(job_id: str, zip_filename: str): def _mount_runtime_apps() -> None: """Mount runtime sub-apps after API routes so they do not shadow /api.""" - app.mount("/mcp", mcp_server) + if mcp_server is not None: + app.mount("/mcp", mcp_server) if STATIC_DIR.is_dir(): app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="frontend") diff --git a/mcp_tools.py b/mcp_tools.py index 85ecd9e..a2bc791 100644 --- a/mcp_tools.py +++ b/mcp_tools.py @@ -7,7 +7,7 @@ # Import shared components from schemas import JobSettings, JobResponse from worker import celery_app -from config import TEMP_UPLOADS_DIR +from config import MCP_SOURCE_ROOTS, TEMP_UPLOADS_DIR def _dump_job_settings(job_settings: JobSettings) -> dict: @@ -31,13 +31,23 @@ async def package_photos(input: MCPPackagePhotosInput) -> JobResponse: job_id = str(uuid.uuid4()) job_dir = TEMP_UPLOADS_DIR / job_id + if not MCP_SOURCE_ROOTS: + return JobResponse( + job_id=job_id, + status="failed", + message="MCP source roots are not configured." + ) + # 1. Create a unique directory for the job and copy files try: job_dir.mkdir(parents=True, exist_ok=True) for src_file in input.source_files: - if not src_file.is_file(): + resolved_source = src_file.resolve() + if not any(resolved_source.is_relative_to(root) for root in MCP_SOURCE_ROOTS): + raise PermissionError(f"Source path is outside the configured MCP roots: {src_file}") + if not resolved_source.is_file(): raise FileNotFoundError(f"Source file not found: {src_file}") - shutil.copy(src_file, job_dir / src_file.name) + shutil.copy(resolved_source, job_dir / resolved_source.name) except Exception as e: # Cleanup if setup fails if job_dir.exists(): diff --git a/requirements.txt b/requirements.txt index 7ef969c..39706d6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ -fastapi==0.111.0 +fastapi==0.139.2 +starlette==1.3.1 celery==5.4.0 redis==5.0.4 uvicorn[standard]==0.51.0 diff --git a/tests/test_api.py b/tests/test_api.py index e58d382..9569ebe 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -50,6 +50,40 @@ def test_create_job_invalid_settings(): assert response.status_code == 400 + +def _valid_settings(): + return { + "output_format": "JPEG", "output_quality": 85, "resize_enabled": False, + "max_width": 1920, "max_height": 1080, "strip_exif": True, + "watermark_enabled": False, "watermark_text": "", "skip_export": False, + } + + +def test_create_job_rejects_too_many_files(monkeypatch): + import main + monkeypatch.setattr(main, "MAX_UPLOAD_FILES", 1) + response = client.post( + "/api/jobs", + data={"settings": json.dumps(_valid_settings())}, + files=[("files", ("one.jpg", b"a", "image/jpeg")), ("files", ("two.jpg", b"b", "image/jpeg"))], + ) + assert response.status_code == 413 + + +def test_create_job_rejects_oversized_file(monkeypatch): + import main + monkeypatch.setattr(main, "MAX_UPLOAD_FILE_BYTES", 2) + response = client.post( + "/api/jobs", + data={"settings": json.dumps(_valid_settings())}, + files=[("files", ("large.jpg", b"abc", "image/jpeg"))], + ) + assert response.status_code == 413 + + +def test_mcp_is_not_mounted_by_default(): + assert all(getattr(route, "path", None) != "/mcp" for route in app.routes) + def test_get_job_status(): job_id = str(uuid.uuid4()) response = client.get(f"/api/jobs/{job_id}/status")