chengyongruandXubin Ren
4ddd639e67
maintainer edit: handle review feedback by treating bootstrap responses without api_token as auth-required, and remove the obsolete issue_token(api_token=...) compatibility path now that API tokens are issued separately.
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
"""Token state for the embedded WebUI gateway."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from websockets.http11 import Request as WsRequest
|
|
|
|
from nanobot.webui.http_utils import bearer_token, parse_query, query_first
|
|
|
|
|
|
@dataclass
|
|
class GatewayTokenStore:
|
|
"""Own short-lived WebSocket and WebUI API tokens for one gateway process."""
|
|
|
|
max_tokens: int = 10_000
|
|
issued_tokens: dict[str, float] = field(default_factory=dict)
|
|
api_tokens: dict[str, float] = field(default_factory=dict)
|
|
|
|
def check_api_token(self, request: WsRequest) -> bool:
|
|
self._purge_expired_api_tokens()
|
|
token = bearer_token(request.headers) or query_first(
|
|
parse_query(request.path), "token"
|
|
)
|
|
if not token:
|
|
return False
|
|
expiry = self.api_tokens.get(token)
|
|
if expiry is None or time.monotonic() > expiry:
|
|
self.api_tokens.pop(token, None)
|
|
return False
|
|
return True
|
|
|
|
def can_issue(self, *, include_api_token: bool = False) -> bool:
|
|
self._purge_expired_issued_tokens()
|
|
self._purge_expired_api_tokens()
|
|
if len(self.issued_tokens) >= self.max_tokens:
|
|
return False
|
|
if include_api_token and len(self.api_tokens) >= self.max_tokens:
|
|
return False
|
|
return True
|
|
|
|
def issue_token(self, ttl_s: int | float) -> str:
|
|
token_value = f"nbwt_{secrets.token_urlsafe(32)}"
|
|
expiry = time.monotonic() + float(ttl_s)
|
|
self.issued_tokens[token_value] = expiry
|
|
return token_value
|
|
|
|
def issue_api_token(self, ttl_s: int | float) -> str:
|
|
token_value = f"nbwt_{secrets.token_urlsafe(32)}"
|
|
expiry = time.monotonic() + float(ttl_s)
|
|
self.api_tokens[token_value] = expiry
|
|
return token_value
|
|
|
|
def take_issued_token_if_valid(self, token_value: str | None) -> bool:
|
|
if not token_value:
|
|
return False
|
|
self._purge_expired_issued_tokens()
|
|
expiry = self.issued_tokens.pop(token_value, None)
|
|
if expiry is None:
|
|
return False
|
|
if time.monotonic() > expiry:
|
|
return False
|
|
return True
|
|
|
|
def clear(self) -> None:
|
|
self.issued_tokens.clear()
|
|
self.api_tokens.clear()
|
|
|
|
def _purge_expired_api_tokens(self) -> None:
|
|
now = time.monotonic()
|
|
for token_key, expiry in list(self.api_tokens.items()):
|
|
if now > expiry:
|
|
self.api_tokens.pop(token_key, None)
|
|
|
|
def _purge_expired_issued_tokens(self) -> None:
|
|
now = time.monotonic()
|
|
for token_key, expiry in list(self.issued_tokens.items()):
|
|
if now > expiry:
|
|
self.issued_tokens.pop(token_key, None)
|
|
|
|
|
|
def token_response_payload(token: str, expires_in: Any) -> dict[str, Any]:
|
|
return {"token": token, "expires_in": expires_in}
|