diff --git a/src/app/app.py b/src/app/app.py index fd24948..8ee2815 100644 --- a/src/app/app.py +++ b/src/app/app.py @@ -7,11 +7,18 @@ from tornado.web import Application from tornado.ioloop import IOLoop from config import WEB_PORT -from handlers import ZMQWebSocketHandler +from handlers import FSMManagingSocketHandler +from sql_injection_fsm import SQLInjectionFSM + + +class SQLInjectionFSMManagingSocketHandler(FSMManagingSocketHandler): + def __init__(self, application, request, **kwargs): + super().__init__(SQLInjectionFSM, application, request, **kwargs) + if __name__ == '__main__': routes = [ - (r'/ws', ZMQWebSocketHandler), + (r'/ws', SQLInjectionFSMManagingSocketHandler), ] application = Application( routes, diff --git a/src/app/handlers/__init__.py b/src/app/handlers/__init__.py index c7ae86c..9ba4aa7 100644 --- a/src/app/handlers/__init__.py +++ b/src/app/handlers/__init__.py @@ -1 +1 @@ -from .zmq_websocket_handler import ZMQWebSocketHandler +from .zmq_websocket_handler import FSMManagingSocketHandler diff --git a/src/app/handlers/zmq_websocket_handler.py b/src/app/handlers/zmq_websocket_handler.py index 267434e..d2e4ce5 100644 --- a/src/app/handlers/zmq_websocket_handler.py +++ b/src/app/handlers/zmq_websocket_handler.py @@ -3,8 +3,6 @@ import json from tornado.websocket import WebSocketHandler from util import parse_anchor_from_message -# from buttons import fsm -from sql_injection_fsm import fsm import event_handler_connector @@ -23,12 +21,10 @@ class ZMQWebSocketHandler(WebSocketHandler): def on_message(self, message): logging.debug('Received on WebSocket: {}'.format(message)) - fsm.trigger(parse_anchor_from_message(message), message=message) - anchor = 'FSMUpdate' - response = { 'anchor': anchor, - 'data': { 'current_state': fsm.state, - 'valid_transitions': [{'trigger': trigger} for trigger in fsm.machine.get_triggers()]}} - self.send_message(json.dumps(response), anchor) + self.send_message(*self.make_response(message)) + + def make_response(self, message): + raise NotImplementedError def send_message(self, message: str, anchor: str = None): event_handler_connector.send_message(message, anchor) @@ -40,3 +36,17 @@ class ZMQWebSocketHandler(WebSocketHandler): def check_origin(self, origin): return True + +class FSMManagingSocketHandler(ZMQWebSocketHandler): + def __init__(self, FSM, application, request, **kwargs): + super().__init__(application, request, **kwargs) + self.fsm = FSM() + + def make_response(self, message): + self.fsm.trigger(parse_anchor_from_message(message), message=message) + anchor = 'FSMUpdate' + response = {'anchor': anchor, + 'data': {'current_state': self.fsm.state, + 'valid_transitions': + [{'trigger': trigger} for trigger in self.fsm.machine.get_triggers()]}} + return json.dumps(response), anchor diff --git a/src/app/sql_injection_fsm.py b/src/app/sql_injection_fsm.py index 54b428d..59696bf 100644 --- a/src/app/sql_injection_fsm.py +++ b/src/app/sql_injection_fsm.py @@ -24,6 +24,3 @@ class SQLInjectionFSM(FSMBase): def __init__(self): super().__init__('start') - - -fsm = SQLInjectionFSM()