""" WebSocket client for Stockbit real-time data. Handles connection, reconnection, protobuf subscription, and message decoding. """ import asyncio import json import logging from typing import Optional, Callable, Dict, Any, List from dataclasses import dataclass, field try: import websockets except ModuleNotFoundError: websockets = None # Optional; install websockets package for WS support from .auth import AuthManager from .models import LivePrice, OrderBook, OrderBookEntry, RunningTrade logger = logging.getLogger(__name__) @dataclass class SubscriptionChannels: watchlist: List[str] = field(default_factory=list) order_book: List[str] = field(default_factory=list) running_trade: List[str] = field(default_factory=list) is_hotlist: bool = False running_trade_batch: List[str] = field(default_factory=list) liveprice: List[str] = field(default_factory=list) liveprice_v3: List[str] = field(default_factory=list) iepiev: List[str] = field(default_factory=list) intraday: List[str] = field(default_factory=list) intraday_v3: List[str] = field(default_factory=list) best_bid_offer: List[str] = field(default_factory=list) order_book_v3: List[str] = field(default_factory=list) def to_dict(self) -> dict: d = {} for key, value in self.__dict__.items(): if value: d[key] = value return d class StockbitWebSocket: WS_URL = "wss://wss-trading.stockbit.com/ws" def __init__(self, auth: AuthManager, ws_type: str = "", channels: Optional[SubscriptionChannels] = None): self._auth = auth self._ws: Optional[websockets.WebSocketClientProtocol] = None self._callbacks: Dict[str, List[Callable]] = {} self._running = False self._channels: Optional[SubscriptionChannels] = channels or SubscriptionChannels() self._ws_type = ws_type # "chart" or empty for main def on(self, channel: str, callback: Callable): if channel not in self._callbacks: self._callbacks[channel] = [] self._callbacks[channel].append(callback) def on_message(self, callback: Callable): self.on("_all", callback) def subscribe(self, channels: SubscriptionChannels): self._channels = channels def subscribe_dict(self, **kwargs): channels = SubscriptionChannels() for key, value in kwargs.items(): if hasattr(channels, key): setattr(channels, key, value) self.subscribe(channels) def _ws_url(self) -> str: if self._ws_type: return f"{self.WS_URL}?type={self._ws_type}" return self.WS_URL async def connect(self): wskey = self._get_wskey() token = self._auth.main_token self._ws = await websockets.connect( self._ws_url(), subprotocols=["web"], ping_interval=30, ping_timeout=10, ) if self._channels: await self._send_subscription(wskey, token) logger.info("WebSocket connected (%s)", self._ws_url()) def _get_wskey(self) -> str: import requests resp = requests.get( f"{self._auth.EXODUS_BASE}/auth/websocket/key", headers={"Authorization": f"Bearer {self._auth.main_token}"}, ) resp.raise_for_status() return resp.json()["data"]["key"] async def _send_subscription(self, wskey: str, token: str): if not self._ws: return uid = self._get_user_id() channels = self._channels or SubscriptionChannels() try: from stockbit_api.proto.stockbit_pb2 import WebsocketRequest auth_msg = WebsocketRequest() auth_msg.user_id = uid auth_msg.key = wskey auth_msg.access_token = token sub_msg = WebsocketRequest() sub_msg.user_id = uid sub_msg.key = wskey sub_msg.access_token = token if channels.liveprice: sub_msg.channel.liveprice.extend(channels.liveprice) if channels.running_trade_batch: sub_msg.channel.running_trade_batch.extend(channels.running_trade_batch) if channels.order_book: sub_msg.channel.order_book.extend(channels.order_book) if channels.iepiev: sub_msg.channel.iepiev.extend(channels.iepiev) await self._ws.send(auth_msg.SerializeToString()) await self._ws.send(sub_msg.SerializeToString()) except ImportError: logger.warning("protobuf not available, sending JSON") payload = { "user_id": uid, "channel": channels.to_dict(), "key": wskey, "access_token": token, } await self._ws.send(json.dumps(payload).encode()) def _get_user_id(self) -> str: import jwt as pyjwt try: decoded = pyjwt.decode(self._auth.main_token, options={"verify_signature": False}) return str(decoded.get("uid", "")) except Exception: return "" async def run(self): self._running = True reconnect_attempts = 0 ping_task = None while self._running: try: await self.connect() reconnect_attempts = 0 if self._ws_type == "chart": ping_task = asyncio.create_task(self._chart_ping_loop()) async for message in self._ws: await self._handle_message(message) except websockets.exceptions.ConnectionClosed as e: logger.warning(f"WebSocket closed: {e.code} {e.reason}") except Exception as e: logger.error(f"WebSocket error: {e}") if not self._running: break reconnect_attempts += 1 delay = min(2 ** reconnect_attempts, 60) logger.info(f"Reconnecting in {delay}s (attempt {reconnect_attempts})") await asyncio.sleep(delay) async def _handle_message(self, message: bytes): try: if isinstance(message, bytes): data = self._decode_frame(message) else: data = json.loads(message) channel = data.get("channel", data.get("case", "unknown")) self._emit(channel, data) self._emit("_all", data) except Exception as e: logger.error(f"Failed to handle message: {e}") def _decode_frame(self, frame: bytes) -> dict: try: text = frame.decode("utf-8") return json.loads(text) except (UnicodeDecodeError, json.JSONDecodeError): pass return self._decode_protobuf_frame(frame) def _decode_protobuf_frame(self, frame: bytes) -> dict: try: import sys sys.path.insert(0, ".") try: from stockbit_api.proto.stockbit_pb2 import (WebsocketWrapMessage, IEPIEV) msg = WebsocketWrapMessage() msg.ParseFromString(frame) which = msg.WhichOneof("event") result = {"event": which} if which == "liveprice": lp = msg.liveprice result["stock"] = lp.stock_code result["lastprice"] = lp.lastprice result["volume"] = lp.volume result["high"] = lp.high result["low"] = lp.low result["prev"] = lp.prev elif which == "running_trade": trades = [] for t in msg.running_trade.trades: trades.append({ "stock": t.stock, "price": t.price, "volume": t.volume, "action": t.action, }) result["trades"] = trades elif which == "orderbook": ob = msg.orderbook result["stock"] = ob.symbol result["orderbook_data"] = ob.orderbook_data.hex() result["timestamp"] = ob.timestamp return result except ImportError: pass except Exception: pass except Exception: pass # Fallback: try to decode as IEPIEV directly try: iepiev = IEPIEV() iepiev.ParseFromString(frame) if iepiev.stock_code: return { "event": "iepiev", "stock": iepiev.stock_code, "iep": iepiev.iep, "iev": iepiev.iev, "iep_change": iepiev.iep_change, "total_buy_lot": iepiev.total_buy_lot, "total_sell_lot": iepiev.total_sell_lot, } except Exception: pass return {"channel": "unknown", "raw": frame.hex()} def _emit(self, channel: str, data: dict): for cb in self._callbacks.get(channel, []): try: cb(data) except Exception as e: logger.error(f"Callback error on {channel}: {e}") async def _chart_ping_loop(self): try: from stockbit_api.proto.stockbit_pb2 import WebsocketRequest ping_req = WebsocketRequest() ping_req.ping.message = "ping" while self._running and self._ws: await asyncio.sleep(30) if self._ws: await self._ws.send(ping_req.SerializeToString()) except asyncio.CancelledError: pass async def close(self): self._running = False if self._ws: await self._ws.close() @classmethod def for_chart(cls, auth: AuthManager, symbols: List[str]) -> "StockbitWebSocket": channels = SubscriptionChannels(liveprice=symbols, iepiev=symbols) return cls(auth, ws_type="chart", channels=channels) async def __aenter__(self): return self async def __aexit__(self, *args): await self.close()