server.py

#!/usr/bin/env python3
# Combined static file server + Token API for Hugging Face Spaces (Docker)
# Serves token-service.html at '/'
# API endpoints under '/api':
#   - POST /api/set_token    { token: str, timestamp?: int, source?: str }
#   - GET  /api/token        { token: str|null, nonce: str, updated_at: int|null }
#   - GET  /api/status       { is_set: bool, updated_at: int|null }
# Uses only Python standard library.

import json
import os
import time
import uuid
from http.server import SimpleHTTPRequestHandler, HTTPServer
from urllib.parse import urlparse
from pathlib import Path

ROOT = Path(__file__).parent.resolve()
STATIC_DIR = ROOT / "static"
PORT = int(os.environ.get("PORT", "7860"))  # HF Spaces typically sets PORT=7860
HOST = os.environ.get("HOST", "0.0.0.0")
SET_SECRET = os.environ.get("TOKEN_SET_SECRET")  # optional; require header X-Token-Set-Secret

storage = {
    "token": None,
    "updated_at": None,
}


def now_ms() -> int:
    return int(time.time() * 1000)


class Handler(SimpleHTTPRequestHandler):
    # Serve from STATIC_DIR by default
    def translate_path(self, path):
        # Map '/' to token-service.html
        parsed = urlparse(path)
        p = parsed.path
        if p == "/":
            return str(STATIC_DIR / "token-service.html")
        # Ensure path traversal safe
        safe = Path(p.lstrip("/"))
        return str((STATIC_DIR / safe).resolve())

    # CORS helpers
    def _set_cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "*, Content-Type, X-Token-Set-Secret")
        self.send_header("Access-Control-Max-Age", "600")

    def _send_json(self, code: int, obj):
        body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self._set_cors()
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_OPTIONS(self):
        if self.path.startswith("/api/"):
            self.send_response(204)
            self._set_cors()
            self.end_headers()
        else:
            super().do_OPTIONS()

    def do_GET(self):
        if self.path.startswith("/api/"):
            parsed = urlparse(self.path)
            if parsed.path == "/api/token":
                self._send_json(200, {
                    "token": storage["token"],
                    "nonce": str(uuid.uuid4()),
                    "updated_at": storage["updated_at"],
                })
                return
            if parsed.path == "/api/status":
                self._send_json(200, {
                    "is_set": storage["token"] is not None,
                    "updated_at": storage["updated_at"],
                })
                return
            self._send_json(404, {"error": "not_found"})
            return
        # Static files
        return super().do_GET()

    def do_POST(self):
        if self.path.startswith("/api/"):
            parsed = urlparse(self.path)
            if parsed.path == "/api/set_token":
                if SET_SECRET:
                    header_secret = self.headers.get("X-Token-Set-Secret")
                    if header_secret != SET_SECRET:
                        self._send_json(403, {"error": "forbidden"})
                        return
                try:
                    length = int(self.headers.get("Content-Length") or 0)
                    data = json.loads(self.rfile.read(length) or b"{}")
                except Exception:
                    self._send_json(400, {"error": "invalid_json"})
                    return
                token = data.get("token")
                if not token or not isinstance(token, str):
                    self._send_json(400, {"error": "missing_token"})
                    return
                storage["token"] = token
                storage["updated_at"] = int(data.get("timestamp") or now_ms())
                print(f"[token-server] token set at {storage['updated_at']}")
                self._send_json(200, {"success": True})
                return
            self._send_json(404, {"error": "not_found"})
            return
        # Not an API path
        self.send_error(405, "Method Not Allowed")

    # Silence access log a bit
    def log_message(self, fmt, *args):
        if self.path.startswith("/api/"):
            return
        super().log_message(fmt, *args)


def main():
    # Ensure static dir exists
    STATIC_DIR.mkdir(parents=True, exist_ok=True)
    httpd = HTTPServer((HOST, PORT), Handler)
    print(f"[hf-token-service] listening on http://{HOST}:{PORT}")
    print(" - Static: / -> static/token-service.html")
    print(" - API:    /api/token, /api/status, /api/set_token")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\n[hf-token-service] shutting down")


if __name__ == "__main__":
    main()