baseimage-tutorial-framework/lib/tfw/event_handler_base.py

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

from abc import ABC, abstractmethod

from tfw.networking.event_handlers import ServerConnector
from tfw.crypto import message_checksum, KeyManager, verify_message
from tfw.config.logs import logging

LOG = logging.getLogger(__name__)


class EventHandlerBase(ABC):
    """
    Abstract base class for all Python based EventHandlers. Useful implementation template
    for other languages.

    Derived classes must implement the handle_event() method
    """
    def __init__(self, key):
        self.server_connector = ServerConnector()
        self.keys = [key]
        self.subscribe(*self.keys)
        self.server_connector.register_callback(self.event_handler_callback)

    @property
    def key(self):
        """
        Returns the oldest key this EventHandler was subscribed to.
        """
        return self.keys[0]

    def event_handler_callback(self, message):
        """
        Callback that is invoked when receiving a message.
        Dispatches messages to handler methods and sends
        a response back in case the handler returned something.
        This is subscribed in __init__().
        """
        if not self.check_key(message):
            return

        response = self.dispatch_handling(message)
        if response:
            self.server_connector.send(response)

    def check_key(self, message):
        """
        Checks whether the message is intended for this
        EventHandler.

        This is necessary because ZMQ handles PUB - SUB
        connetions with pattern matching (e.g. someone
        subscribed to 'fsm' will receive 'fsm_update'
        messages as well.
        """
        return message['key'] in self.keys

    def dispatch_handling(self, message):
        """
        Used to dispatch messages to their specific handlers.

        :param message: the message received
        :returns: the message to send back
        """
        return self.handle_event(message)

    @abstractmethod
    def handle_event(self, message):
        """
        Abstract method that implements the handling of messages.

        :param message: the message received
        :returns: the message to send back
        """
        raise NotImplementedError

    def subscribe(self, *keys):
        """
        Subscribe this EventHandler to receive events for given keys.
        Note that you can subscribe to the same key several times in which
        case you will need to unsubscribe multiple times in order to stop
        receiving events.

        :param keys: list of keys to subscribe to
        """
        for key in keys:
            self.server_connector.subscribe(key)
            self.keys.append(key)

    def unsubscribe(self, *keys):
        """
        Unsubscribe this eventhandler from the given keys.

        :param keys: list of keys to unsubscribe from
        """
        for key in keys:
            self.server_connector.unsubscribe(key)
            self.keys.remove(key)

    def cleanup(self):
        """
        Perform cleanup actions such as releasing database
        connections and stuff like that.
        """
        pass


class FSMAwareEventHandler(EventHandlerBase, ABC):
    # pylint: disable=abstract-method
    """
    Abstract base class for EventHandlers which automatically
    keep track of the state of the TFW FSM.
    """
    def __init__(self, key):
        super().__init__(key)
        self.subscribe('fsm_update')
        self.fsm_state = None
        self.in_accepted_state = False
        self._auth_key = KeyManager().auth_key

    def dispatch_handling(self, message):
        if message['key'] == 'fsm_update':
            if verify_message(self._auth_key, message):
                self._handle_fsm_update(message)
            return None
        return super().dispatch_handling(message)

    def _handle_fsm_update(self, message):
        try:
            new_state = message['data']['current_state']
            trigger = message['data']['last_trigger']
            if self.fsm_state != new_state:
                self.handle_fsm_step(self.fsm_state, new_state, trigger)
            self.fsm_state = new_state
            self.in_accepted_state = message['data']['in_accepted_state']
        except KeyError:
            LOG.error('Invalid fsm_update message received!')

    def handle_fsm_step(self, from_state, to_state, trigger):
        """
        Called in case the TFW FSM has stepped.
        """
        pass


class BroadcastingEventHandler(EventHandlerBase, ABC):
    # pylint: disable=abstract-method
    """
    Abstract base class for EventHandlers which broadcast responses
    and intelligently ignore their own broadcasted messages they receive.
    """
    def __init__(self, key):
        super().__init__(key)
        self.own_message_hashes = []

    def event_handler_callback(self, message):
        message_hash = message_checksum(message)

        if message_hash in self.own_message_hashes:
            self.own_message_hashes.remove(message_hash)
            return

        response = self.dispatch_handling(message)
        if response:
            self.own_message_hashes.append(message_checksum(response))
            self.server_connector.broadcast(response)