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