mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 08:52:45 +00:00
453 lines
15 KiB
Python
453 lines
15 KiB
Python
"""
|
|
S58 — Bridge Token Lifecycle v2.
|
|
|
|
Manages bridge token issuance, rotation with overlap windows, revocation,
|
|
and expiry enforcement. All mutations emit structured audit events.
|
|
|
|
Security properties:
|
|
- Tokens have bounded lifetimes (expires_at)
|
|
- Rotation provides a controlled overlap window for seamless handover
|
|
- Revocation takes immediate effect
|
|
- Lifecycle decisions are deterministic and auditable
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import secrets
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
try:
|
|
from .sidecar.bridge_contract import BridgeScope, DeviceToken, TokenStatus
|
|
except ImportError:
|
|
from services.sidecar.bridge_contract import ( # type: ignore
|
|
BridgeScope,
|
|
DeviceToken,
|
|
TokenStatus,
|
|
)
|
|
|
|
logger = logging.getLogger("ComfyUI-OpenClaw.services.bridge_token_lifecycle")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Defaults
|
|
# ---------------------------------------------------------------------------
|
|
|
|
DEFAULT_TTL_SEC = 3600 # 1 hour
|
|
MAX_TTL_SEC = 86400 # 24 hours
|
|
DEFAULT_OVERLAP_SEC = 300 # 5 minutes
|
|
MAX_OVERLAP_SEC = 1800 # 30 minutes
|
|
TOKEN_BYTE_LENGTH = 32 # 256-bit random tokens
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Validation result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class TokenValidationResult:
|
|
"""Result of token validation."""
|
|
|
|
ok: bool
|
|
token: Optional[DeviceToken] = None
|
|
reject_reason: str = ""
|
|
is_overlap: bool = False # True if token accepted within overlap window
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audit events
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class TokenAuditEvent:
|
|
"""Structured audit event for token lifecycle operations."""
|
|
|
|
timestamp: float
|
|
action: str # issue, rotate, revoke, expire, validate_reject
|
|
token_id: str
|
|
device_id: str
|
|
details: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"timestamp": self.timestamp,
|
|
"action": self.action,
|
|
"token_id": self.token_id,
|
|
"device_id": self.device_id,
|
|
"details": self.details,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bridge Token Store
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class BridgeTokenStore:
|
|
"""
|
|
In-memory + persisted bridge token registry.
|
|
|
|
All tokens are stored by token_id and resolved by bounded constant-time scan.
|
|
Persistence is optional (state_dir may be None for test usage).
|
|
"""
|
|
|
|
MAX_TOKENS_PER_DEVICE = 5 # Cap active tokens per device
|
|
MAX_AUDIT_TRAIL = 200 # Cap total audit entries
|
|
|
|
def __init__(self, state_dir: Optional[str] = None):
|
|
self._tokens: Dict[str, DeviceToken] = {} # token_id → DeviceToken
|
|
self._audit_trail: List[TokenAuditEvent] = []
|
|
self._state_dir = state_dir
|
|
self._store_path: Optional[Path] = None
|
|
if state_dir:
|
|
self._store_path = Path(state_dir) / "bridge_tokens.json"
|
|
self._load()
|
|
|
|
# --- Persistence ---
|
|
|
|
def _load(self) -> None:
|
|
"""Load persisted tokens from disk."""
|
|
if not self._store_path or not self._store_path.exists():
|
|
return
|
|
try:
|
|
with open(self._store_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
for td in data.get("tokens", []):
|
|
token = DeviceToken(
|
|
device_id=td["device_id"],
|
|
device_token=td["device_token"],
|
|
scopes=[BridgeScope(s) for s in td.get("scopes", [])],
|
|
expires_at=td.get("expires_at"),
|
|
token_id=td.get("token_id", ""),
|
|
issued_at=td.get("issued_at", 0.0),
|
|
status=td.get("status", TokenStatus.ACTIVE.value),
|
|
replaces=td.get("replaces", ""),
|
|
overlap_until=td.get("overlap_until"),
|
|
)
|
|
self._tokens[token.token_id] = token
|
|
logger.info(f"S58: Loaded {len(self._tokens)} bridge tokens")
|
|
except Exception as e:
|
|
logger.error(f"S58: Failed to load bridge tokens: {e}")
|
|
|
|
def _save(self) -> None:
|
|
"""Persist tokens to disk."""
|
|
if not self._store_path:
|
|
return
|
|
try:
|
|
os.makedirs(self._store_path.parent, exist_ok=True)
|
|
tokens_data = []
|
|
for t in self._tokens.values():
|
|
tokens_data.append(
|
|
{
|
|
"device_id": t.device_id,
|
|
"device_token": t.device_token,
|
|
"scopes": [
|
|
s.value if isinstance(s, BridgeScope) else s
|
|
for s in t.scopes
|
|
],
|
|
"expires_at": t.expires_at,
|
|
"token_id": t.token_id,
|
|
"issued_at": t.issued_at,
|
|
"status": t.status,
|
|
"replaces": t.replaces,
|
|
"overlap_until": t.overlap_until,
|
|
}
|
|
)
|
|
temp = self._store_path.with_suffix(".tmp")
|
|
with open(temp, "w", encoding="utf-8") as f:
|
|
json.dump({"tokens": tokens_data}, f, indent=2)
|
|
temp.replace(self._store_path)
|
|
except Exception as e:
|
|
logger.error(f"S58: Failed to persist bridge tokens: {e}")
|
|
|
|
# --- Token lookup ---
|
|
|
|
def _resolve_token_for_value(
|
|
self, token_value: str
|
|
) -> Tuple[Optional[str], Optional[DeviceToken]]:
|
|
# IMPORTANT: keep lookup on bounded constant-time comparison instead of
|
|
# a derived hash index; hashing the presented token is the residual sink.
|
|
for token_id, token in self._tokens.items():
|
|
if secrets.compare_digest(token.device_token, token_value):
|
|
return token_id, token
|
|
return None, None
|
|
|
|
# --- Audit ---
|
|
|
|
def _emit_audit(
|
|
self, action: str, token_id: str, device_id: str, **details: Any
|
|
) -> None:
|
|
"""Record a structured audit event."""
|
|
event = TokenAuditEvent(
|
|
timestamp=time.time(),
|
|
action=action,
|
|
token_id=token_id,
|
|
device_id=device_id,
|
|
details=details,
|
|
)
|
|
self._audit_trail.append(event)
|
|
# Cap trail size
|
|
if len(self._audit_trail) > self.MAX_AUDIT_TRAIL:
|
|
self._audit_trail = self._audit_trail[-self.MAX_AUDIT_TRAIL :]
|
|
logger.info(f"S58: Audit [{action}] token={token_id} device={device_id}")
|
|
|
|
def get_audit_trail(self, *, device_id: Optional[str] = None) -> List[dict]:
|
|
"""Return audit trail, optionally filtered by device_id."""
|
|
trail = self._audit_trail
|
|
if device_id:
|
|
trail = [e for e in trail if e.device_id == device_id]
|
|
return [e.to_dict() for e in trail]
|
|
|
|
# --- Issue ---
|
|
|
|
def issue_token(
|
|
self,
|
|
device_id: str,
|
|
scopes: Optional[List[BridgeScope]] = None,
|
|
ttl_sec: int = DEFAULT_TTL_SEC,
|
|
) -> DeviceToken:
|
|
"""
|
|
Issue a new bridge token for a device.
|
|
|
|
Args:
|
|
device_id: Target device identifier
|
|
scopes: Authorized scopes (defaults to read-only)
|
|
ttl_sec: Token lifetime in seconds (capped at MAX_TTL_SEC)
|
|
"""
|
|
ttl_sec = min(max(ttl_sec, 60), MAX_TTL_SEC) # Bounded 60s..24h
|
|
now = time.time()
|
|
|
|
token_id = f"bt_{secrets.token_hex(8)}"
|
|
token_value = secrets.token_urlsafe(TOKEN_BYTE_LENGTH)
|
|
|
|
token = DeviceToken(
|
|
device_id=device_id,
|
|
device_token=token_value,
|
|
scopes=scopes or [BridgeScope.JOB_STATUS, BridgeScope.CONFIG_READ],
|
|
expires_at=now + ttl_sec,
|
|
token_id=token_id,
|
|
issued_at=now,
|
|
status=TokenStatus.ACTIVE.value,
|
|
)
|
|
|
|
self._tokens[token_id] = token
|
|
self._emit_audit("issue", token_id, device_id, ttl_sec=ttl_sec)
|
|
self._save()
|
|
|
|
return token
|
|
|
|
# --- Rotate ---
|
|
|
|
def rotate_token(
|
|
self,
|
|
old_token_id: str,
|
|
overlap_sec: int = DEFAULT_OVERLAP_SEC,
|
|
ttl_sec: int = DEFAULT_TTL_SEC,
|
|
scopes: Optional[List[BridgeScope]] = None,
|
|
) -> Tuple[DeviceToken, DeviceToken]:
|
|
"""
|
|
Rotate a bridge token, returning (new_token, old_token_updated).
|
|
|
|
The old token stays valid until min(old.expires_at, now + overlap_sec).
|
|
After overlap_until, the old token is deterministically rejected.
|
|
"""
|
|
old_token = self._tokens.get(old_token_id)
|
|
if not old_token:
|
|
raise ValueError(f"Token {old_token_id} not found")
|
|
if old_token.status != TokenStatus.ACTIVE.value:
|
|
raise ValueError(
|
|
f"Cannot rotate non-active token (status={old_token.status})"
|
|
)
|
|
|
|
overlap_sec = min(max(overlap_sec, 30), MAX_OVERLAP_SEC)
|
|
now = time.time()
|
|
|
|
# Set overlap window on old token
|
|
old_expires = old_token.expires_at or (now + 86400)
|
|
old_token.overlap_until = min(now + overlap_sec, old_expires)
|
|
|
|
# Issue new token linked to old
|
|
new_token = self.issue_token(
|
|
device_id=old_token.device_id,
|
|
scopes=scopes or old_token.scopes,
|
|
ttl_sec=ttl_sec,
|
|
)
|
|
new_token.replaces = old_token_id
|
|
|
|
self._tokens[old_token_id] = old_token
|
|
self._tokens[new_token.token_id] = new_token
|
|
self._emit_audit(
|
|
"rotate",
|
|
new_token.token_id,
|
|
old_token.device_id,
|
|
old_token_id=old_token_id,
|
|
overlap_sec=overlap_sec,
|
|
)
|
|
self._save()
|
|
|
|
return new_token, old_token
|
|
|
|
# --- Revoke ---
|
|
|
|
def revoke_token(self, token_id: str, reason: str = "") -> DeviceToken:
|
|
"""
|
|
Immediately revoke a token. Takes effect on next validation.
|
|
"""
|
|
token = self._tokens.get(token_id)
|
|
if not token:
|
|
raise ValueError(f"Token {token_id} not found")
|
|
|
|
token.status = TokenStatus.REVOKED.value
|
|
self._tokens[token_id] = token
|
|
self._emit_audit(
|
|
"revoke",
|
|
token_id,
|
|
token.device_id,
|
|
reason=reason,
|
|
)
|
|
self._save()
|
|
|
|
return token
|
|
|
|
# --- Validate ---
|
|
|
|
def validate_token(
|
|
self, token_value: str, required_scope: Optional[str] = None
|
|
) -> TokenValidationResult:
|
|
"""
|
|
Validate a token value against the store.
|
|
|
|
Checks (in order):
|
|
1. Token exists
|
|
2. Token not revoked
|
|
3. Token not expired (respects overlap_until window)
|
|
4. Required scope (if specified)
|
|
|
|
Returns TokenValidationResult with ok, reject_reason, and token metadata.
|
|
"""
|
|
token_id, token = self._resolve_token_for_value(token_value)
|
|
if token_id is None or token is None:
|
|
return TokenValidationResult(ok=False, reject_reason="unknown_token")
|
|
|
|
now = time.time()
|
|
|
|
# Check revocation (immediate, non-negotiable)
|
|
if token.status == TokenStatus.REVOKED.value:
|
|
self._emit_audit(
|
|
"validate_reject", token_id, token.device_id, reason="revoked"
|
|
)
|
|
return TokenValidationResult(
|
|
ok=False, token=token, reject_reason="token_revoked"
|
|
)
|
|
|
|
# Check expiry
|
|
if token.expires_at and now > token.expires_at:
|
|
token.status = TokenStatus.EXPIRED.value
|
|
self._tokens[token_id] = token
|
|
self._emit_audit("expire", token_id, token.device_id)
|
|
return TokenValidationResult(
|
|
ok=False, token=token, reject_reason="token_expired"
|
|
)
|
|
|
|
# Check overlap window (old token in rotation — still valid but soon to expire)
|
|
is_overlap = False
|
|
if token.overlap_until:
|
|
if now > token.overlap_until:
|
|
# Overlap window passed — reject deterministically
|
|
token.status = TokenStatus.EXPIRED.value
|
|
self._tokens[token_id] = token
|
|
self._emit_audit(
|
|
"validate_reject",
|
|
token_id,
|
|
token.device_id,
|
|
reason="overlap_window_expired",
|
|
)
|
|
return TokenValidationResult(
|
|
ok=False, token=token, reject_reason="overlap_window_expired"
|
|
)
|
|
is_overlap = True
|
|
|
|
# Check scope
|
|
if required_scope:
|
|
token_scopes = {
|
|
s.value if isinstance(s, BridgeScope) else s for s in token.scopes
|
|
}
|
|
if required_scope not in token_scopes:
|
|
self._emit_audit(
|
|
"validate_reject",
|
|
token_id,
|
|
token.device_id,
|
|
reason="insufficient_scope",
|
|
required=required_scope,
|
|
)
|
|
return TokenValidationResult(
|
|
ok=False, token=token, reject_reason="insufficient_scope"
|
|
)
|
|
|
|
return TokenValidationResult(ok=True, token=token, is_overlap=is_overlap)
|
|
|
|
# --- Listing ---
|
|
|
|
def list_tokens(
|
|
self, *, device_id: Optional[str] = None, active_only: bool = False
|
|
) -> List[DeviceToken]:
|
|
"""List tokens, optionally filtered."""
|
|
tokens = list(self._tokens.values())
|
|
if device_id:
|
|
tokens = [t for t in tokens if t.device_id == device_id]
|
|
if active_only:
|
|
tokens = [t for t in tokens if t.status == TokenStatus.ACTIVE.value]
|
|
return tokens
|
|
|
|
# --- Cleanup ---
|
|
|
|
def cleanup_expired(self) -> int:
|
|
"""Remove expired/revoked tokens. Returns count removed."""
|
|
now = time.time()
|
|
to_remove = []
|
|
for tid, token in self._tokens.items():
|
|
if token.status == TokenStatus.REVOKED.value:
|
|
to_remove.append(tid)
|
|
elif token.expires_at and now > token.expires_at:
|
|
to_remove.append(tid)
|
|
elif token.overlap_until and now > token.overlap_until:
|
|
to_remove.append(tid)
|
|
for tid in to_remove:
|
|
self._tokens.pop(tid)
|
|
if to_remove:
|
|
self._save()
|
|
logger.info(f"S58: Cleaned up {len(to_remove)} expired/revoked tokens")
|
|
return len(to_remove)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module-level singleton
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_store: Optional[BridgeTokenStore] = None
|
|
|
|
|
|
def get_token_store(state_dir: Optional[str] = None) -> BridgeTokenStore:
|
|
"""Get or create the global bridge token store."""
|
|
global _store
|
|
if _store is None:
|
|
if state_dir is None:
|
|
try:
|
|
from .state_dir import get_state_dir
|
|
|
|
state_dir = get_state_dir()
|
|
except ImportError:
|
|
try:
|
|
from services.state_dir import get_state_dir # type: ignore
|
|
|
|
state_dir = get_state_dir()
|
|
except ImportError:
|
|
state_dir = None
|
|
_store = BridgeTokenStore(state_dir)
|
|
return _store
|