baseimage-tutorial-framework/lib/tfw/internals/server/zmq_websocket_router.py

import json
import logging

from tornado.websocket import WebSocketHandler

from tfw.internals.networking import Scope

LOG = logging.getLogger(__name__)


class ZMQWebSocketRouter(WebSocketHandler):
    # pylint: disable=abstract-method,attribute-defined-outside-init
    instances = set()

    def initialize(self, **kwargs):
        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.pop('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)