Create main.py

main.py

@@ -0,0 +1,212 @@
+import os
+import shutil
+import time
+import glob
+import asyncio
+import mimetypes
+import pathlib
+import re
+from datetime import datetime, timedelta
+
+from fastapi import FastAPI, UploadFile, File, Form, Request, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+
+# CONFIG
+API_VERSION = "1.0"
+TMP_DIR = os.environ.get("TMP_DIR", "/tmp")  # container /tmp by default
+MAX_BYTES = 2 * 1024 * 1024 * 1024  # 2GB
+CLEANUP_INTERVAL_SECONDS = 600  # run cleanup every 10 minutes
+EXPIRE_SECONDS = 3 * 3600  # 3 hours
+CHUNK_SIZE = 1024 * 1024  # 1MB chunks
+
+# Blacklist of extensions (lowercase, without dot)
+DISALLOWED_EXT = {
+    "bat", "exe", "cmd", "sh", "msi", "ps1", "com", "scr"
+}
+
+# ensure tmp dir exists
+os.makedirs(TMP_DIR, exist_ok=True)
+
+app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)  # restrict /docs
+
+templates = Jinja2Templates(directory="templates")
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+
+def sanitize_slug(s: str) -> str:
+    s = re.sub(r"[^\w\-\.]", "", s)  # allow letters, numbers, underscore, hyphen, dot
+    return s[:128]
+
+
+def file_exists_for_slug(slug: str):
+    pattern = os.path.join(TMP_DIR, f"{slug}.*")
+    matches = glob.glob(pattern)
+    return matches[0] if matches else None
+
+
+def make_file_path(slug: str, filename: str):
+    _, ext = os.path.splitext(filename)
+    ext = ext.lower()
+    return os.path.join(TMP_DIR, f"{slug}{ext}")
+
+
+def gen_slug(length=8):
+    import secrets, string
+    alphabet = string.ascii_lowercase + string.digits
+    return ''.join(secrets.choice(alphabet) for _ in range(length))
+
+
+async def save_upload_to_tmp(upload_file: UploadFile, dest_path: str):
+    total = 0
+    # write in chunks
+    with open(dest_path, "wb") as f:
+        while True:
+            chunk = await upload_file.read(CHUNK_SIZE)
+            if not chunk:
+                break
+            f.write(chunk)
+            total += len(chunk)
+            if total > MAX_BYTES:
+                # cleanup partial file
+                f.close()
+                try:
+                    os.remove(dest_path)
+                except Exception:
+                    pass
+                raise HTTPException(status_code=413, detail="File exceeds max size (2GB).")
+    return total
+
+
+@app.on_event("startup")
+async def startup_event():
+    # launch cleanup background task
+    loop = asyncio.get_event_loop()
+    loop.create_task(cleaner_task())
+
+
+async def cleaner_task():
+    """
+    Periodically remove files older than EXPIRE_SECONDS to keep /tmp tidy.
+    """
+    while True:
+        try:
+            now = time.time()
+            for path in glob.glob(os.path.join(TMP_DIR, "*")):
+                try:
+                    # only remove files (and ignore directories)
+                    if os.path.isfile(path):
+                        mtime = os.path.getmtime(path)
+                        if now - mtime > EXPIRE_SECONDS:
+                            os.remove(path)
+                except Exception:
+                    continue
+        except Exception:
+            pass
+        await asyncio.sleep(CLEANUP_INTERVAL_SECONDS)
+
+
+@app.get("/", response_class=HTMLResponse)
+async def index(request: Request):
+    return templates.TemplateResponse("index.html", {
+        "request": request,
+        "api_version": API_VERSION,
+        "max_bytes": MAX_BYTES,
+        "expire_seconds": EXPIRE_SECONDS
+    })
+
+
+@app.post("/api/upload")
+async def api_upload(file: UploadFile = File(...), custom_slug: str = Form(None)):
+    filename = file.filename or "upload"
+    _, ext = os.path.splitext(filename)
+    ext_l = ext.lower().lstrip(".")
+    if ext_l in DISALLOWED_EXT:
+        raise HTTPException(status_code=400, detail=f"Disallowed file type: {ext}")
+
+    # choose slug
+    if custom_slug:
+        slug = sanitize_slug(custom_slug)
+        if not slug:
+            raise HTTPException(status_code=400, detail="Invalid custom slug.")
+        if file_exists_for_slug(slug):
+            raise HTTPException(status_code=409, detail="Slug already exists.")
+    else:
+        # generate until free
+        for _ in range(8):
+            slug = gen_slug(8)
+            if not file_exists_for_slug(slug):
+                break
+        else:
+            # fallback long slug
+            slug = gen_slug(16)
+
+    dest = make_file_path(slug, filename)
+
+    # save file (enforces max size)
+    try:
+        bytes_written = await save_upload_to_tmp(file, dest)
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        # cleanup if any partial file
+        try:
+            if os.path.exists(dest):
+                os.remove(dest)
+        except Exception:
+            pass
+        raise HTTPException(status_code=500, detail="Failed to save file.")
+
+    # set mtime so cleanup knows created time (already set)
+    expires_at = datetime.utcnow() + timedelta(seconds=EXPIRE_SECONDS)
+    url = f"/f/{slug}"
+    return JSONResponse({
+        "slug": slug,
+        "url": url,
+        "filename": filename,
+        "size": bytes_written,
+        "expires_at": int(expires_at.timestamp())
+    })
+
+
+@app.get("/f/{slug}")
+async def serve_file(slug: str, dl: int = 0):
+    # find file by slug
+    path = file_exists_for_slug(slug)
+    if not path:
+        raise HTTPException(status_code=404, detail="File not found or expired.")
+    # serve with correct media_type
+    mime_type, _ = mimetypes.guess_type(path)
+    headers = {}
+    filename = os.path.basename(path)
+    if dl:
+        # force download
+        return FileResponse(path, media_type=mime_type or "application/octet-stream",
+                            filename=filename)
+    # decide inline vs attachment by mime
+    inline_media = {"image", "video", "audio", "text", "application/pdf"}
+    mt = mime_type or ""
+    if any(mt.startswith(p) for p in inline_media) or mt == "application/pdf":
+        return FileResponse(path, media_type=mime_type or "application/octet-stream",
+                            filename=filename)
+    else:
+        return FileResponse(path, media_type=mime_type or "application/octet-stream",
+                            filename=filename)
+
+
+@app.get("/api/info")
+async def api_info():
+    curl_example = (
+        "curl -X POST -H \"Accept: application/json\" "
+        "-F \"file=@/path/to/file\" "
+        "https://your-space-url/f --output -"
+    )
+    return {
+        "version": API_VERSION,
+        "upload_endpoint": "/api/upload",
+        "file_endpoint_example": "/f/<slug>",
+        "max_size_bytes": MAX_BYTES,
+        "expiry_seconds": EXPIRE_SECONDS,
+        "curl_example": curl_example,
+    }