def __init__(self):
self.playbook1 = Playbook(
'play1', workflows=[Workflow('work1'),
Workflow('work2')])
self.playbook2 = Playbook(
'play2', workflows=[Workflow('work1'),
Workflow('work3')])
def test_load_playbook(self):
workflows = [Workflow('something'), Workflow('something2')]
playbook = Playbook('test', workflows=workflows)
filepath = os.path.join(test_data_path, 'test.json')
with open(filepath, 'w') as file_out:
file_out.write(json.dumps(playbook.read()))
loaded = Loader.load_playbook(filepath)
self.assertIsInstance(loaded, Playbook)
self.assertEqual(loaded.name, 'test')
def test_get_branch_invalid_action(self):
flag = Condition('HelloWorld',
action_name='regMatch',
arguments=[Argument('regex', value='aaa')])
branch = Branch(source_uid="1",
destination_uid='next',
conditions=[flag])
action = Action('HelloWorld', 'helloWorld', uid="2")
action._output = ActionResult(result='bbb', status='Success')
workflow = Workflow(actions=[action], branches=[branch])
self.assertIsNone(workflow.get_branch(action, {}))
def test_load_playbooks_some_invalid_playbooks_in_directory(self):
workflows = [Workflow('something'), Workflow('something2')]
playbook = Playbook('test', workflows=workflows)
filepath = os.path.join(test_data_path, 'test.playbook')
with open(filepath, 'w') as file_out:
file_out.write(json.dumps(playbook.read()))
test_invalid_json = 'something not json'
filepath = os.path.join(test_data_path, 'test2.playbook')
with open(filepath, 'w') as file_out:
file_out.write(test_invalid_json)
loaded = Loader.load_playbooks(test_data_path)
self.assertEqual(len(loaded), 1)
self.assertEqual(loaded[0].name, 'test')
def test_load_multiple_workflows(self):
workflows = [Workflow('something'), Workflow('something2')]
playbook = Playbook('test', workflows=workflows)
filepath = os.path.join(test_data_path, 'test.playbook')
with open(filepath, 'w') as file_out:
file_out.write(json.dumps(playbook.read()))
workflows = [Workflow('something'), Workflow('something2')]
playbook = Playbook('test2', workflows=workflows)
filepath = os.path.join(test_data_path, 'test2.playbook')
with open(filepath, 'w') as file_out:
file_out.write(json.dumps(playbook.read()))
loaded = Loader.load_playbooks(test_data_path)
self.assertEqual(len(loaded), 2)
self.assertSetEqual({playbook.name
for playbook in loaded}, {'test', 'test2'})
def recreate_workflow(workflow_json):
"""Recreates a workflow from a JSON to prepare for it to be executed.
Args:
workflow_json (JSON dict): The input workflow JSON, with some other fields as well.
Returns:
(Workflow object, start_input): A tuple containing the reconstructed Workflow object, and the input to
the start step.
"""
uid = workflow_json['uid']
del workflow_json['uid']
execution_uid = workflow_json['execution_uid']
del workflow_json['execution_uid']
start = workflow_json['start']
start_input = ''
if 'start_input' in workflow_json:
start_input = workflow_json['start_input']
del workflow_json['start_input']
workflow = Workflow.create(workflow_json)
workflow.uid = uid
workflow.set_execution_uid(execution_uid)
workflow.start = start
return workflow, start_input
def test_create_playbook_nonempty_store_with_workflows(self):
self.store.playbooks['play1'] = self.loader.playbook1
workflow = Workflow('work1')
self.store.create_playbook('play2', [workflow])
self.assertStoreKeysEqual({'play1', 'play2'})
self.assertListEqual(
list(self.store.playbooks['play2'].workflows.values()), [workflow])
def recreate_workflow(workflow_json):
"""Recreates a workflow from a JSON to prepare for it to be executed.
Args:
workflow_json (JSON dict): The input workflow JSON, with some other fields as well.
Returns:
(Workflow object, start_arguments): A tuple containing the reconstructed Workflow object, and the arguments to
the start action.
"""
uid = workflow_json['uid']
del workflow_json['uid']
execution_uid = workflow_json['execution_uid']
del workflow_json['execution_uid']
start = workflow_json['start']
start_arguments = {}
if 'start_arguments' in workflow_json:
start_arguments = [
Argument(**arg) for arg in workflow_json['start_arguments']
]
workflow_json.pop("start_arguments")
workflow = Workflow.create(workflow_json)
workflow.uid = uid
workflow.set_execution_uid(execution_uid)
workflow.start = start
return workflow, start_arguments
def test_branch_with_priority(self):
flag = Condition('HelloWorld',
action_name='regMatch',
arguments=[Argument('regex', value='aaa')])
branch_one = Branch(source_uid="1",
destination_uid='five',
conditions=[flag],
priority="5")
branch_two = Branch(source_uid="1",
destination_uid='one',
conditions=[flag],
priority="1")
action = Action('HelloWorld', 'helloWorld', uid="1")
action._output = ActionResult(result='aaa', status='Success')
workflow = Workflow(actions=[action],
branches=[branch_one, branch_two])
self.assertEqual(workflow.get_branch(action, {}), "one")
def test_add_workflow_playbook_not_in_store(self):
self.store.playbooks['play1'] = self.loader.playbook1
workflow = Workflow('work3')
self.store.add_workflow('play2', workflow)
self.assertSetEqual(set(self.store.playbooks.keys()),
{'play1', 'play2'})
self.assertListEqual(
list(self.store.playbooks['play2'].workflows.values()), [workflow])
def create_workflow(self, playbook_name, workflow_name):
"""
Creates an empty workflow
Args:
playbook_name (str): Name of the playbook to add
workflow_name (str): The name of the workflow to add
"""
workflow = Workflow(name=workflow_name)
self.add_workflow(playbook_name, workflow)
def test_get_branch(self):
flag = Condition('HelloWorld',
action_name='regMatch',
arguments=[Argument('regex', value='aaa')])
branch = Branch(source_uid="1", destination_uid="2", conditions=[flag])
action = Action('HelloWorld', 'helloWorld', uid="1")
action._output = ActionResult(result='aaa', status='Success')
workflow = Workflow(actions=[action], branches=[branch])
result = {'triggered': False}
def validate_sent_data(sender, **kwargs):
if isinstance(sender, Branch):
self.assertIn('event', kwargs)
self.assertEqual(kwargs['event'], WalkoffEvent.BranchTaken)
result['triggered'] = True
WalkoffEvent.CommonWorkflowSignal.connect(validate_sent_data)
self.assertEqual(workflow.get_branch(action, {}), '2')
self.assertTrue(result['triggered'])
def test_simple_risk(self):
workflow = Workflow(name='workflow')
workflow.create_step(name="stepOne",
action='helloWorld',
app='HelloWorld',
risk=1)
workflow.create_step(name="stepTwo",
action='helloWorld',
app='HelloWorld',
risk=2)
workflow.create_step(name="stepThree",
action='helloWorld',
app='HelloWorld',
risk=3)
self.assertEqual(workflow._total_risk, 6)
def load_workflow(resource, workflow_name):
"""Loads a workflow from a file.
Args:
resource (str): Path to the workflow.
workflow_name (str): Name of the workflow to load.
Returns:
True on success, False otherwise.
"""
try:
playbook_file = open(resource, 'r')
except (IOError, OSError) as e:
logger.error(
'Could not load workflow from {0}. Reason: {1}'.format(
resource, format_exception_message(e)))
return None
else:
with playbook_file:
workflow_loaded = playbook_file.read()
try:
playbook_json = json.loads(workflow_loaded)
playbook_name = playbook_json['name']
workflow_json = next(
(workflow for workflow in playbook_json['workflows']
if workflow['name'] == workflow_name), None)
if workflow_json is None:
logger.warning(
'Workflow {0} not found in playbook {0}. '
'Cannot load.'.format(workflow_name,
playbook_name))
return None
workflow = Workflow.create(workflow_json)
return playbook_name, workflow
except ValueError as e:
logger.error('Cannot parse {0}. Reason: {1}'.format(
resource, format_exception_message(e)))
except (InvalidInput, UnknownApp, UnknownAppAction,
UnknownFilter, UnknownFlag) as e:
logger.error(
'Error constructing workflow {0}. Reason: {1}'.format(
workflow_name, format_exception_message(e)))
return None
except KeyError as e:
logger.error(
'Invalid Playbook JSON format. Details: {}'.format(e))
return None
def test_create_playbook_playbook_already_exists(self):
self.store.playbooks['play1'] = self.loader.playbook1
workflow = Workflow('work3')
self.store.create_playbook('play1', workflows=[workflow])
self.assertStoreKeysEqual({'play1'})
self.assertPlaybookWorkflowNamesEqual('play1', {'work3'})
def test_add_workflow_playbook_in_store(self):
self.store.playbooks['play1'] = self.loader.playbook1
workflow = Workflow('work3')
self.store.add_workflow('play1', workflow)
self.assertStoreKeysEqual({'play1'})
self.assertPlaybookHasWorkflowName('play1', 'work3')
def test_add_workflow_no_playbooks(self):
workflow = Workflow('work1')
self.store.add_workflow('play1', workflow)
self.assertStoreKeysEqual({'play1'})
self.assertListEqual(list(self.store.playbooks['play1'].workflows.values()), [workflow])
def test_accumulated_risk_with_error(self):
workflow = Workflow(name='workflow')
workflow._execution_uid = 'some_uid'
step1 = Step(name="step_one", app='HelloWorld', action='Buggy', risk=1)
step2 = Step(name="step_two", app='HelloWorld', action='Buggy', risk=2)
step3 = Step(name="step_three",
app='HelloWorld',
action='Buggy',
risk=3.5)
workflow.steps = {
'step_one': step1,
'step_two': step2,
'step_three': step3
}
workflow._total_risk = 6.5
instance = AppInstance.create(app_name='HelloWorld',
device_name='test_device_name')
workflow._Workflow__execute_step(workflow.steps["step_one"], instance)
self.assertAlmostEqual(workflow.accumulated_risk, 1.0 / 6.5)
workflow._Workflow__execute_step(workflow.steps["step_two"], instance)
self.assertAlmostEqual(workflow.accumulated_risk,
(1.0 / 6.5) + (2.0 / 6.5))
workflow._Workflow__execute_step(workflow.steps["step_three"],
instance)
self.assertAlmostEqual(workflow.accumulated_risk, 1.0)
def test_get_branch_no_branchs(self):
workflow = Workflow()
self.assertIsNone(workflow.get_branch(None, {}))
