import logging

import taskflow.engines
from taskflow.patterns import graph_flow as gf
from taskflow import task
from taskflow.listeners import logging as logging_listener

logging.basicConfig(level=logging.DEBUG)


# Definition of attack module execution
class Attack(task.Task):
    def __init__(self, name, arg, pro=None, req=None, *args, **kwargs):
        super(Attack, self).__init__(name=name, **kwargs)
            
        self.name = name
        self.arg = arg
        self.provides = pro
        self.requires = req

    def execute(self, *args, **kwargs):
        print("Executing attack with arguments: %s." % self.arg)
        return {'return_code': self.arg}


# Creating attack Steps
step1 = Attack(name='t1', arg=0)
step2 = Attack(name='t2', arg=-1)
step3 = Attack(name='t3', arg=0)

# Creating flow (Stage)
flow = gf.Flow('stage1').add(
    step1,
    step2,
    step3
)


def decide(history):
    print(history)
    for k in history:
        if history.get(k).get('return_code') != 0:
            return False
    return True 


# Set dependencies between steps
flow.link(step1, step3, decider=decide)
flow.link(step1, step2, decider=decide)
flow.link(step2, step3, decider=decide)

# Load the flow
e = taskflow.engines.load(flow)

# Run the flow with logging on
with logging_listener.DynamicLoggingListener(e):
    e.run()