Blame for app.py

"""tracebin — a pastebin for agent traces, secured by unguessable URLs.

The link *is* the credential (a "capability URL"). That only holds if the app
never leaks a link and never lets anyone enumerate one, so:

* ids are 128 bits from ``secrets.token_urlsafe`` — not uuids, not hashes of the
  content, not anything derived from time
* there is no index, no listing endpoint, and no id appears in any response the
  caller didn't already have the id for
* ``Referrer-Policy: no-referrer`` and a CSP that forbids outbound requests, so
  a trace full of links can't hand the URL to a third party
* unknown and malformed ids get the same flat 404
* reads are rate limited per IP so the keyspace can't be scanned
"""

from __future__ import annotations

import json
import logging
import os
import re
import secrets
import time
import urllib.parse
from collections import defaultdict, deque
from datetime import datetime, timezone

from fastapi import FastAPI, Form, Header, HTTPException, Request
from fastapi.responses import (
    HTMLResponse,
    JSONResponse,
    PlainTextResponse,
    RedirectResponse,
    Response,
)
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

import auth
import render
from storage import NotFound, build_store

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s %(message)s")
log = logging.getLogger("tracebin")

ID_BYTES = 16  # 128 bits -> 22 url-safe chars
ID_RE = re.compile(r"^[A-Za-z0-9_-]{22}$")
MAX_BYTES = int(os.environ.get("MAX_TRACE_BYTES", 5 * 1024 * 1024))
UPLOAD_TOKEN = os.environ.get("UPLOAD_TOKEN", "").strip()
BASE_URL = os.environ.get("BASE_URL", "").strip().rstrip("/")

app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")
store = build_store()


# --------------------------------------------------------------------------
# rate limiting
# --------------------------------------------------------------------------

_hits: dict[str, deque[float]] = defaultdict(deque)

LIMITS = {"read": (120, 60.0), "write": (20, 60.0)}


def _client_ip(request: Request) -> str:
    # Spaces sit behind a proxy; the leftmost XFF entry is the real client.
    forwarded = request.headers.get("x-forwarded-for", "")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


def rate_limit(request: Request, bucket: str) -> None:
    limit, window = LIMITS[bucket]
    key = f"{bucket}:{_client_ip(request)}"
    now = time.monotonic()
    hits = _hits[key]
    while hits and now - hits[0] > window:
        hits.popleft()
    if len(hits) >= limit:
        raise HTTPException(status_code=429, detail="slow down")
    hits.append(now)
    if len(_hits) > 10_000:  # crude bound; single-process app
        for stale_key in [k for k, v in _hits.items() if not v or now - v[-1] > 600]:
            del _hits[stale_key]


# --------------------------------------------------------------------------
# security headers
# --------------------------------------------------------------------------

CSP = (
    "default-src 'none'; "
    "style-src 'self'; "
    "script-src 'self'; "
    "img-src 'self' data:; "
    "font-src 'self'; "
    "base-uri 'none'; "
    "form-action 'self'; "
    "frame-ancestors 'none'"
)


@app.middleware("http")
async def security_headers(request: Request, call_next):
    response = await call_next(request)
    # no-referrer is load-bearing: without it, clicking any link inside a trace
    # sends this page's secret URL to whatever site was linked.
    response.headers["Referrer-Policy"] = "no-referrer"
    response.headers["X-Robots-Tag"] = "noindex, nofollow, noarchive, nosnippet"
    response.headers["Content-Security-Policy"] = CSP
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Cache-Control"] = "private, no-store"
    return response


# --------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------


def check_upload_token(supplied: str | None) -> None:
    """Writes are gated when UPLOAD_TOKEN is set, so a public Space can't be
    used as free storage by anyone who finds it."""
    if not UPLOAD_TOKEN:
        return
    if not supplied or not secrets.compare_digest(supplied.strip(), UPLOAD_TOKEN):
        raise HTTPException(status_code=401, detail="bad or missing upload token")


def new_id() -> str:
    return secrets.token_urlsafe(ID_BYTES)


def trace_url(request: Request, trace_id: str) -> str:
    if BASE_URL:
        return f"{BASE_URL}/t/{trace_id}"
    return str(request.url_for("view_trace", trace_id=trace_id))


def save(payload, title: str, allowed: list[str] | None = None) -> str:
    body = json.dumps(payload, ensure_ascii=False, default=str)
    if len(body.encode("utf-8")) > MAX_BYTES:
        raise HTTPException(status_code=413, detail=f"trace exceeds {MAX_BYTES} bytes")
    if allowed and not auth.oauth_enabled():
        raise HTTPException(
            status_code=400,
            detail="restricted traces need OAuth; this deployment has none configured",
        )
    trace_id = new_id()
    store.put(
        trace_id,
        {
            "id": trace_id,
            "created": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "title": (title or "").strip()[:200],
            "allowed": allowed or None,
            "payload": payload,
        },
    )
    # Deliberately not logging the id — Space logs are a plaintext copy of every
    # capability URL otherwise.
    log.info("stored trace (%d bytes)", len(body))
    return trace_id


def not_found() -> HTTPException:
    """One response for unknown ids, malformed ids, and anything else missing,
    so a probe learns nothing from the difference."""
    return HTTPException(status_code=404, detail="not found")


def load(trace_id: str) -> dict:
    if not ID_RE.match(trace_id):
        raise not_found()
    try:
        return store.get(trace_id)
    except NotFound:
        raise not_found()


def safe_next(target: str | None) -> str:
    """Only ever redirect back into this site — an attacker must not be able to
    turn /login into an open redirect."""
    if not target or not target.startswith("/") or target.startswith("//"):
        return "/"
    return target


def redirect_uri(request: Request) -> str:
    if BASE_URL:
        return f"{BASE_URL}/auth/callback"
    return str(request.url_for("oauth_callback"))


def agent_prompt(url: str, restricted: bool) -> str:
    """A ready-to-paste instruction for handing a trace to another agent."""
    note = (
        " This link requires a Hugging Face sign-in, so fetch it with a browser "
        "session rather than a plain HTTP request."
        if restricted
        else ""
    )
    return (
        f"Download the agent trace at {url} (JSON) and analyse it: summarise what "
        f"the agent was asked to do, the steps and tool calls it made, where it "
        f"went wrong, and what it should have done instead.{note}"
    )


# --------------------------------------------------------------------------
# routes
# --------------------------------------------------------------------------


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse(
        request,
        "index.html",
        {
            "needs_token": bool(UPLOAD_TOKEN),
            "oauth": auth.oauth_enabled(),
            "viewer": auth.session_username(request),
            "api_url": f"{BASE_URL}/api/traces" if BASE_URL else "/api/traces",
        },
    )


@app.post("/", response_class=HTMLResponse)
async def create_from_form(
    request: Request,
    content: str = Form(...),
    title: str = Form(""),
    token: str = Form(""),
    allowed: str = Form(""),
):
    rate_limit(request, "write")
    check_upload_token(token or None)

    content = content.strip()
    if not content:
        raise HTTPException(status_code=400, detail="nothing to paste")
    try:
        payload = json.loads(content)
    except ValueError:
        payload = content  # plain text is a perfectly good trace

    users = auth.normalize_users(allowed)
    trace_id = save(payload, title, users)
    url = trace_url(request, trace_id)
    return templates.TemplateResponse(
        request,
        "created.html",
        {
            "url": url,
            "raw_url": url + "/raw",
            "allowed": users,
            "prompt": agent_prompt(url, bool(users)),
        },
        status_code=201,
    )


@app.post("/api/traces")
async def create_from_api(
    request: Request,
    x_upload_token: str | None = Header(default=None),
    authorization: str | None = Header(default=None),
):
    rate_limit(request, "write")

    supplied = x_upload_token
    if not supplied and authorization and authorization.lower().startswith("bearer "):
        supplied = authorization[7:]
    check_upload_token(supplied)

    raw = await request.body()
    if len(raw) > MAX_BYTES:
        raise HTTPException(status_code=413, detail=f"trace exceeds {MAX_BYTES} bytes")
    try:
        body = json.loads(raw)
    except ValueError:
        raise HTTPException(status_code=400, detail="body must be JSON")

    # {"trace": ..., "title": ...} is the explicit form; anything else is taken
    # as the trace itself.
    if isinstance(body, dict) and "trace" in body:
        payload = body["trace"]
        title = str(body.get("title", ""))
        users = auth.normalize_users(body.get("allowed_users"))
    else:
        payload, title, users = body, "", []

    trace_id = save(payload, title, users)
    url = trace_url(request, trace_id)
    return JSONResponse(
        {
            "id": trace_id,
            "url": url,
            "raw_url": url + "/raw",
            "allowed_users": users or None,
        },
        status_code=201,
    )


@app.get("/t/{trace_id}", response_class=HTMLResponse, name="view_trace")
async def view_trace(request: Request, trace_id: str):
    rate_limit(request, "read")
    record = load(trace_id)

    allowed = record.get("allowed")
    viewer = auth.session_username(request)
    if not auth.is_allowed(viewer, allowed):
        if viewer is None and auth.oauth_enabled():
            return RedirectResponse(
                f"/login?next={urllib.parse.quote(f'/t/{trace_id}', safe='')}",
                status_code=302,
            )
        # Signed in but not on the list: indistinguishable from a trace that
        # does not exist. No point confirming it to someone who can't read it.
        raise not_found()

    payload = record.get("payload")
    url = trace_url(request, trace_id)
    return templates.TemplateResponse(
        request,
        "trace.html",
        {
            "trace_id": trace_id,
            "title": record.get("title") or "",
            "created": record.get("created", ""),
            "summary": render.summarize(payload),
            "view": render.normalize(payload),
            "url": url,
            "raw_url": url + "/raw",
            "prompt": agent_prompt(url, bool(allowed)),
            "allowed": allowed,
            "viewer": viewer,
        },
    )


@app.get("/t/{trace_id}/raw")
async def raw_trace(request: Request, trace_id: str, download: int = 0):
    rate_limit(request, "read")
    record = load(trace_id)
    if not auth.is_allowed(auth.session_username(request), record.get("allowed")):
        # A data endpoint has nowhere useful to send a browser, so no redirect.
        raise not_found()

    payload = record.get("payload")
    if isinstance(payload, str):
        body, media = payload, "text/plain; charset=utf-8"
        suffix = "txt"
    else:
        body = json.dumps(payload, indent=2, ensure_ascii=False, default=str)
        media, suffix = "application/json", "json"

    headers = {}
    if download:
        headers["Content-Disposition"] = (
            f'attachment; filename="trace-{trace_id}.{suffix}"'
        )
    return Response(body, media_type=media, headers=headers)


# --------------------------------------------------------------------------
# sign-in with Hugging Face
# --------------------------------------------------------------------------


@app.get("/login")
def login(request: Request, next: str = "/"):
    if not auth.oauth_enabled():
        raise HTTPException(status_code=503, detail="sign-in is not configured")
    state = auth.new_state()
    response = RedirectResponse(
        auth.authorize_url(redirect_uri(request), state), status_code=302
    )
    response.set_cookie(
        auth.STATE_COOKIE,
        auth.sign({"s": state, "n": safe_next(next), "exp": int(time.time()) + 600}),
        httponly=True,
        secure=request.url.scheme == "https",
        samesite="lax",
        max_age=600,
        path="/",
    )
    return response


@app.get("/auth/callback", name="oauth_callback")
def oauth_callback(request: Request, code: str = "", state: str = ""):
    if not auth.oauth_enabled():
        raise HTTPException(status_code=503, detail="sign-in is not configured")
    pending = auth.unsign(request.cookies.get(auth.STATE_COOKIE))
    if not (code and state and pending and secrets.compare_digest(pending.get("s", ""), state)):
        # Covers CSRF, a replayed callback, and an expired attempt alike.
        raise HTTPException(status_code=400, detail="sign-in failed, please retry")

    try:
        username = auth.fetch_username(auth.exchange_code(code, redirect_uri(request)))
    except Exception:
        log.exception("oauth callback failed")
        raise HTTPException(status_code=502, detail="sign-in failed, please retry")

    response = RedirectResponse(safe_next(pending.get("n")), status_code=302)
    response.set_cookie(
        auth.SESSION_COOKIE,
        auth.make_session(username),
        httponly=True,
        secure=request.url.scheme == "https",
        samesite="lax",
        max_age=auth.SESSION_TTL,
        path="/",
    )
    response.delete_cookie(auth.STATE_COOKIE, path="/")
    return response


@app.get("/logout")
async def logout(request: Request, next: str = "/"):
    response = RedirectResponse(safe_next(next), status_code=302)
    response.delete_cookie(auth.SESSION_COOKIE, path="/")
    return response


@app.get("/robots.txt", response_class=PlainTextResponse)
async def robots():
    return "User-agent: *\nDisallow: /\n"


@app.get("/healthz", response_class=PlainTextResponse)
async def healthz():
    return "ok"


@app.exception_handler(HTTPException)
async def http_error(request: Request, exc: HTTPException):
    accepts_html = "text/html" in request.headers.get("accept", "")
    if accepts_html and not request.url.path.startswith("/api/"):
        return templates.TemplateResponse(
            request,
            "error.html",
            {"code": exc.status_code, "detail": exc.detail},
            status_code=exc.status_code,
        )
    return JSONResponse({"error": exc.detail}, status_code=exc.status_code)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))