From e05ae079ade46304d0b326faad27aecbda4376f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= <kristof.toth@avatao.com>
Date: Mon, 29 Jan 2018 14:09:40 +0100
Subject: [PATCH] Refactor FSM instance creation

---
 src/app/app.py                            | 11 ++++++++--
 src/app/handlers/__init__.py              |  2 +-
 src/app/handlers/zmq_websocket_handler.py | 26 ++++++++++++++++-------
 src/app/sql_injection_fsm.py              |  3 ---
 4 files changed, 28 insertions(+), 14 deletions(-)

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()