Added file download tool

server.py

@@ -17,14 +17,16 @@ import logging
 import asyncio
 import subprocess
 import tempfile
-from typing import Optional
+import secrets
+import time
+from typing import Dict, Optional
 from urllib.parse import urlparse
 
 from fastmcp import FastMCP
 from fastmcp.server.middleware import Middleware, MiddlewareContext
 from starlette.middleware import Middleware as StarletteMiddleware
 from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import JSONResponse
+from starlette.responses import FileResponse, JSONResponse
 
 # ---------------------------------------------------------------------------
 # Logging (JSON Format)
@@ -70,11 +72,18 @@ WORKSPACE_MOUNT = os.path.abspath(os.environ.get("WORKSPACE_MOUNT", "/workspace"
 
 logger.info(f"Workspace Root configured to: {WORKSPACE_MOUNT}")
 
+PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "http://localhost:8000")
+DOWNLOAD_DEFAULT_TTL = 10
+DOWNLOAD_MIN_TTL = 1
+DOWNLOAD_MAX_TTL = 60
+
 # ---------------------------------------------------------------------------
 # FastMCP app
 # ---------------------------------------------------------------------------
 mcp = FastMCP("turnstone-mcp")
 
+TOKEN_STORE: Dict[str, dict] = {}
+
 # ---------------------------------------------------------------------------
 # Playwright lifecycle (stateful browser session)
 # ---------------------------------------------------------------------------
@@ -121,6 +130,26 @@ def _resolve_safe_path(relative_path: str) -> str:
         raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
     return target
 
+
+def _validate_download_path(raw_path: str) -> str:
+    canonical_path = os.path.realpath(raw_path)
+    if not canonical_path.startswith(WORKSPACE_MOUNT + "/"):
+        raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
+    if not os.path.exists(canonical_path):
+        raise FileNotFoundError(f"File not found: {raw_path}")
+    if os.path.isdir(canonical_path):
+        raise IsADirectoryError(
+            "Target path is a directory. Please compress it to an archive first."
+        )
+    return canonical_path
+
+
+def _purge_expired_tokens():
+    now = time.time()
+    expired_keys = [k for k, v in TOKEN_STORE.items() if now > v["expires_at"]]
+    for k in expired_keys:
+        TOKEN_STORE.pop(k, None)
+
 # ---------------------------------------------------------------------------
 # Health endpoint
 # ---------------------------------------------------------------------------
@@ -128,6 +157,38 @@ def _resolve_safe_path(relative_path: str) -> str:
 async def health_check(request):
     return JSONResponse({"status": "healthy", "service": "turnstone-mcp"})
 
+
+@mcp.custom_route("/download", methods=["GET"])
+async def download_file(request):
+    _purge_expired_tokens()
+
+    token = request.query_params.get("token")
+    if not token:
+        return JSONResponse({"detail": "Missing token parameter."}, status_code=400)
+
+    record = TOKEN_STORE.get(token)
+    if not record:
+        return JSONResponse(
+            {"detail": "Invalid or previously used download link."}, status_code=403
+        )
+
+    del TOKEN_STORE[token]
+
+    if time.time() > record["expires_at"]:
+        return JSONResponse({"detail": "Download link has expired."}, status_code=410)
+
+    file_path = record["file_path"]
+    if not os.path.exists(file_path):
+        return JSONResponse(
+            {"detail": "File no longer exists on host."}, status_code=404
+        )
+
+    return FileResponse(
+        path=file_path,
+        filename=os.path.basename(file_path),
+        media_type="application/octet-stream",
+    )
+
 # ---------------------------------------------------------------------------
 # Semantic-logging middleware
 # ---------------------------------------------------------------------------
@@ -351,6 +412,46 @@ def extract_to_markdown(relative_path: str) -> str:
     except Exception as e:
         return f"Extraction error: {e}"
 
+# ===================================================================
+# Tool 5 — generate_download_link
+# ===================================================================
+
+@mcp.tool()
+def generate_download_link(path: str, ttl_minutes: int = DOWNLOAD_DEFAULT_TTL) -> str:
+    """Generates a secure, single-use, temporary HTTPS download link for a file
+    stored within the /workspace directory.
+
+    Parameters
+    ----------
+    path : str
+        The absolute path to the file inside the container (must reside within /workspace).
+    ttl_minutes : int
+        Validity period of the link in minutes. Default is 10. Clamped to 1–60.
+
+    Returns
+    -------
+    str
+        A clickable download URL, or an error message if validation fails.
+    """
+    try:
+        abs_path = _validate_download_path(path)
+    except (PermissionError, FileNotFoundError, IsADirectoryError) as e:
+        return f"Error: {e}"
+
+    _purge_expired_tokens()
+
+    effective_ttl = max(DOWNLOAD_MIN_TTL, min(ttl_minutes, DOWNLOAD_MAX_TTL))
+    token = secrets.token_urlsafe(32)
+    expires_at = time.time() + (effective_ttl * 60)
+
+    TOKEN_STORE[token] = {
+        "file_path": abs_path,
+        "expires_at": expires_at,
+    }
+
+    return f"{PUBLIC_BASE_URL}/download?token={token}"
+
+
 # ===================================================================
 # Execution
 # ===================================================================