You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
PurpleDome/tests/test_calderacontrol.py

455 lines
18 KiB
Python

import unittest
from unittest.mock import patch
from app.calderacontrol import CalderaControl
from simplejson.errors import JSONDecodeError
from app.exceptions import CalderaError
from app.attack_log import AttackLog
# https://docs.python.org/3/library/unittest.html
class TestExample(unittest.TestCase):
def setUp(self) -> None:
self.attack_logger = AttackLog(0)
self.cc = CalderaControl("https://localhost", attack_logger=self.attack_logger, apikey="123")
def tearDown(self) -> None:
pass
# List links sends the right commands and post
def test_list_links(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_links("asd")
mock_method.assert_called_once_with({"index": "link", "op_id": "asd"})
# List links gets an Exception and does not handle it (as expected)
def test_list_links_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_links("asd")
# list results sends the right commands and post
def test_list_results(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_results("asd")
mock_method.assert_called_once_with({"index": "result", "link_id": "asd"})
# List results gets an Exception and does not handle it (as expected)
def test_list_results_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_results("asd")
# list_operations
def test_list_operations(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_operations()
mock_method.assert_called_once_with({"index": "operations"})
# list operations gets the expected exception
def test_list_operations_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_operations()
# list_abilities
def test_list_abilities(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_abilities()
mock_method.assert_called_once_with({"index": "abilities"})
# list abilities gets the expected exception
def test_list_abilities_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_abilities()
# list_agents
def test_list_agents(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_agents()
mock_method.assert_called_once_with({"index": "agents"})
# list agents gets the expected exception
def test_list_agents_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_agents()
# list_adversaries
def test_list_adversaries(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_adversaries()
mock_method.assert_called_once_with({"index": "adversaries"})
# list adversaries gets the expected exception
def test_list_adversaries_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_adversaries()
# list_objectives
def test_list_objectives(self):
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.list_objectives()
mock_method.assert_called_once_with({"index": "objectives"})
# list objectives gets the expected exception
def test_list_objectives_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.list_objectives()
# Get operation (working)
def test_get_operation(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_operations", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_operation("bar")
self.assertEqual(op, a)
# Get operation (no matching name)
def test_get_operation_not_available(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_operations", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_operation("baaaaar")
self.assertEqual(op, None)
# get_adversary
def test_get_adversary(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_adversaries", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_adversary("bar")
self.assertEqual(op, a)
# get_adversary (no matching name)
def test_get_adversary_not_available(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_adversaries", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_adversary("baaaar")
self.assertEqual(op, None)
# get_objective
def test_get_objective(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_objectives", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_objective("bar")
self.assertEqual(op, a)
# get_objective (no matching name)
def test_get_objective_not_available(self):
a = {"name": "bar", "bar": "baz"}
with patch.object(self.cc, "list_objectives", return_value=[a, {"name": "zigg", "bar": "baz"}]):
op = self.cc.get_objective("baaaar")
self.assertEqual(op, None)
# get_ability
def test_get_ability(self):
a = {"ability_id": "bar", "bar": "baz"}
with patch.object(self.cc, "list_abilities", return_value=[a, {"ability_id": "zigg", "bar": "baz"}]):
op = self.cc.get_ability("bar")
self.assertEqual(op, [a])
# get_ability (no matching name)
def test_get_ability_not_available(self):
a = {"ability_id": "bar", "bar": "baz"}
with patch.object(self.cc, "list_abilities", return_value=[a, {"ability_id": "zigg", "bar": "baz"}]):
op = self.cc.get_ability("baaaar")
self.assertEqual(op, [])
# get_operation_by_id
def test_get_operation_by_id(self):
opid = "FooBar"
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.get_operation_by_id(opid)
mock_method.assert_called_once_with({"index": "operations", "id": opid})
# get_operation_by_id gets the expected exception
def test_get_operation_by_id_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.get_result_by_id("FooBar")
# get_result_by_id
def test_get_result_by_id(self):
opid = "FooBar"
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.get_result_by_id(opid)
mock_method.assert_called_once_with({"index": "result", "link_id": opid})
# get_result_by_id gets the expected exception
def test_get_result_by_id_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.get_result_by_id("FooBar")
# get_linkid
def test_get_linkid(self):
paw = "FooBar"
ability_id = "AID"
alink = {"paw": paw,
"ability": {"ability_id": ability_id},
"id": "Getme"}
op = [{"chain": [alink]}]
with patch.object(self.cc, "get_operation_by_id", return_value=op):
res = self.cc.get_linkid("Foo", paw, ability_id)
self.assertEqual(res, "Getme")
# get missing link id
def test_get_linkid_missing(self):
paw = "FooBar"
ability_id = "AID"
alink = {"paw": paw,
"ability": {"ability_id": ability_id},
"id": "Getme"}
op = [{"chain": [alink]}]
with patch.object(self.cc, "get_operation_by_id", return_value=op):
res = self.cc.get_linkid("Foo", "Bar", ability_id)
self.assertEqual(res, None)
# view_operation_report
def test_view_operation_report(self):
opid = "FooBar"
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.view_operation_report(opid)
mock_method.assert_called_once_with({"index": "operation_report", "op_id": opid, "agent_output": 1})
# get_result_by_id gets the expected exception
def test_view_operation_report_with_exception(self):
with self.assertRaises(JSONDecodeError):
with patch.object(self.cc, "__contact_server__", side_effect=JSONDecodeError("foo", "bar", 2)):
self.cc.view_operation_report("FooBar")
# view_operation_output
def test_view_operation_output(self):
opid = "OPID_123"
paw = "PAW_123"
ability_id = "AID_123"
or_ret = {"steps": {paw: {"steps": [{"ability_id": ability_id, "output": "red"}]}}}
with patch.object(self.cc, "view_operation_report", return_value=or_ret):
res = self.cc.view_operation_output(opid, paw, ability_id)
self.assertEqual(res, "red")
# view_operation_output paw missing
def test_view_operation_output_paw_missing(self):
opid = "OPID_123"
paw = "PAW_123"
ability_id = "AID_123"
or_ret = {"steps": {paw: {"steps": [{"ability_id": ability_id, "output": "red"}]}}}
with self.assertRaises(CalderaError):
with patch.object(self.cc, "view_operation_report", return_value=or_ret):
self.cc.view_operation_output(opid, "missing", ability_id)
# view_operation_output ability_id missing
def test_view_operation_output_ability_id_missing(self):
opid = "OPID_123"
paw = "PAW_123"
ability_id = "AID_123"
or_ret = {"steps": {paw: {"steps": [{"ability_id": ability_id, "output": "red"}]}}}
with patch.object(self.cc, "view_operation_report", return_value=or_ret):
res = self.cc.view_operation_output(opid, paw, "missing")
self.assertEqual(res, None)
# add_operation
def test_add_operation(self):
name = "test_name"
state = "test_state"
group = "test_group"
advid = "test_id"
exp = {"index": "operations",
"name": name,
"state": state,
"autonomous": 1,
'obfuscator': 'plain-text',
'auto_close': '1',
'jitter': '4/8',
'source': 'Alice Filters',
'visibility': '50',
"group": group,
"planner": "atomic",
"adversary_id": advid,
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_operation(name, advid, group, state)
mock_method.assert_called_once_with(exp, method="put")
# add_operation defaults
def test_add_operation_defaults(self):
name = "test_name"
advid = "test_id"
exp = {"index": "operations",
"name": name,
"state": "running", # default
"autonomous": 1,
'obfuscator': 'plain-text',
'auto_close': '1',
'jitter': '4/8',
'source': 'Alice Filters',
'visibility': '50',
"group": "red", # default
"planner": "atomic",
"adversary_id": advid,
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_operation(name, advid)
mock_method.assert_called_once_with(exp, method="put")
# add_adversary
def test_add_adversary(self):
name = "test_name"
ability = "test_ability"
description = "test_descritption"
exp = {"index": "adversaries",
"name": name,
"description": description,
"atomic_ordering": [{"id": ability}],
#
"objective": '495a9828-cab1-44dd-a0ca-66e58177d8cc' # default objective
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_adversary(name, ability, description)
mock_method.assert_called_once_with(exp, method="put")
def test_add_adversary_default(self):
name = "test_name"
ability = "test_ability"
exp = {"index": "adversaries",
"name": name,
"description": "created automatically",
"atomic_ordering": [{"id": ability}],
#
"objective": '495a9828-cab1-44dd-a0ca-66e58177d8cc' # default objective
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.add_adversary(name, ability)
mock_method.assert_called_once_with(exp, method="put")
# execute_ability
def test_execute_ability(self):
paw = "test_paw"
ability_id = "test_ability"
obfuscator = "plain-text"
exp = {"paw": paw,
"ability_id": ability_id,
"obfuscator": obfuscator}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.execute_ability(paw, ability_id, obfuscator)
mock_method.assert_called_once_with(exp, rest_path="plugin/access/exploit_ex")
def test_execute_ability_default(self):
paw = "test_paw"
ability_id = "test_ability"
exp = {"paw": paw,
"ability_id": ability_id,
"obfuscator": "plain-text"}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.execute_ability(paw, ability_id)
mock_method.assert_called_once_with(exp, rest_path="plugin/access/exploit_ex")
# execute_operation
def test_execute_operation(self):
operation_id = "test_opid"
state = "paused"
exp = {"index": "operation",
"op_id": operation_id,
"state": state}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.execute_operation(operation_id, state)
mock_method.assert_called_once_with(exp)
# not supported state
def test_execute_operation_not_supported(self):
operation_id = "test_opid"
state = "not supported"
with self.assertRaises(ValueError):
with patch.object(self.cc, "__contact_server__", return_value=None):
self.cc.execute_operation(operation_id, state)
def test_execute_operation_default(self):
operation_id = "test_opid"
exp = {"index": "operation",
"op_id": operation_id,
"state": "running" # default
}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.execute_operation(operation_id)
mock_method.assert_called_once_with(exp)
# delete_operation
def test_delete_operation(self):
opid = "test_opid"
exp = {"index": "operations",
"id": opid}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.delete_operation(opid)
mock_method.assert_called_once_with(exp, method="delete")
# delete_adversary
def test_delete_adversary(self):
adid = "test_adid"
exp = {"index": "adversaries",
"adversary_id": [{"adversary_id": adid}]}
with patch.object(self.cc, "__contact_server__", return_value=None) as mock_method:
self.cc.delete_adversary(adid)
mock_method.assert_called_once_with(exp, method="delete")
# is_operation_finished
def test_is_operation_finished_true(self):
opdata = [{"state": "finished", "chain": [{"status": 0}]}]
opid = "does not matter"
with patch.object(self.cc, "get_operation_by_id", return_value=opdata):
res = self.cc.is_operation_finished(opid)
self.assertEqual(res, True)
def test_is_operation_finished_false(self):
opdata = [{"state": "running", "chain": [{"status": 1}]}]
opid = "does not matter"
with patch.object(self.cc, "get_operation_by_id", return_value=opdata):
res = self.cc.is_operation_finished(opid)
self.assertEqual(res, False)
def test_is_operation_finished_exception(self):
opdata = [{"chain": [{"statusa": 1}]}]
opid = "does not matter"
with self.assertRaises(CalderaError):
with patch.object(self.cc, "get_operation_by_id", return_value=opdata):
self.cc.is_operation_finished(opid)
def test_is_operation_finished_exception2(self):
opdata = []
opid = "does not matter"
with self.assertRaises(CalderaError):
with patch.object(self.cc, "get_operation_by_id", return_value=opdata):
self.cc.is_operation_finished(opid)
# TODO attack (lots of complexity !)
# TODO test fetch_client
# TODO test __contact_server__
if __name__ == '__main__':
unittest.main()