Spaces:
Running
Running
File size: 5,023 Bytes
2b6ef22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | 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")
@dataclass
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
|