import os, secrets, hmac, hashlib, httpx, re
import gradio as gr
import spaces
from fastapi import Request, Form
from fastapi.responses import HTMLResponse, RedirectResponse
ADMIN_TOKEN = os.environ.get("ADMIN_TOKEN", "")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
PRIVATE_URL = os.environ.get("PRIVATE_URL", "")
AUTH_HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
@spaces.GPU(duration=30)
def _gpu():
return ""
def make_token() -> str:
raw = secrets.token_hex(16)
key = ADMIN_TOKEN.encode() or b""
return raw + "." + hmac.new(key, raw.encode(), hashlib.sha256).hexdigest()
def check_token(t: str) -> bool:
if not t or "." not in t: return False
p = t.split(".")
k = ADMIN_TOKEN.encode() or b""
return hmac.compare_digest(p[1], hmac.new(k, p[0].encode(), hashlib.sha256).hexdigest())
def rewrite(html: str) -> str:
for q in ['"', "'"]:
for a in ["href", "src", "action"]:
html = html.replace(f'{a}={q}/', f'{a}={q}/px/')
return html
LOGIN = 'VESTA
VESTA
Acceso al ecosistema
'
with gr.Blocks(title="VESTA") as demo:
gr.Markdown("")
app = demo.app
def cookie(req: Request) -> str:
return req.cookies.get("vesta_session", "")
@app.get("/login")
async def login_get(req: Request):
t = cookie(req)
if check_token(t):
return RedirectResponse(url="/px/")
return HTMLResponse(LOGIN)
@app.post("/login")
async def do_login(user: str = Form(...), passw: str = Form(...)):
if user.lower() != "admin" or passw != ADMIN_TOKEN:
return RedirectResponse(url="/login?error=1", status_code=303)
token = make_token()
r = RedirectResponse(url="/px/", status_code=303)
r.set_cookie("vesta_session", token, max_age=86400, path="/", httponly=True)
return r
@app.get("/logout")
async def logout():
r = RedirectResponse(url="/login")
r.delete_cookie("vesta_session", path="/")
return r
@app.get("/px/{rest:path}")
async def proxy(req: Request, rest: str = ""):
t = cookie(req)
if not check_token(t):
return RedirectResponse(url="/login")
path = "/" + rest
try:
async with httpx.AsyncClient(timeout=30) as c:
r = await c.get(f"{PRIVATE_URL}{path}", headers=AUTH_HEADERS)
content = r.text
if "text/html" in r.headers.get("content-type", ""):
content = rewrite(content)
return HTMLResponse(content)
except Exception as e:
return HTMLResponse(f"Error: {e}", 502)
@app.post("/px/proxy/predict")
async def predict(req: Request):
if not check_token(cookie(req)):
return HTMLResponse("NO_AUTH", 401)
body = await req.json()
async with httpx.AsyncClient(timeout=120) as c:
r = await c.post(f"{PRIVATE_URL}/proxy/predict", json=body, headers=AUTH_HEADERS)
return HTMLResponse(r.text, r.status_code)
demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False, head="")