baseimage-tutorial-framework/lib/tfw/builtins/fsm_managing_event_handler.py

import logging

from tfw.crypto import KeyManager, sign_message, verify_message
from tfw.networking import Scope
from tfw.components import FSMUpdater

from .event_handler import EventHandler

LOG = logging.getLogger(__name__)


class FSMManagingEventHandler(EventHandler):
    """
    EventHandler responsible for managing the state machine of
    the framework (TFW FSM).

    tfw.networking.TFWServer instances automatically send 'trigger'
    commands to the event handler listening on the 'fsm' key,
    which should be an instance of this event handler.

    This event handler accepts messages that have a
    data['command'] key specifying a command to be executed.

    An 'fsm_update' message is broadcasted after every successful
    command.
    """
    def __init__(self, key, fsm_type, require_signature=False):
        super().__init__(key, scope=Scope.WEBSOCKET)
        self.fsm = fsm_type()
        self._fsm_updater = FSMUpdater(self.fsm)
        self.auth_key = KeyManager().auth_key
        self._require_signature = require_signature

        self.command_handlers = {
            'trigger': self.handle_trigger,
            'update':  self.handle_update
        }

    def handle_event(self, message):
        try:
            message = self.command_handlers[message['data']['command']](message)
            if message:
                fsm_update_message = self._fsm_updater.fsm_update
                sign_message(self.auth_key, message)
                sign_message(self.auth_key, fsm_update_message)
                self.server_connector.send_message(fsm_update_message, Scope.BROADCAST)
                self.send_message(message)
        except KeyError:
            LOG.error('IGNORING MESSAGE: Invalid message received: %s', message)

    def handle_trigger(self, message):
        """
        Attempts to step the FSM with the supplied trigger.

        :param message: TFW message with a data field containing
                        the action to try triggering in data['value']
        """
        trigger = message['data']['value']
        if self._require_signature:
            if not verify_message(self.auth_key, message):
                LOG.error('Ignoring unsigned trigger command: %s', message)
                return None
        if self.fsm.step(trigger):
            return message
        return None

    def handle_update(self, message):
        """
        Does nothing, but triggers an 'fsm_update' message.
        """
        # pylint: disable=no-self-use
        return message