diff --git a/lib/tfw/components/fsm_managing_event_handler.py b/lib/tfw/components/fsm_managing_event_handler.py index ee95c6f..47761a0 100644 --- a/lib/tfw/components/fsm_managing_event_handler.py +++ b/lib/tfw/components/fsm_managing_event_handler.py @@ -2,18 +2,19 @@ # All Rights Reserved. See LICENSE file for details. from tfw import EventHandlerBase -from tfw.crypto import KeyManager, sign_message +from tfw.crypto import KeyManager, sign_message, verify_message from tfw.config.logs import logging LOG = logging.getLogger(__name__) class FSMManagingEventHandler(EventHandlerBase): - def __init__(self, key, fsm_type): + def __init__(self, key, fsm_type, require_signature=False): super().__init__(key) 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, @@ -22,24 +23,29 @@ class FSMManagingEventHandler(EventHandlerBase): def handle_event(self, message): try: - data = message['data'] - message['data'] = self.command_handlers[data['command']](data) - fsm_update_message = self._fsm_updater.generate_fsm_update() - sign_message(self.auth_key, message) - sign_message(self.auth_key, fsm_update_message) - self.server_connector.broadcast(fsm_update_message) + message = self.command_handlers[message['data']['command']](message) + if message: + fsm_update_message = self._fsm_updater.generate_fsm_update() + sign_message(self.auth_key, message) + sign_message(self.auth_key, fsm_update_message) + self.server_connector.broadcast(fsm_update_message) return message except KeyError: LOG.error('IGNORING MESSAGE: Invalid message received: %s', message) - def handle_trigger(self, data): - trigger = data['value'] - self.fsm.step(trigger) - return data + def handle_trigger(self, message): + 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, data): + def handle_update(self, message): # pylint: disable=no-self-use - return data + return message class FSMUpdater: diff --git a/lib/tfw/fsm_base.py b/lib/tfw/fsm_base.py index 3cc2af5..7ae8157 100644 --- a/lib/tfw/fsm_base.py +++ b/lib/tfw/fsm_base.py @@ -62,5 +62,7 @@ class FSMBase(Machine, CallbackMixin): try: self.trigger(trigger) self.trigger_history.append(trigger) + return True except (AttributeError, MachineError): LOG.debug('FSM failed to execute nonexistent trigger: "%s"', trigger) + return False diff --git a/lib/tfw/networking/server/tfw_server.py b/lib/tfw/networking/server/tfw_server.py index 5f297df..c0d98f5 100644 --- a/lib/tfw/networking/server/tfw_server.py +++ b/lib/tfw/networking/server/tfw_server.py @@ -9,6 +9,7 @@ from tornado.web import Application from tfw.networking.event_handlers import ServerUplinkConnector from tfw.networking.server import EventHandlerConnector from tfw.networking import MessageSender +from tfw.crypto import KeyManager, verify_message, sign_message from tfw.config.logs import logging from .zmq_websocket_proxy import ZMQWebSocketProxy @@ -24,6 +25,7 @@ class TFWServer: def __init__(self): self._event_handler_connector = EventHandlerConnector() self._uplink_connector = ServerUplinkConnector() + self._auth_key = KeyManager().auth_key self.application = Application([( r'/ws', ZMQWebSocketProxy, { @@ -37,13 +39,16 @@ class TFWServer: def handle_trigger(self, message): if 'trigger' in message: LOG.debug('Executing handler for trigger "%s"', message.get('trigger', '')) - self._uplink_connector.send_to_eventhandler({ + fsm_eh_command = { 'key': 'fsm', 'data': { 'command': 'trigger', - 'value': message.get('trigger', '') + 'value': message['trigger'] } - }) + } + if verify_message(self._auth_key, message): + sign_message(self._auth_key, fsm_eh_command) + self._uplink_connector.send_to_eventhandler(fsm_eh_command) def handle_recover(self, message): if message['key'] == 'recover':