eip-statemachine-test.py

#!/usr/bin/env python
# encoding: utf-8
"""
A Finite State Machine for the EIP service, based on
https://leap.se/code/issues/5616#note-2
"""
import random
from pprint import pprint

# see: http://xworkflows.readthedocs.org/en/latest/reference.html
import xworkflows


class EIPWorkflow(xworkflows.Workflow):
    """
    EIP workflow definition of states and transitions.
    """

    states = (
        ('off', "Everything is off"),
        ('on', "Everything is on"),
        ('error', "Fatal error, can't recover."),

        ('fw_start', "Firewall is starting."),
        ('fw_stop', "firewall is stopping."),

        ('eip_start', "Openvpn is starting."),
        ('eip_stop', "Openvpn is stopping."),
        ('eip_retry', "Retrying openvpn start."),
        ('eip_fail', "Gave up on openvpn, may be tried again."),
    )

    transitions = (
        ('start', 'off', 'fw_start'),

        ('fw_ok', 'fw_start', 'eip_start'),
        ('fw_error', 'fw_start', 'error'),

        ('eip_ok', 'eip_start', 'on'),
        ('eip_cancel', 'eip_start', 'eip_stop'),
        ('eip_error', 'eip_start', 'eip_retry'),
        ('eip_failed', 'eip_retry', 'error'),

        ('stop', 'on', 'eip_stop'),

        ('eip_stop_ok', 'eip_stop', 'fw_stop'),
        ('eip_stop_error', 'eip_stop', 'error'),

        ('fw_stop_ok', 'fw_stop', 'off'),
        ('fw_stop_error', 'fw_stop', 'error'),

        ('reset', 'error', 'off'),
    )

    initial_state = 'off'


class EIPMachine(xworkflows.WorkflowEnabled):
    """
    The actual state machine to be used to do the transitions and check states.
    """
    state = EIPWorkflow()

    @xworkflows.transition()
    def start(self):
        # print("EIPMachine: State is %s" % self.state.name)
        pass


class EIPManager(object):
    """
    The EIP Manager object, this handles the service actions and switch between
    states.
    """
    def __init__(self):
        self._statemachine = EIPMachine()

    def __getattribute__(self, name):
        mocked_methods = ['_eip_start', '_eip_stop',
                          '_firewall_start', '_firewall_stop']

        if name in mocked_methods:
            if random.randint(0, 5) == 0:
                raise Exception("{0}: some imaginary failure.".format(name))
            else:
                return lambda: {}
        else:  # the default:
            return object.__getattribute__(self, name)

    def status(self):
        return self._statemachine.state

    def start(self):
        self._statemachine.start()

        try:
            self._firewall_start()
        except Exception as e:
            print "Exception: {0!r}".format(e)
            self._statemachine.fw_error()
            return False
        else:
            self._statemachine.fw_ok()

        try:
            self._eip_start()
        except Exception as e:
            print "Exception: {0!r}".format(e)
            self._statemachine.eip_error()
            return False

        self._statemachine.eip_ok()
        return True

    def stop(self):
        self._statemachine.stop()

        try:
            self._eip_stop()
        except Exception as e:
            print "Exception: {0!r}".format(e)
            self._statemachine.eip_stop_error()
            return False
        else:
            self._statemachine.eip_stop_ok()

        try:
            self._firewall_stop()
        except Exception as e:
            print "Exception: {0!r}".format(e)
            self._statemachine.fw_stop_error()
            return False

        self._statemachine.fw_stop_ok()
        return True


def debug():
    """
    This shows us some details about the internals of the xworkflows object.
    """
    em = EIPMachine()

    print dir(em)
    print
    print dir(em.state)
    print
    pprint(list(em.state.transitions()))
    print


def simple_test():
    """A simple test of the transitions/states"""
    em = EIPMachine()

    print em.state.title
    em.start()
    print em.state.title
    em.fw_ok()
    print em.state.title
    em.eip_error()
    print em.state.title
    em.eip_failed()
    print em.state.title


def test_eip_manager():
    em = EIPManager()
    if em.start():
        print "EIP Started"
    else:
        print "Problem starting EIP:", em.status()


def test_random_paths():
    """
    This run a random sequence of valid transitions for each state until
    reaches a final state (error/on) or no more transitions can be done.
    """
    em = EIPMachine()

    # go randomly through available transitions
    retries = 0
    max_retries = 3
    while retries <= max_retries:
        try:
            # get available transitions from here
            ts = list(em.state.transitions())

            state = em.state.name
            print "State:", state

            # Stop in case of final state reached
            if state == 'error':
                retries += 1
                if retries <= max_retries:
                    print "Retrying!"

            if state == 'on':
                print "-"*50
                print "Final state reached. ON"
                break

            if len(ts) == 0:
                print "-"*50
                print "ERROR: no available transitions. State:", state
                break

            # pick one transition randomly
            random.shuffle(ts)
            transition = ts[0].name

            # do the transition
            t = em.__getattribute__(transition)
            t()
        except Exception as e:
            print "Exception caught:", repr(e)

    if state == 'error':
        print "-"*50
        print "Final state reached. Error"


if __name__ == '__main__':
    # test_random_paths()
    test_eip_manager()