chengyongruandXubin Ren
2a98360105
Maintainer edit for PR 4115: rebase onto origin/main and split gateway HTTP routing from token, media, and workspace services so WebSocketChannel depends on explicit gateway services instead of GatewayHTTPHandler internals. Preserve file edit channel capabilities and restore tools.restrict_to_workspace wiring through ChannelManager.
83 lines
2.7 KiB
Python
83 lines
2.7 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, *, api_token: bool = False) -> str:
|
|
token_value = f"nbwt_{secrets.token_urlsafe(32)}"
|
|
expiry = time.monotonic() + float(ttl_s)
|
|
self.issued_tokens[token_value] = expiry
|
|
if api_token:
|
|
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}
|