Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import hashlib | |
| import hmac | |
| import time | |
| import uuid | |
| from dataclasses import dataclass, field | |
| from typing import Any, Dict, List, Optional | |
| from app.core.logger import get_logger | |
| log = get_logger("webhook-socket") | |
| class Channel: | |
| channel_id: str | |
| secret: Optional[str] = None | |
| buffer_size: int = 0 | |
| created_at: float = field(default_factory=time.time) | |
| message_count: int = 0 | |
| last_activity: float = field(default_factory=time.time) | |
| subscribers: Dict[Any, asyncio.Queue] = field(default_factory=dict) | |
| history: List[dict] = field(default_factory=list) | |
| class ChannelManager: | |
| def __init__(self, default_buffer: int = 0): | |
| self.channels: Dict[str, Channel] = {} | |
| self.default_buffer = default_buffer | |
| self.total_messages = 0 | |
| def create_channel( | |
| self, | |
| channel_id: Optional[str] = None, | |
| secret: Optional[str] = None, | |
| buffer_size: Optional[int] = None, | |
| ) -> Channel: | |
| ch_id = channel_id or uuid.uuid4().hex[:16] | |
| buf = buffer_size if buffer_size is not None else self.default_buffer | |
| ch = Channel(channel_id=ch_id, secret=secret, buffer_size=buf) | |
| self.channels[ch_id] = ch | |
| log.info("Channel created id=%s buffer=%d secret=%s", ch_id, buf, "yes" if secret else "no") | |
| return ch | |
| def get_channel(self, channel_id: str) -> Optional[Channel]: | |
| return self.channels.get(channel_id) | |
| def delete_channel(self, channel_id: str) -> bool: | |
| if channel_id in self.channels: | |
| for q in self.channels[channel_id].subscribers.values(): | |
| q.put_nowait({"event": "channel_deleted", "channel": channel_id}) | |
| del self.channels[channel_id] | |
| log.info("Channel deleted id=%s", channel_id) | |
| return True | |
| return False | |
| def subscribe(self, channel_id: str, ws: Any) -> Optional[asyncio.Queue]: | |
| ch = self.channels.get(channel_id) | |
| if not ch: | |
| return None | |
| q: asyncio.Queue = asyncio.Queue() | |
| ch.subscribers[ws] = q | |
| log.info( | |
| "Subscriber joined channel=%s total_subs=%d", | |
| channel_id, len(ch.subscribers), | |
| ) | |
| for msg in ch.history: | |
| q.put_nowait(msg) | |
| return q | |
| def unsubscribe(self, channel_id: str, ws: Any) -> None: | |
| ch = self.channels.get(channel_id) | |
| if ch and ws in ch.subscribers: | |
| del ch.subscribers[ws] | |
| log.info( | |
| "Subscriber left channel=%s total_subs=%d", | |
| channel_id, len(ch.subscribers), | |
| ) | |
| async def publish( | |
| self, | |
| channel_id: str, | |
| payload: Any, | |
| headers: Optional[Dict[str, str]] = None, | |
| ) -> int: | |
| ch = self.channels.get(channel_id) | |
| if not ch: | |
| return -1 | |
| message = { | |
| "event": "message", | |
| "channel": channel_id, | |
| "timestamp": time.time(), | |
| "id": uuid.uuid4().hex[:12], | |
| "payload": payload, | |
| "headers": headers or {}, | |
| } | |
| ch.message_count += 1 | |
| ch.last_activity = time.time() | |
| self.total_messages += 1 | |
| if ch.buffer_size > 0: | |
| ch.history.append(message) | |
| while len(ch.history) > ch.buffer_size: | |
| ch.history.pop(0) | |
| dead: List[Any] = [] | |
| sent = 0 | |
| for ws, q in list(ch.subscribers.items()): | |
| if getattr(ws, "closed", False): | |
| dead.append(ws) | |
| continue | |
| await q.put(message) | |
| sent += 1 | |
| for ws in dead: | |
| del ch.subscribers[ws] | |
| log.info( | |
| "Published channel=%s subs=%d msg_total=%d", | |
| channel_id, sent, ch.message_count, | |
| ) | |
| return sent | |
| def stats(self) -> dict: | |
| return { | |
| "channels": len(self.channels), | |
| "total_messages": self.total_messages, | |
| "total_subscribers": sum(len(c.subscribers) for c in self.channels.values()), | |
| "channels_detail": { | |
| cid: { | |
| "subscribers": len(ch.subscribers), | |
| "messages": ch.message_count, | |
| "buffered": len(ch.history), | |
| "last_activity": ch.last_activity, | |
| "has_secret": ch.secret is not None, | |
| } | |
| for cid, ch in self.channels.items() | |
| }, | |
| } | |
| def sign_payload(secret: str, raw_body: bytes) -> str: | |
| return "sha256=" + hmac.new(secret.encode(), raw_body, hashlib.sha256).hexdigest() | |
| def verify_signature(secret: str, raw_body: bytes, signature: str) -> bool: | |
| expected = sign_payload(secret, raw_body) | |
| return hmac.compare_digest(expected, signature) | |
| _manager: Optional[ChannelManager] = None | |
| def get_manager() -> ChannelManager: | |
| global _manager | |
| if _manager is None: | |
| _manager = ChannelManager() | |
| return _manager | |