From d3db21f26e067cd565ca7e1964cc3e5009d4cf84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A1lint=20Bokros?= <balint.bokros@avatao.com>
Date: Wed, 6 Dec 2017 01:29:09 +0100
Subject: [PATCH] Move common parts of FSM to a baseclass

---
 src/app/buttons.py           | 19 ++++---------------
 src/app/fsm_base.py          | 20 ++++++++++++++++++++
 src/app/sql_injection_fsm.py | 18 +++---------------
 3 files changed, 27 insertions(+), 30 deletions(-)
 create mode 100644 src/app/fsm_base.py

diff --git a/src/app/buttons.py b/src/app/buttons.py
index 4bdffee..0aa6210 100644
--- a/src/app/buttons.py
+++ b/src/app/buttons.py
@@ -1,8 +1,7 @@
-from transitions import Machine
+from fsm_base import FSMBase
 
-import event_handler_connector
 
-class Buttons:
+class Buttons(FSMBase):
     states = ['A', 'B', 'C']
     transitions = [
         {'trigger': 'anchor_a', 'source': 'A', 'dest': 'B'},
@@ -11,17 +10,7 @@ class Buttons:
     ]
 
     def __init__(self):
-        self.machine = Machine(model=self,
-                               states=Buttons.states,
-                               transitions=Buttons.transitions,
-                               initial='A',
-                               send_event=True,
-                               ignore_invalid_triggers=True,
-                               after_state_change='forward_message')
-
-    def forward_message(self, event_data):
-        message = event_data.kwargs.get('message')
-        event_handler_connector.send_message(message)
+        super().__init__('A')
 
 
-fsm = Buttons()
\ No newline at end of file
+fsm = Buttons()
diff --git a/src/app/fsm_base.py b/src/app/fsm_base.py
new file mode 100644
index 0000000..f22b06d
--- /dev/null
+++ b/src/app/fsm_base.py
@@ -0,0 +1,20 @@
+from transitions import Machine
+
+import event_handler_connector
+
+
+class FSMBase:
+    states, transitions = [], []
+
+    def __init__(self, initial: str = None):
+        self.machine = Machine(model=self,
+                               states=self.states,
+                               transitions=self.transitions,
+                               initial=initial or self.states[0],
+                               send_event=True,
+                               ignore_invalid_triggers=True,
+                               after_state_change='forward_message')
+
+    def forward_message(self, event_data):
+        message = event_data.kwargs.get('message')
+        event_handler_connector.send_message(message)
diff --git a/src/app/sql_injection_fsm.py b/src/app/sql_injection_fsm.py
index 3df3d45..9c64796 100644
--- a/src/app/sql_injection_fsm.py
+++ b/src/app/sql_injection_fsm.py
@@ -1,9 +1,7 @@
-from transitions import Machine
-
-import event_handler_connector
+from fsm_base import FSMBase
 
 
-class SQLInjectionFSM:
+class SQLInjectionFSM(FSMBase):
     states = ['start', 'stripped_code', 'sql', 'commented_code', 'sql_with_substitutions', 'sql_output', 'end']
     transitions = [
         {'trigger': 'anchor_webide', 'source': 'start', 'dest': 'stripped_code'},
@@ -16,17 +14,7 @@ class SQLInjectionFSM:
     ]
 
     def __init__(self):
-        self.machine = Machine(model=self,
-                               states=SQLInjectionFSM.states,
-                               transitions=SQLInjectionFSM.transitions,
-                               initial='start',
-                               send_event=True,
-                               ignore_invalid_triggers=True,
-                               after_state_change='forward_message')
-
-    def forward_message(self, event_data):
-        message = event_data.kwargs.get('message')
-        event_handler_connector.send_message(message)
+        super().__init__('start')
 
 
 fsm = SQLInjectionFSM()