From 4a810b4229259f43c43c856c862cde0bbb6446ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A1lint=20Bokros?= <balint.bokros@avatao.com>
Date: Mon, 27 Nov 2017 18:47:33 +0100
Subject: [PATCH] Create new FSM for SQL injection tutorial

---
 src/app/handlers/zmq_websocket_handler.py |  3 ++-
 src/app/sql_injection_fsm.py              | 32 +++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 src/app/sql_injection_fsm.py

diff --git a/src/app/handlers/zmq_websocket_handler.py b/src/app/handlers/zmq_websocket_handler.py
index eecbc01..9519cff 100644
--- a/src/app/handlers/zmq_websocket_handler.py
+++ b/src/app/handlers/zmq_websocket_handler.py
@@ -3,7 +3,8 @@ import logging
 from tornado.websocket import WebSocketHandler
 from util import parse_anchor_from_message
 
-from buttons import fsm
+# from buttons import fsm
+from sql_injection_fsm import fsm
 import component_connector
 
 
diff --git a/src/app/sql_injection_fsm.py b/src/app/sql_injection_fsm.py
new file mode 100644
index 0000000..196c12a
--- /dev/null
+++ b/src/app/sql_injection_fsm.py
@@ -0,0 +1,32 @@
+from transitions import Machine
+
+import component_connector
+
+
+class SQLInjectionFSM:
+    states = ['start', 'stripped_code', 'sql', 'commented_code', 'sql_with_substitutions', 'sql_output', 'end']
+    transitions = [
+        {'trigger': 'anchor_webide', 'source': 'start', 'dest': 'stripped_code'},
+        {'trigger': 'anchor_login', 'source': 'stripped_code', 'dest': 'sql'},
+        {'trigger': 'anchor_logger', 'source': 'sql', 'dest': 'commented_code'},
+        {'trigger': 'anchor_webide', 'source': 'commented_code', 'dest': 'sql_with_substitutions'},
+        {'trigger': 'anchor_logger', 'source': 'sql_with_substitutions', 'dest': 'sql_output'},
+        {'trigger': 'anchor_logger', 'source': 'sql_output', 'dest': 'end'},
+        {'trigger': 'reset', 'source': '*', 'dest': 'start'},
+    ]
+
+    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')
+        component_connector.send_message(message)
+
+
+fsm = SQLInjectionFSM()