import json import logging from tornado.websocket import WebSocketHandler from tfw.networking import Scope LOG = logging.getLogger(__name__) class ZMQWebSocketRouter(WebSocketHandler): # pylint: disable=abstract-method instances = set() def initialize(self, **kwargs): # pylint: disable=arguments-differ self.event_handler_connector = kwargs['event_handler_connector'] self.tfw_router = TFWRouter(self.send_to_zmq, self.send_to_websockets) def send_to_zmq(self, message): self.event_handler_connector.send_message(message) @classmethod def send_to_websockets(cls, message): for instance in cls.instances: instance.write_message(message) def prepare(self): type(self).instances.add(self) def on_close(self): type(self).instances.remove(self) def open(self, *args, **kwargs): LOG.debug('WebSocket connection initiated!') self.event_handler_connector.register_callback(self.zmq_callback) def zmq_callback(self, message): LOG.debug('Received on ZMQ pull socket: %s', message) self.tfw_router.route(message) def on_message(self, message): message = json.loads(message) LOG.debug('Received on WebSocket: %s', message) self.tfw_router.route(message) # much secure, very cors, wow def check_origin(self, origin): return True class TFWRouter: def __init__(self, send_to_zmq, send_to_websockets): self.send_to_zmq = send_to_zmq self.send_to_websockets = send_to_websockets def route(self, message): scope = Scope(message.get('scope', 'zmq')) routing_table = { Scope.ZMQ: self.send_to_zmq, Scope.WEBSOCKET: self.send_to_websockets, Scope.BROADCAST: self.broadcast } action = routing_table[scope] action(message) def broadcast(self, message): self.send_to_zmq(message) self.send_to_websockets(message)