Blame for invite.py

"""Invite codes for the Emberwick game. A player with a code gets live decisions billed to the Space owner's Hugging Face
account (PRO ZeroGPU quota first, then prepaid credits) instead of their own small anonymous quota. Daily caps keep
that spend bounded; the counters live in memory and start again each UTC day or when the Space restarts."""

import hmac
import threading
from datetime import datetime, timezone


def norm(code: str) -> bytes:
    """Case, spaces and punctuation do not matter: "Emberwick invite" and "EMBERWICK-INVITE" are the same code."""
    return "".join(ch for ch in str(code or "").lower() if ch.isalnum()).encode()


def utc_day() -> str:
    return datetime.now(timezone.utc).date().isoformat()


class Invites:
    def __init__(self, codes: str, daily_units: int, ip_daily_units: int, ip_daily_misses: int = 30, today=utc_day):
        self.codes = [c for c in (norm(x) for x in codes.split(",")) if c]
        self.daily_units, self.ip_daily_units, self.ip_daily_misses = daily_units, ip_daily_units, ip_daily_misses
        self._today = today
        self._lock = threading.Lock()
        self._day, self._units, self._ip_units, self._misses = "", 0, {}, {}

    def _roll(self) -> None:
        if (day := self._today()) != self._day:
            self._day, self._units, self._ip_units, self._misses = day, 0, {}, {}

    def valid(self, code: str, ip: str) -> bool:
        """True for a known code. Wrong codes are counted per address; past the limit every check fails for the day."""
        given = norm(code)
        ok = bool(given) and any(hmac.compare_digest(given, c) for c in self.codes)
        with self._lock:
            self._roll()
            if self._misses.get(ip, 0) >= self.ip_daily_misses:
                return False
            if not ok:
                self._misses[ip] = self._misses.get(ip, 0) + 1
        return ok

    def charge(self, ip: str, units: int) -> bool:
        """Book `units` for today. False when the Space-wide or the per-address cap would be passed."""
        with self._lock:
            self._roll()
            if self._units + units > self.daily_units or self._ip_units.get(ip, 0) + units > self.ip_daily_units:
                return False
            self._units += units
            self._ip_units[ip] = self._ip_units.get(ip, 0) + units
            return True