cloudflare-keepalive-setup.py

#!/usr/bin/env python3
from __future__ import annotations

"""Create or reuse a Cloudflare Worker for Space keep-awake.

Vendored verbatim from github.com/somratpro/HuggingMes.
"""

import json
import os
import re
import sys
import time
import urllib.request
import urllib.error
from pathlib import Path

API_BASE = "https://api.cloudflare.com/client/v4"
KEEPALIVE_STATUS_FILE = Path("/tmp/huggingmes-cloudflare-keepalive-status.json")


def cf_request(method: str, path: str, token: str, body: bytes | None = None, content_type: str = "application/json"):
    req = urllib.request.Request(
        f"{API_BASE}{path}",
        data=body,
        method=method,
        headers={"Authorization": f"Bearer {token}", "Content-Type": content_type},
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        try:
            error_body = json.loads(e.read().decode("utf-8"))
            errors = error_body.get("errors") or [{"message": "Unknown error"}]
            error_msg = errors[0].get("message", "Unknown error") if errors else "Unknown error"
        except Exception:
            error_msg = f"HTTP {e.code}: {e.reason}"
        raise RuntimeError(f"Cloudflare API {e.code}: {error_msg}")
    if not payload.get("success"):
        errors = payload.get("errors") or [{"message": "Unknown Cloudflare API error"}]
        raise RuntimeError(errors[0].get("message", "Unknown Cloudflare API error"))
    return payload["result"]


def slugify(value: str) -> str:
    cleaned = re.sub(r"[^a-z0-9-]+", "-", value.lower()).strip("-")
    cleaned = re.sub(r"-{2,}", "-", cleaned)
    return (cleaned or "huggingmes-keepalive")[:63].rstrip("-")


def get_space_host() -> str:
    space_host = os.environ.get("SPACE_HOST", "").strip()
    if space_host:
        return space_host

    author = os.environ.get("SPACE_AUTHOR_NAME", "").strip()
    repo = os.environ.get("SPACE_REPO_NAME", "").strip()
    if author and repo:
        return f"{author}-{repo}.hf.space".lower()

    return ""


def derive_keepalive_worker_name() -> str:
    explicit = os.environ.get("CLOUDFLARE_KEEPALIVE_WORKER_NAME", "").strip()
    if explicit:
        return slugify(explicit)
    space_host = get_space_host()
    if space_host:
        return slugify(f"{space_host.replace('.hf.space', '')}-keepalive")
    return "huggingmes-keepalive"


def render_keepalive_worker(target_url: str) -> str:
    return f"""addEventListener("fetch", (event) => {{
  event.respondWith(handleRequest(event.request));
}});

addEventListener("scheduled", (event) => {{
  event.waitUntil(ping("cron"));
}});

const TARGET_URL = {json.dumps(target_url)};

async function ping(source) {{
  const startedAt = new Date().toISOString();
  try {{
    const response = await fetch(TARGET_URL, {{
      method: "GET",
      headers: {{
        "user-agent": "HuggingMes Cloudflare KeepAlive",
        "cache-control": "no-cache"
      }},
      cf: {{ cacheTtl: 0, cacheEverything: false }}
    }});
    return {{
      ok: response.ok,
      status: response.status,
      source,
      target: TARGET_URL,
      timestamp: startedAt
    }};
  }} catch (error) {{
    return {{
      ok: false,
      status: 0,
      source,
      target: TARGET_URL,
      timestamp: startedAt,
      error: error.message
    }};
  }}
}}

async function handleRequest(request) {{
  const url = new URL(request.url);
  if (url.pathname === "/" || url.pathname === "/health" || url.pathname === "/ping") {{
    const result = await ping("manual");
    return new Response(JSON.stringify(result, null, 2), {{
      status: result.ok ? 200 : 502,
      headers: {{ "content-type": "application/json; charset=utf-8" }}
    }});
  }}
  return new Response("Not found", {{ status: 404 }});
}}
"""


def write_keepalive_status(payload: dict) -> None:
    payload = {
        **payload,
        "timestamp": payload.get("timestamp") or time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    KEEPALIVE_STATUS_FILE.write_text(json.dumps(payload), encoding="utf-8")
    try:
        KEEPALIVE_STATUS_FILE.chmod(0o600)
    except OSError:
        pass


def resolve_account_and_subdomain(api_token: str) -> tuple[str, str]:
    account_id = os.environ.get("CLOUDFLARE_ACCOUNT_ID", "").strip()
    if not account_id:
        accounts = cf_request("GET", "/accounts", api_token)
        if not accounts:
            raise RuntimeError("No Cloudflare account is available for this token.")
        account_id = accounts[0]["id"]

    subdomain_info = cf_request("GET", f"/accounts/{account_id}/workers/subdomain", api_token)
    subdomain = (subdomain_info or {}).get("subdomain", "").strip()
    if not subdomain:
        raise RuntimeError("Cloudflare Workers subdomain is not configured. Enable workers.dev first.")
    return account_id, subdomain


def setup_keepalive_worker(api_token: str, account_id: str, subdomain: str) -> None:
    enabled = os.environ.get("CLOUDFLARE_KEEPALIVE_ENABLED", "true").strip().lower()
    if enabled in {"0", "false", "no", "off"}:
        write_keepalive_status({"configured": False, "status": "disabled", "message": "Cloudflare keep-awake is disabled."})
        return

    space_host = get_space_host()
    if not space_host:
        write_keepalive_status({"configured": False, "status": "skipped", "message": "SPACE_HOST could not be determined."})
        return

    cron = os.environ.get("CLOUDFLARE_KEEPALIVE_CRON", "*/10 * * * *").strip()
    space_host = space_host.removeprefix("https://").removeprefix("http://").split("/")[0]
    target_url = os.environ.get("CLOUDFLARE_KEEPALIVE_URL", f"https://{space_host}/health").strip()
    worker_name = derive_keepalive_worker_name()
    worker_source = render_keepalive_worker(target_url)

    cf_request(
        "PUT",
        f"/accounts/{account_id}/workers/scripts/{worker_name}",
        api_token,
        body=worker_source.encode("utf-8"),
        content_type="application/javascript",
    )
    cf_request(
        "POST",
        f"/accounts/{account_id}/workers/scripts/{worker_name}/subdomain",
        api_token,
        body=json.dumps({"enabled": True, "previews_enabled": True}).encode("utf-8"),
    )
    cf_request(
        "PUT",
        f"/accounts/{account_id}/workers/scripts/{worker_name}/schedules",
        api_token,
        body=json.dumps([{"cron": cron}]).encode("utf-8"),
    )

    worker_url = f"https://{worker_name}.{subdomain}.workers.dev"
    write_keepalive_status(
        {
            "configured": True,
            "status": "configured",
            "workerName": worker_name,
            "workerUrl": worker_url,
            "targetUrl": target_url,
            "cron": cron,
            "message": f"Cloudflare Worker cron pings {target_url} on {cron}.",
        }
    )


def main() -> int:
    api_token = os.environ.get("CLOUDFLARE_WORKERS_TOKEN", "").strip()

    if not api_token:
        return 0

    try:
        account_id, subdomain = resolve_account_and_subdomain(api_token)
        setup_keepalive_worker(api_token, account_id, subdomain)
        return 0
    except Exception as exc:
        print(f"Cloudflare keepalive setup failed: {exc}", file=sys.stderr)
        write_keepalive_status({"configured": False, "status": "error", "message": str(exc)})
        return 1


if __name__ == "__main__":
    raise SystemExit(main())