Spaces:
Sleeping
Sleeping
| import hashlib | |
| import hmac | |
| import os | |
| import time | |
| from typing import Optional | |
| from urllib.parse import quote, urlencode | |
| from fastapi import FastAPI, Header, HTTPException, Query, Request | |
| from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse | |
| from huggingface_hub import HfApi, HfFileSystem | |
| HF_TOKEN = os.environ["HF_TOKEN"] | |
| SIGN_SECRET = os.environ["SIGN_SECRET"].encode() | |
| ADMIN_TOKEN = os.environ["PROXY_ADMIN_TOKEN"] | |
| ALLOWED_BUCKETS = {b.strip() for b in os.environ.get("ALLOWED_BUCKETS", "").split(",") if b.strip()} | |
| DEFAULT_TTL = int(os.environ.get("DEFAULT_TTL", "86400")) | |
| MAX_TTL = int(os.environ.get("MAX_TTL", str(7 * 86400))) | |
| CHUNK = 1024 * 1024 # 1 MiB | |
| fs = HfFileSystem(token=HF_TOKEN) | |
| api = HfApi(token=HF_TOKEN) | |
| app = FastAPI(title="bucket-proxy") | |
| def _sign(bucket: str, path: str, exp: int) -> str: | |
| msg = f"{bucket}|{path}|{exp}".encode() | |
| return hmac.new(SIGN_SECRET, msg, hashlib.sha256).hexdigest() | |
| def _check_bucket(bucket: str) -> None: | |
| if ALLOWED_BUCKETS and bucket not in ALLOWED_BUCKETS: | |
| raise HTTPException(403, f"bucket {bucket} not in allowlist") | |
| def _check_admin(token_query: Optional[str], authorization: Optional[str]) -> None: | |
| token = token_query | |
| if not token and authorization and authorization.startswith("Bearer "): | |
| token = authorization[7:] | |
| if not token: | |
| raise HTTPException(401, "missing admin token (query 't' or Bearer header)") | |
| if not hmac.compare_digest(token, ADMIN_TOKEN): | |
| raise HTTPException(403, "bad admin token") | |
| def health(): | |
| return {"ok": True, "ts": int(time.time())} | |
| INDEX_HTML = """<!doctype html> | |
| <html><head><meta charset=utf-8><title>bucket-proxy</title> | |
| <style> | |
| :root{color-scheme:dark light} | |
| body{font-family:system-ui,sans-serif;max-width:780px;margin:2rem auto;padding:0 1rem} | |
| h1{margin:0 0 1rem;font-size:1.4rem} | |
| label{display:block;margin:.6rem 0 .2rem;font-size:.85rem;opacity:.8} | |
| input,select,button{font:inherit;padding:.5rem .6rem;border:1px solid #888;border-radius:6px;background:transparent;color:inherit;width:100%;box-sizing:border-box} | |
| button{cursor:pointer;background:#3b82f6;color:#fff;border:none;font-weight:600} | |
| button:hover{background:#2563eb} | |
| .row{display:flex;gap:.6rem} | |
| .row>*{flex:1} | |
| .out{margin-top:1rem;padding:.8rem;border:1px solid #888;border-radius:6px;word-break:break-all;font-family:ui-monospace,monospace;font-size:.85rem;white-space:pre-wrap} | |
| .err{border-color:#ef4444;color:#ef4444} | |
| .copy{margin-top:.5rem;width:auto;padding:.3rem .8rem;font-size:.8rem} | |
| small{opacity:.6} | |
| </style></head><body> | |
| <h1>bucket-proxy · sign URL</h1> | |
| <label>Admin token</label> | |
| <input id=t type=password placeholder="PROXY_ADMIN_TOKEN" autocomplete=off> | |
| <label>Bucket <small>(namespace/name)</small></label> | |
| <input id=b placeholder="Philasil/comfyui-models"> | |
| <label>Path <small>(inside bucket)</small></label> | |
| <input id=p placeholder="loras/foo.safetensors"> | |
| <div class=row> | |
| <div><label>TTL (seconds)</label><input id=ttl type=number value=86400 min=60></div> | |
| <div style="display:flex;align-items:end"><button id=go>Sign</button></div> | |
| </div> | |
| <div id=out class=out style=display:none></div> | |
| <button id=cp class=copy style=display:none>Copy URL</button> | |
| <script> | |
| const $=id=>document.getElementById(id); | |
| ['t','b'].forEach(k=>{const v=localStorage.getItem('bp_'+k);if(v)$(k).value=v;}); | |
| $('go').onclick=async()=>{ | |
| const t=$('t').value.trim(),b=$('b').value.trim(),p=$('p').value.trim(),ttl=$('ttl').value; | |
| if(!t||!b||!p)return; | |
| localStorage.setItem('bp_t',t);localStorage.setItem('bp_b',b); | |
| $('out').style.display='block';$('out').className='out';$('out').textContent='...';$('cp').style.display='none'; | |
| const u=`/sign?bucket=${encodeURIComponent(b)}&path=${encodeURIComponent(p)}&ttl=${encodeURIComponent(ttl)}&t=${encodeURIComponent(t)}`; | |
| try{ | |
| const r=await fetch(u);const j=await r.json(); | |
| if(!r.ok){$('out').className='out err';$('out').textContent=JSON.stringify(j,null,2);return;} | |
| $('out').textContent=`URL:\\n${j.url}\\n\\nexpires: ${new Date(j.expires_at*1000).toLocaleString()}\\nsize: ${j.size} bytes`; | |
| $('cp').style.display='inline-block';$('cp').dataset.url=j.url; | |
| }catch(e){$('out').className='out err';$('out').textContent=String(e);} | |
| }; | |
| $('cp').onclick=()=>{navigator.clipboard.writeText($('cp').dataset.url);$('cp').textContent='Copied!';setTimeout(()=>$('cp').textContent='Copy URL',1200);}; | |
| </script> | |
| </body></html>""" | |
| def index(): | |
| return INDEX_HTML | |
| def sign( | |
| request: Request, | |
| bucket: str = Query(..., description="namespace/bucket-name"), | |
| path: str = Query(..., description="path inside the bucket"), | |
| ttl: int = Query(DEFAULT_TTL, ge=60, le=MAX_TTL), | |
| t: Optional[str] = Query(None, description="admin token (alternative to Bearer)"), | |
| authorization: Optional[str] = Header(None), | |
| ): | |
| _check_admin(t, authorization) | |
| _check_bucket(bucket) | |
| path = path.lstrip("/") | |
| # Try model repo first (no prefix), then datasets/, then buckets/ | |
| candidates = [f"{bucket}/{path}", f"datasets/{bucket}/{path}", f"buckets/{bucket}/{path}"] | |
| full = None | |
| for c in candidates: | |
| try: | |
| if fs.exists(c): | |
| full = c | |
| break | |
| except Exception: | |
| continue | |
| if full is None: | |
| raise HTTPException(404, f"file not found: {bucket}/{path}") | |
| exp = int(time.time()) + ttl | |
| sig = _sign(bucket, path, exp) | |
| host = request.headers.get("x-forwarded-host") or request.url.netloc | |
| scheme = request.headers.get("x-forwarded-proto", "https") | |
| base = f"{scheme}://{host}" | |
| qs = urlencode({"b": bucket, "f": path, "exp": exp, "sig": sig}, quote_via=quote) | |
| return { | |
| "url": f"{base}/d?{qs}", | |
| "expires_at": exp, | |
| "ttl": ttl, | |
| "size": fs.info(full).get("size"), | |
| } | |
| def download( | |
| b: str = Query(..., description="bucket id namespace/name"), | |
| f: str = Query(..., description="path inside bucket"), | |
| exp: int = Query(...), | |
| sig: str = Query(...), | |
| ): | |
| _check_bucket(b) | |
| if exp < int(time.time()): | |
| raise HTTPException(410, "link expired") | |
| expected = _sign(b, f, exp) | |
| if not hmac.compare_digest(expected, sig): | |
| raise HTTPException(403, "bad signature") | |
| candidates = [f"{b}/{f}", f"datasets/{b}/{f}", f"buckets/{b}/{f}"] | |
| full = None | |
| info = None | |
| for c in candidates: | |
| try: | |
| info = fs.info(c) | |
| full = c | |
| break | |
| except Exception: | |
| continue | |
| if full is None: | |
| raise HTTPException(404, "file not found") | |
| size = info.get("size") | |
| def stream(): | |
| with fs.open(full, "rb") as src: | |
| while True: | |
| chunk = src.read(CHUNK) | |
| if not chunk: | |
| return | |
| yield chunk | |
| filename = f.rsplit("/", 1)[-1] | |
| headers = { | |
| "Content-Disposition": f'attachment; filename="{filename}"', | |
| "Cache-Control": "private, max-age=300", | |
| } | |
| if size: | |
| headers["Content-Length"] = str(size) | |
| return StreamingResponse(stream(), media_type="application/octet-stream", headers=headers) | |
| def http_exc(_: Request, exc: HTTPException): | |
| return JSONResponse({"error": exc.detail}, status_code=exc.status_code) | |