mirror of
				https://github.com/avatao-content/baseimage-tutorial-framework
				synced 2025-10-31 20:02:55 +00:00 
			
		
		
		
	Implement require_signature mode in FSMManagingEH
This commit is contained in:
		| @@ -2,18 +2,19 @@ | |||||||
| # All Rights Reserved. See LICENSE file for details. | # All Rights Reserved. See LICENSE file for details. | ||||||
|  |  | ||||||
| from tfw import EventHandlerBase | 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 | from tfw.config.logs import logging | ||||||
|  |  | ||||||
| LOG = logging.getLogger(__name__) | LOG = logging.getLogger(__name__) | ||||||
|  |  | ||||||
|  |  | ||||||
| class FSMManagingEventHandler(EventHandlerBase): | class FSMManagingEventHandler(EventHandlerBase): | ||||||
|     def __init__(self, key, fsm_type): |     def __init__(self, key, fsm_type, require_signature=False): | ||||||
|         super().__init__(key) |         super().__init__(key) | ||||||
|         self.fsm = fsm_type() |         self.fsm = fsm_type() | ||||||
|         self._fsm_updater = FSMUpdater(self.fsm) |         self._fsm_updater = FSMUpdater(self.fsm) | ||||||
|         self.auth_key = KeyManager().auth_key |         self.auth_key = KeyManager().auth_key | ||||||
|  |         self._require_signature = require_signature | ||||||
|  |  | ||||||
|         self.command_handlers = { |         self.command_handlers = { | ||||||
|             'trigger': self.handle_trigger, |             'trigger': self.handle_trigger, | ||||||
| @@ -22,8 +23,8 @@ class FSMManagingEventHandler(EventHandlerBase): | |||||||
|  |  | ||||||
|     def handle_event(self, message): |     def handle_event(self, message): | ||||||
|         try: |         try: | ||||||
|             data = message['data'] |             message = self.command_handlers[message['data']['command']](message) | ||||||
|             message['data'] = self.command_handlers[data['command']](data) |             if message: | ||||||
|                 fsm_update_message = self._fsm_updater.generate_fsm_update() |                 fsm_update_message = self._fsm_updater.generate_fsm_update() | ||||||
|                 sign_message(self.auth_key, message) |                 sign_message(self.auth_key, message) | ||||||
|                 sign_message(self.auth_key, fsm_update_message) |                 sign_message(self.auth_key, fsm_update_message) | ||||||
| @@ -32,14 +33,19 @@ class FSMManagingEventHandler(EventHandlerBase): | |||||||
|         except KeyError: |         except KeyError: | ||||||
|             LOG.error('IGNORING MESSAGE: Invalid message received: %s', message) |             LOG.error('IGNORING MESSAGE: Invalid message received: %s', message) | ||||||
|  |  | ||||||
|     def handle_trigger(self, data): |     def handle_trigger(self, message): | ||||||
|         trigger = data['value'] |         trigger = message['data']['value'] | ||||||
|         self.fsm.step(trigger) |         if self._require_signature: | ||||||
|         return data |             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 |         # pylint: disable=no-self-use | ||||||
|         return data |         return message | ||||||
|  |  | ||||||
|  |  | ||||||
| class FSMUpdater: | class FSMUpdater: | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user