From b3e8af20245057bded2cdff244941e893646185d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= Date: Mon, 23 Jul 2018 17:16:37 +0200 Subject: [PATCH] Implement require_signature mode in FSMManagingEH --- .../components/fsm_managing_event_handler.py | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) 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: