Refactor FSM instance creation

This commit is contained in:
Kristóf Tóth 2018-01-29 14:09:40 +01:00
parent 22aef7dd7b
commit e05ae079ad
4 changed files with 28 additions and 14 deletions

View File

@ -7,11 +7,18 @@ from tornado.web import Application
from tornado.ioloop import IOLoop from tornado.ioloop import IOLoop
from config import WEB_PORT 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__': if __name__ == '__main__':
routes = [ routes = [
(r'/ws', ZMQWebSocketHandler), (r'/ws', SQLInjectionFSMManagingSocketHandler),
] ]
application = Application( application = Application(
routes, routes,

View File

@ -1 +1 @@
from .zmq_websocket_handler import ZMQWebSocketHandler from .zmq_websocket_handler import FSMManagingSocketHandler

View File

@ -3,8 +3,6 @@ import json
from tornado.websocket import WebSocketHandler from tornado.websocket import WebSocketHandler
from util import parse_anchor_from_message from util import parse_anchor_from_message
# from buttons import fsm
from sql_injection_fsm import fsm
import event_handler_connector import event_handler_connector
@ -23,12 +21,10 @@ class ZMQWebSocketHandler(WebSocketHandler):
def on_message(self, message): def on_message(self, message):
logging.debug('Received on WebSocket: {}'.format(message)) logging.debug('Received on WebSocket: {}'.format(message))
fsm.trigger(parse_anchor_from_message(message), message=message) self.send_message(*self.make_response(message))
anchor = 'FSMUpdate'
response = { 'anchor': anchor, def make_response(self, message):
'data': { 'current_state': fsm.state, raise NotImplementedError
'valid_transitions': [{'trigger': trigger} for trigger in fsm.machine.get_triggers()]}}
self.send_message(json.dumps(response), anchor)
def send_message(self, message: str, anchor: str = None): def send_message(self, message: str, anchor: str = None):
event_handler_connector.send_message(message, anchor) event_handler_connector.send_message(message, anchor)
@ -40,3 +36,17 @@ class ZMQWebSocketHandler(WebSocketHandler):
def check_origin(self, origin): def check_origin(self, origin):
return True 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

View File

@ -24,6 +24,3 @@ class SQLInjectionFSM(FSMBase):
def __init__(self): def __init__(self):
super().__init__('start') super().__init__('start')
fsm = SQLInjectionFSM()