baseimage-tutorial-framework/lib/tfw/networking/server/zmq_websocket_handler.py

# Copyright (C) 2018 Avatao.com Innovative Learning Kft.
# All Rights Reserved. See LICENSE file for details.

import json
from abc import ABC, abstractmethod

from tornado.websocket import WebSocketHandler

from tfw.networking import deserialize_all
from tfw.config.logs import logging

LOG = logging.getLogger(__name__)


class ZMQWebSocketHandler(WebSocketHandler, ABC):
    instances = set()

    def initialize(self, **kwargs): # pylint: disable=arguments-differ
        self._event_handler_connector = kwargs['event_handler_connector']

    def prepare(self):
        ZMQWebSocketHandler.instances.add(self)

    def on_close(self):
        ZMQWebSocketHandler.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, msg_parts):
        keyhandlers = {'mirror': self.mirror}

        key, message = deserialize_all(*msg_parts)
        LOG.debug('Received on pull socket: %s', message)
        self.handle_trigger(message)
        if key not in keyhandlers:
            for instance in ZMQWebSocketHandler.instances:
                instance.write_message(message)
        else:
            try:
                keyhandlers[key](message['data'])
            except KeyError:
                LOG.error('Invalid mirror message format! Ignoring.')

    def mirror(self, data):
        key = data['key']
        self._event_handler_connector.send_message({'data': data}, key)

    def on_message(self, message):
        LOG.debug('Received on WebSocket: %s', message)
        self.send_message(self.make_response(message))

    @abstractmethod
    def make_response(self, message):
        raise NotImplementedError

    def send_message(self, message: dict, key: str = None):
        self._event_handler_connector.send_message(message, key)

    @abstractmethod
    def handle_trigger(self, message):
        raise NotImplementedError

    # much secure, very cors, wow
    def check_origin(self, origin):
        return True


class ZMQWebSocketProxy(ZMQWebSocketHandler):
    # pylint: disable=abstract-method
    def initialize(self, **kwargs): # pylint: disable=arguments-differ
        super(ZMQWebSocketProxy, self).initialize(**kwargs)
        self._make_response = kwargs['make_response']
        self._proxy_filter = kwargs['proxy_filter']
        self._handle_trigger = kwargs['handle_trigger']

    def on_message(self, message):
        message = json.loads(message)
        if self._proxy_filter(message):
            super().on_message(message)

    def make_response(self, message):
        return self._make_response(message)

    def handle_trigger(self, message):
        self._handle_trigger(message)