hsbatakurki/HuggingMes

1#!/usr/bin/env python32from __future__ import annotations3 4"""Create or reuse a Cloudflare Worker for Space keep-awake."""5 6import json7import os8import re9import sys10import time11import urllib.request12import urllib.error13from pathlib import Path14 15API_BASE = "https://api.cloudflare.com/client/v4"16KEEPALIVE_STATUS_FILE = Path("/tmp/huggingmes-cloudflare-keepalive-status.json")17 18 19def cf_request(method: str, path: str, token: str, body: bytes | None = None, content_type: str = "application/json"):20    req = urllib.request.Request(21        f"{API_BASE}{path}",22        data=body,23        method=method,24        headers={"Authorization": f"Bearer {token}", "Content-Type": content_type},25    )26    try:27        with urllib.request.urlopen(req, timeout=30) as response:28            payload = json.loads(response.read().decode("utf-8"))29    except urllib.error.HTTPError as e:30        try:31            error_body = json.loads(e.read().decode("utf-8"))32            errors = error_body.get("errors") or [{"message": "Unknown error"}]33            error_msg = errors[0].get("message", "Unknown error") if errors else "Unknown error"34        except:35            error_msg = f"HTTP {e.code}: {e.reason}"36        raise RuntimeError(f"Cloudflare API {e.code}: {error_msg}")37    if not payload.get("success"):38        errors = payload.get("errors") or [{"message": "Unknown Cloudflare API error"}]39        raise RuntimeError(errors[0].get("message", "Unknown Cloudflare API error"))40    return payload["result"]41 42 43def slugify(value: str) -> str:44    cleaned = re.sub(r"[^a-z0-9-]+", "-", value.lower()).strip("-")45    cleaned = re.sub(r"-{2,}", "-", cleaned)46    return (cleaned or "huggingmes-proxy")[:63].rstrip("-")47 48 49def get_space_host() -> str:50    space_host = os.environ.get("SPACE_HOST", "").strip()51    if space_host:52        return space_host53    54    author = os.environ.get("SPACE_AUTHOR_NAME", "").strip()55    repo = os.environ.get("SPACE_REPO_NAME", "").strip()56    if author and repo:57        return f"{author}-{repo}.hf.space".lower()58    59    return ""60 61 62def derive_keepalive_worker_name() -> str:63    explicit = os.environ.get("CLOUDFLARE_KEEPALIVE_WORKER_NAME", "").strip()64    if explicit:65        return slugify(explicit)66    space_host = get_space_host()67    if space_host:68        return slugify(f"{space_host.replace('.hf.space', '')}-keepalive")69    return "huggingmes-keepalive"70 71 72def render_keepalive_worker(target_url: str) -> str:73    return f"""addEventListener("fetch", (event) => {{74  event.respondWith(handleRequest(event.request));75}});76 77addEventListener("scheduled", (event) => {{78  event.waitUntil(ping("cron"));79}});80 81const TARGET_URL = {json.dumps(target_url)};82 83async function ping(source) {{84  const startedAt = new Date().toISOString();85  try {{86    const response = await fetch(TARGET_URL, {{87      method: "GET",88      headers: {{89        "user-agent": "HuggingMes Cloudflare KeepAlive",90        "cache-control": "no-cache"91      }},92      cf: {{ cacheTtl: 0, cacheEverything: false }}93    }});94    return {{95      ok: response.ok,96      status: response.status,97      source,98      target: TARGET_URL,99      timestamp: startedAt100    }};101  }} catch (error) {{102    return {{103      ok: false,104      status: 0,105      source,106      target: TARGET_URL,107      timestamp: startedAt,108      error: error.message109    }};110  }}111}}112 113async function handleRequest(request) {{114  const url = new URL(request.url);115  if (url.pathname === "/" || url.pathname === "/health" || url.pathname === "/ping") {{116    const result = await ping("manual");117    return new Response(JSON.stringify(result, null, 2), {{118      status: result.ok ? 200 : 502,119      headers: {{ "content-type": "application/json; charset=utf-8" }}120    }});121  }}122  return new Response("Not found", {{ status: 404 }});123}}124"""125 126 127def write_keepalive_status(payload: dict) -> None:128    payload = {129        **payload,130        "timestamp": payload.get("timestamp") or time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),131    }132    KEEPALIVE_STATUS_FILE.write_text(json.dumps(payload), encoding="utf-8")133    try:134        KEEPALIVE_STATUS_FILE.chmod(0o600)135    except OSError:136        pass137 138 139def resolve_account_and_subdomain(api_token: str) -> tuple[str, str]:140    account_id = os.environ.get("CLOUDFLARE_ACCOUNT_ID", "").strip()141    if not account_id:142        accounts = cf_request("GET", "/accounts", api_token)143        if not accounts:144            raise RuntimeError("No Cloudflare account is available for this token.")145        account_id = accounts[0]["id"]146 147    subdomain_info = cf_request("GET", f"/accounts/{account_id}/workers/subdomain", api_token)148    subdomain = (subdomain_info or {}).get("subdomain", "").strip()149    if not subdomain:150        raise RuntimeError("Cloudflare Workers subdomain is not configured. Enable workers.dev first.")151    return account_id, subdomain152 153 154def setup_keepalive_worker(api_token: str, account_id: str, subdomain: str) -> None:155    enabled = os.environ.get("CLOUDFLARE_KEEPALIVE_ENABLED", "true").strip().lower()156    if enabled in {"0", "false", "no", "off"}:157        write_keepalive_status({"configured": False, "status": "disabled", "message": "Cloudflare keep-awake is disabled."})158        return159 160    space_host = get_space_host()161    if not space_host:162        write_keepalive_status({"configured": False, "status": "skipped", "message": "SPACE_HOST could not be determined."})163        return164 165    cron = os.environ.get("CLOUDFLARE_KEEPALIVE_CRON", "*/10 * * * *").strip()166    space_host = space_host.removeprefix("https://").removeprefix("http://").split("/")[0]167    target_url = os.environ.get("CLOUDFLARE_KEEPALIVE_URL", f"https://{space_host}/health").strip()168    worker_name = derive_keepalive_worker_name()169    worker_source = render_keepalive_worker(target_url)170 171    cf_request(172        "PUT",173        f"/accounts/{account_id}/workers/scripts/{worker_name}",174        api_token,175        body=worker_source.encode("utf-8"),176        content_type="application/javascript",177    )178    cf_request(179        "POST",180        f"/accounts/{account_id}/workers/scripts/{worker_name}/subdomain",181        api_token,182        body=json.dumps({"enabled": True, "previews_enabled": True}).encode("utf-8"),183    )184    cf_request(185        "PUT",186        f"/accounts/{account_id}/workers/scripts/{worker_name}/schedules",187        api_token,188        body=json.dumps([{"cron": cron}]).encode("utf-8"),189    )190 191    worker_url = f"https://{worker_name}.{subdomain}.workers.dev"192    write_keepalive_status(193        {194            "configured": True,195            "status": "configured",196            "workerName": worker_name,197            "workerUrl": worker_url,198            "targetUrl": target_url,199            "cron": cron,200            "message": f"Cloudflare Worker cron pings {target_url} on {cron}.",201        }202    )203 204 205def main() -> int:206    api_token = os.environ.get("CLOUDFLARE_WORKERS_TOKEN", "").strip()207 208    if not api_token:209        return 0210 211    try:212        account_id, subdomain = resolve_account_and_subdomain(api_token)213        setup_keepalive_worker(api_token, account_id, subdomain)214        return 0215    except Exception as exc:216        print(f"Cloudflare keepalive setup failed: {exc}", file=sys.stderr)217        write_keepalive_status({"configured": False, "status": "error", "message": str(exc)})218        return 1219 220 221if __name__ == "__main__":222    raise SystemExit(main())223