mirror of
https://github.com/BigBodyCobain/Shadowbroker.git
synced 2026-08-18 08:27:18 +02:00
60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
"""Short-lived, single-use WebSocket bootstrap tokens for the agent shell."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import time
|
|
from threading import Lock
|
|
|
|
_TOKEN_TTL_SECONDS = 60.0
|
|
_MAX_ACTIVE_TOKENS = 256
|
|
|
|
_store: dict[str, float] = {}
|
|
_lock = Lock()
|
|
|
|
|
|
def _purge_expired_locked(now: float) -> None:
|
|
"""Drop expired tokens while the caller holds ``_lock``."""
|
|
expired = [token for token, expires in _store.items() if expires <= now]
|
|
for token in expired:
|
|
_store.pop(token, None)
|
|
|
|
|
|
def _purge_expired() -> None:
|
|
now = time.time()
|
|
with _lock:
|
|
_purge_expired_locked(now)
|
|
|
|
|
|
def mint_agent_shell_ws_token() -> tuple[str, int]:
|
|
"""Return (token, expires_in_seconds) without exceeding the store bound."""
|
|
token = secrets.token_urlsafe(32)
|
|
now = time.time()
|
|
expires_at = now + _TOKEN_TTL_SECONDS
|
|
with _lock:
|
|
_purge_expired_locked(now)
|
|
if len(_store) >= _MAX_ACTIVE_TOKENS:
|
|
# Evict the token that will expire soonest. This keeps the store
|
|
# bounded without recursively acquiring the non-reentrant lock.
|
|
oldest = min(_store, key=_store.get)
|
|
_store.pop(oldest, None)
|
|
_store[token] = expires_at
|
|
return token, int(_TOKEN_TTL_SECONDS)
|
|
|
|
|
|
def consume_agent_shell_ws_token(token: str) -> bool:
|
|
"""Validate and burn a one-time token. Returns True when accepted."""
|
|
cleaned = str(token or "").strip()
|
|
if not cleaned:
|
|
return False
|
|
now = time.time()
|
|
with _lock:
|
|
_purge_expired_locked(now)
|
|
expires_at = _store.pop(cleaned, None)
|
|
return expires_at is not None and expires_at > now
|
|
|
|
|
|
def reset_agent_shell_ws_tokens_for_tests() -> None:
|
|
with _lock:
|
|
_store.clear()
|