From b9bbfac30b29475f1cae054efb1b91a529f328e2 Mon Sep 17 00:00:00 2001 From: Shreya Date: Sat, 23 Jan 2021 00:06:24 +0530 Subject: [PATCH 01/13] Add/modify tests for attack telems --- .../attack/tests/test_attack_telem_classes.py | 63 ++++++ .../attack/tests/test_technique_telems.py | 184 ++++++++++++++++++ .../attack/victim_host_telem_test.py | 29 --- 3 files changed, 247 insertions(+), 29 deletions(-) create mode 100644 monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py create mode 100644 monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py delete mode 100644 monkey/infection_monkey/telemetry/attack/victim_host_telem_test.py diff --git a/monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py b/monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py new file mode 100644 index 000000000..ca850ef9e --- /dev/null +++ b/monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py @@ -0,0 +1,63 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.model import VictimHost +from infection_monkey.telemetry.attack.attack_telem import AttackTelem +from infection_monkey.telemetry.attack.usage_telem import UsageTelem +from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem + +MACHINE = VictimHost('127.0.0.1') +STATUS = ScanStatus.USED +TECHNIQUE = 'T9999' +USAGE = UsageEnum.SMB + + +@pytest.fixture +def attack_telem_test_instance(): + return AttackTelem(TECHNIQUE, STATUS) + + +def test_attack_telem_category(attack_telem_test_instance): + assert attack_telem_test_instance.telem_category == 'attack' + + +def test_attack_telem_get_data(attack_telem_test_instance): + actual_data = attack_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': TECHNIQUE} + assert actual_data == expected_data + + +@pytest.fixture +def usage_telem_test_instance(): + return UsageTelem(TECHNIQUE, STATUS, USAGE) + + +def test_usage_telem_category(usage_telem_test_instance): + assert usage_telem_test_instance.telem_category == 'attack' + + +def test_usage_telem_get_data(usage_telem_test_instance): + actual_data = usage_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': TECHNIQUE, + 'usage': USAGE.name} + assert actual_data == expected_data + + +@pytest.fixture +def victim_host_telem_test_instance(): + return VictimHostTelem(TECHNIQUE, STATUS, MACHINE) + + +def test_victim_host_telem_category(victim_host_telem_test_instance): + assert victim_host_telem_test_instance.telem_category == 'attack' + + +def test_victim_host_telem_get_data(victim_host_telem_test_instance): + actual_data = victim_host_telem_test_instance.get_data() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': TECHNIQUE} + assert actual_data == expected_data diff --git a/monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py b/monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py new file mode 100644 index 000000000..47fd71665 --- /dev/null +++ b/monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py @@ -0,0 +1,184 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.model import VictimHost +from infection_monkey.telemetry.attack.t1005_telem import T1005Telem +from infection_monkey.telemetry.attack.t1035_telem import T1035Telem +from infection_monkey.telemetry.attack.t1064_telem import T1064Telem +from infection_monkey.telemetry.attack.t1105_telem import T1105Telem +from infection_monkey.telemetry.attack.t1106_telem import T1106Telem +from infection_monkey.telemetry.attack.t1107_telem import T1107Telem +from infection_monkey.telemetry.attack.t1129_telem import T1129Telem +from infection_monkey.telemetry.attack.t1197_telem import T1197Telem +from infection_monkey.telemetry.attack.t1222_telem import T1222Telem + +GATHERED_DATA_TYPE = '[Type of data collected]' +INFO = '[Additional info]' +MACHINE = VictimHost('127.0.0.1') +STATUS = ScanStatus.USED +USAGE = UsageEnum.SMB +SRC_IP = '0.0.0.0' +DST_IP = '0.0.0.1' +FILENAME = 'virus.exe' +PATH = 'path/to/file.txt' +COMMAND = 'echo hi' + + +@pytest.fixture +def T1005_telem_test_instance(): + return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO) + + +def test_T1005_telem_category(T1005_telem_test_instance): + assert T1005_telem_test_instance.telem_category == 'attack' + + +def test_T1005_get_data(T1005_telem_test_instance): + actual_data = T1005_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1005', + 'gathered_data_type': GATHERED_DATA_TYPE, + 'info': INFO} + assert actual_data == expected_data + + +@pytest.fixture +def T1035_telem_test_instance(): + return T1035Telem(STATUS, USAGE) + + +def test_T1035_telem_category(T1035_telem_test_instance): + assert T1035_telem_test_instance.telem_category == 'attack' + + +def test_T1035_get_data(T1035_telem_test_instance): + actual_data = T1035_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1035', + 'usage': USAGE.name} + assert actual_data == expected_data + + +@pytest.fixture +def T1064_telem_test_instance(): + return T1064Telem(STATUS, USAGE) + + +def test_T1064_telem_category(T1064_telem_test_instance): + assert T1064_telem_test_instance.telem_category == 'attack' + + +def test_T1064_get_data(T1064_telem_test_instance): + actual_data = T1064_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1064', + 'usage': USAGE} + assert actual_data == expected_data + + +@pytest.fixture +def T1105_telem_test_instance(): + return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME) + + +def test_T1105_telem_category(T1105_telem_test_instance): + assert T1105_telem_test_instance.telem_category == 'attack' + + +def test_T1105_get_data(T1105_telem_test_instance): + actual_data = T1105_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1105', + 'filename': FILENAME, + 'src': SRC_IP, + 'dst': DST_IP} + assert actual_data == expected_data + + +@pytest.fixture +def T1106_telem_test_instance(): + return T1106Telem(STATUS, USAGE) + + +def test_T1106_telem_category(T1106_telem_test_instance): + assert T1106_telem_test_instance.telem_category == 'attack' + + +def test_T1106_get_data(T1106_telem_test_instance): + actual_data = T1106_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1106', + 'usage': USAGE.name} + assert actual_data == expected_data + + +@pytest.fixture +def T1107_telem_test_instance(): + return T1107Telem(STATUS, PATH) + + +def test_T1107_telem_category(T1107_telem_test_instance): + assert T1107_telem_test_instance.telem_category == 'attack' + + +def test_T1107_get_data(T1107_telem_test_instance): + actual_data = T1107_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1107', + 'path': PATH} + assert actual_data == expected_data + + +@pytest.fixture +def T1129_telem_test_instance(): + return T1129Telem(STATUS, USAGE) + + +def test_T1129_telem_category(T1129_telem_test_instance): + assert T1129_telem_test_instance.telem_category == 'attack' + + +def test_T1129_get_data(T1129_telem_test_instance): + actual_data = T1129_telem_test_instance.get_data() + expected_data = {'status': STATUS.value, + 'technique': 'T1129', + 'usage': USAGE.name} + assert actual_data == expected_data + + +@pytest.fixture +def T1197_telem_test_instance(): + return T1197Telem(STATUS, MACHINE, USAGE) + + +def test_T1197_telem_category(T1197_telem_test_instance): + assert T1197_telem_test_instance.telem_category == 'attack' + + +def test_T1197_get_data(T1197_telem_test_instance): + actual_data = T1197_telem_test_instance.get_data() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': 'T1197', + 'usage': USAGE} + assert actual_data == expected_data + + +@pytest.fixture +def T1222_telem_test_instance(): + return T1222Telem(STATUS, COMMAND, MACHINE) + + +def test_T1222_telem_category(T1222_telem_test_instance): + assert T1222_telem_test_instance.telem_category == 'attack' + + +def test_T1222_get_data(T1222_telem_test_instance): + actual_data = T1222_telem_test_instance.get_data() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': 'T1222', + 'command': COMMAND} + assert actual_data == expected_data diff --git a/monkey/infection_monkey/telemetry/attack/victim_host_telem_test.py b/monkey/infection_monkey/telemetry/attack/victim_host_telem_test.py deleted file mode 100644 index 2ccab7483..000000000 --- a/monkey/infection_monkey/telemetry/attack/victim_host_telem_test.py +++ /dev/null @@ -1,29 +0,0 @@ -from unittest import TestCase - -from common.utils.attack_utils import ScanStatus -from infection_monkey.model import VictimHost -from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem - - -class TestVictimHostTelem(TestCase): - def test_get_data(self): - machine = VictimHost('127.0.0.1') - status = ScanStatus.USED - technique = 'T1210' - - telem = VictimHostTelem(technique, status, machine) - - self.assertEqual(telem.telem_category, 'attack') - - expected_data = { - 'machine': { - 'domain_name': machine.domain_name, - 'ip_addr': machine.ip_addr - }, - 'status': status.value, - 'technique': technique - } - - actual_data = telem.get_data() - - self.assertEqual(actual_data, expected_data) From 6b0cc1e36856c9f99693d624a5412ee4ca96f0a1 Mon Sep 17 00:00:00 2001 From: Shreya Date: Wed, 17 Feb 2021 18:05:16 +0530 Subject: [PATCH 02/13] Add tests for other base telems + put all telem tests in one folder --- .../tests/test_attack_telem_classes.py | 1 + .../tests/test_base_telem_classes.py | 142 ++++++++++++++++++ .../tests/test_technique_telems.py | 1 + 3 files changed, 144 insertions(+) rename monkey/infection_monkey/telemetry/{attack => }/tests/test_attack_telem_classes.py (99%) create mode 100644 monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py rename monkey/infection_monkey/telemetry/{attack => }/tests/test_technique_telems.py (99%) diff --git a/monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py similarity index 99% rename from monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py rename to monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py index ca850ef9e..de77c6ec1 100644 --- a/monkey/infection_monkey/telemetry/attack/tests/test_attack_telem_classes.py +++ b/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py @@ -6,6 +6,7 @@ from infection_monkey.telemetry.attack.attack_telem import AttackTelem from infection_monkey.telemetry.attack.usage_telem import UsageTelem from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem + MACHINE = VictimHost('127.0.0.1') STATUS = ScanStatus.USED TECHNIQUE = 'T9999' diff --git a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py new file mode 100644 index 000000000..3d00175e0 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py @@ -0,0 +1,142 @@ +import pytest + +from infection_monkey.exploit.wmiexec import WmiExploiter +from infection_monkey.model.host import VictimHost +from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs +from infection_monkey.telemetry.exploit_telem import ExploitTelem +from infection_monkey.telemetry.post_breach_telem import PostBreachTelem +from infection_monkey.telemetry.scan_telem import ScanTelem +from infection_monkey.telemetry.state_telem import StateTelem +from infection_monkey.telemetry.system_info_telem import SystemInfoTelem +from infection_monkey.telemetry.trace_telem import TraceTelem +from infection_monkey.telemetry.tunnel_telem import TunnelTelem + + +DOMAIN_NAME = 'domain-name' +HOSTNAME = 'hostname' +IP = '0.0.0.0' +IS_DONE = True +MSG = 'message' +RESULT = False +SYSTEM_INFO = {} +VERSION = 'version' +HOST = VictimHost(IP, DOMAIN_NAME) +EXPLOITER = WmiExploiter(HOST) +PBA = ScheduleJobs() + + +@pytest.fixture +def exploit_telem_test_instance(): + return ExploitTelem(EXPLOITER, RESULT) + + +def test_exploit_telem_category(exploit_telem_test_instance): + assert exploit_telem_test_instance.telem_category == 'exploit' + + +def test_exploit_telem_get_data(exploit_telem_test_instance): + actual_data = exploit_telem_test_instance.get_data() + expected_data = {'result': RESULT, + 'machine': HOST.as_dict(), + 'exploiter': EXPLOITER.__class__.__name__, + 'info': EXPLOITER.exploit_info, + 'attempts': EXPLOITER.exploit_attempts} + assert actual_data == expected_data + + +@pytest.fixture +def post_breach_telem_test_instance(mocker): + mocker.patch('infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip', + return_value=(HOSTNAME, IP)) + return PostBreachTelem(PBA, RESULT) + + +def test_post_breach_telem_category(post_breach_telem_test_instance): + assert post_breach_telem_test_instance.telem_category == 'post_breach' + + +def test_post_breach_telem_get_data(post_breach_telem_test_instance): + actual_data = post_breach_telem_test_instance.get_data() + expected_data = {'command': PBA.command, + 'result': RESULT, + 'name': PBA.name, + 'hostname': HOSTNAME, + 'ip': IP} + assert actual_data == expected_data + + +@pytest.fixture +def scan_telem_test_instance(): + return ScanTelem(HOST) + + +def test_scan_telem_category(scan_telem_test_instance): + assert scan_telem_test_instance.telem_category == 'scan' + + +def test_scan_telem_get_data(scan_telem_test_instance): + actual_data = scan_telem_test_instance.get_data() + expected_data = {'machine': HOST.as_dict(), + 'service_count': len(HOST.services)} + assert actual_data == expected_data + + +@pytest.fixture +def state_telem_test_instance(): + return StateTelem(IS_DONE, VERSION) + + +def test_state_telem_category(state_telem_test_instance): + assert state_telem_test_instance.telem_category == 'state' + + +def test_state_telem_get_data(state_telem_test_instance): + actual_data = state_telem_test_instance.get_data() + expected_data = {'done': IS_DONE, + 'version': VERSION} + assert actual_data == expected_data + + +@pytest.fixture +def system_info_telem_test_instance(): + return SystemInfoTelem(SYSTEM_INFO) + + +def test_system_info_telem_category(system_info_telem_test_instance): + assert system_info_telem_test_instance.telem_category == 'system_info' + + +def test_system_info_telem_get_data(system_info_telem_test_instance): + actual_data = system_info_telem_test_instance.get_data() + expected_data = SYSTEM_INFO + assert actual_data == expected_data + + +@pytest.fixture +def trace_telem_test_instance(): + return TraceTelem(MSG) + + +def test_trace_telem_category(trace_telem_test_instance): + assert trace_telem_test_instance.telem_category == 'trace' + + +def test_trace_telem_get_data(trace_telem_test_instance): + actual_data = trace_telem_test_instance.get_data() + expected_data = {'msg': MSG} + assert actual_data == expected_data + + +@pytest.fixture +def tunnel_telem_test_instance(): + return TunnelTelem() + + +def test_tunnel_telem_category(tunnel_telem_test_instance): + assert tunnel_telem_test_instance.telem_category == 'tunnel' + + +def test_tunnel_telem_get_data(tunnel_telem_test_instance): + actual_data = tunnel_telem_test_instance.get_data() + expected_data = {'proxy': None} + assert actual_data == expected_data diff --git a/monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py b/monkey/infection_monkey/telemetry/tests/test_technique_telems.py similarity index 99% rename from monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py rename to monkey/infection_monkey/telemetry/tests/test_technique_telems.py index 47fd71665..907205a0f 100644 --- a/monkey/infection_monkey/telemetry/attack/tests/test_technique_telems.py +++ b/monkey/infection_monkey/telemetry/tests/test_technique_telems.py @@ -12,6 +12,7 @@ from infection_monkey.telemetry.attack.t1129_telem import T1129Telem from infection_monkey.telemetry.attack.t1197_telem import T1197Telem from infection_monkey.telemetry.attack.t1222_telem import T1222Telem + GATHERED_DATA_TYPE = '[Type of data collected]' INFO = '[Additional info]' MACHINE = VictimHost('127.0.0.1') From 42b7fa05d6662d7680dd1f18a2088b0419f546c4 Mon Sep 17 00:00:00 2001 From: Shreya Date: Wed, 17 Feb 2021 18:45:48 +0530 Subject: [PATCH 03/13] Add requirement --- monkey/infection_monkey/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/infection_monkey/requirements.txt b/monkey/infection_monkey/requirements.txt index 0a1dbd282..b81018d95 100644 --- a/monkey/infection_monkey/requirements.txt +++ b/monkey/infection_monkey/requirements.txt @@ -13,5 +13,6 @@ pyftpdlib==1.5.6 pymssql<3.0 pypykatz==0.3.12 pysmb==1.2.5 +pytest-mock==3.5.1 requests>=2.24 wmi==1.5.1 ; sys_platform == 'win32' From 7960529ee91d6883d3b6a285909b4c0e8ffb1fc6 Mon Sep 17 00:00:00 2001 From: Shreya Date: Thu, 18 Feb 2021 19:41:29 +0530 Subject: [PATCH 04/13] Add conftest.py --- .../infection_monkey/telemetry/tests/conftest.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 monkey/infection_monkey/telemetry/tests/conftest.py diff --git a/monkey/infection_monkey/telemetry/tests/conftest.py b/monkey/infection_monkey/telemetry/tests/conftest.py new file mode 100644 index 000000000..dab650174 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/conftest.py @@ -0,0 +1,15 @@ +import pytest + +from infection_monkey.control import ControlClient + + +@pytest.fixture +def spy_send_telemetry(monkeypatch): + def _spy_send_telemetry(telem_category, data): + _spy_send_telemetry.telem_category = telem_category + _spy_send_telemetry.data = data + + _spy_send_telemetry.telem_category = None + _spy_send_telemetry.data = None + monkeypatch.setattr(ControlClient, 'send_telemetry', _spy_send_telemetry) + return _spy_send_telemetry From c2ed31bde89640d4c920d4f9d47c1331029b423f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 18 Feb 2021 09:33:58 -0500 Subject: [PATCH 05/13] telemetry: test `send()` for telemetry classes in `telemetry/` --- .../tests/test_base_telem_classes.py | 73 +++++++------------ 1 file changed, 28 insertions(+), 45 deletions(-) diff --git a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py index 3d00175e0..1bc117b30 100644 --- a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py +++ b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py @@ -30,18 +30,15 @@ def exploit_telem_test_instance(): return ExploitTelem(EXPLOITER, RESULT) -def test_exploit_telem_category(exploit_telem_test_instance): - assert exploit_telem_test_instance.telem_category == 'exploit' - - -def test_exploit_telem_get_data(exploit_telem_test_instance): - actual_data = exploit_telem_test_instance.get_data() +def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): + exploit_telem_test_instance.send() expected_data = {'result': RESULT, 'machine': HOST.as_dict(), 'exploiter': EXPLOITER.__class__.__name__, 'info': EXPLOITER.exploit_info, 'attempts': EXPLOITER.exploit_attempts} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "exploit" @pytest.fixture @@ -55,14 +52,15 @@ def test_post_breach_telem_category(post_breach_telem_test_instance): assert post_breach_telem_test_instance.telem_category == 'post_breach' -def test_post_breach_telem_get_data(post_breach_telem_test_instance): - actual_data = post_breach_telem_test_instance.get_data() +def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry): + post_breach_telem_test_instance.send() expected_data = {'command': PBA.command, 'result': RESULT, 'name': PBA.name, 'hostname': HOSTNAME, 'ip': IP} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "post_breach" @pytest.fixture @@ -70,15 +68,12 @@ def scan_telem_test_instance(): return ScanTelem(HOST) -def test_scan_telem_category(scan_telem_test_instance): - assert scan_telem_test_instance.telem_category == 'scan' - - -def test_scan_telem_get_data(scan_telem_test_instance): - actual_data = scan_telem_test_instance.get_data() +def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry): + scan_telem_test_instance.send() expected_data = {'machine': HOST.as_dict(), 'service_count': len(HOST.services)} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "scan" @pytest.fixture @@ -86,15 +81,12 @@ def state_telem_test_instance(): return StateTelem(IS_DONE, VERSION) -def test_state_telem_category(state_telem_test_instance): - assert state_telem_test_instance.telem_category == 'state' - - -def test_state_telem_get_data(state_telem_test_instance): - actual_data = state_telem_test_instance.get_data() +def test_state_telem_send(state_telem_test_instance, spy_send_telemetry): + state_telem_test_instance.send() expected_data = {'done': IS_DONE, 'version': VERSION} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "state" @pytest.fixture @@ -102,14 +94,11 @@ def system_info_telem_test_instance(): return SystemInfoTelem(SYSTEM_INFO) -def test_system_info_telem_category(system_info_telem_test_instance): - assert system_info_telem_test_instance.telem_category == 'system_info' - - -def test_system_info_telem_get_data(system_info_telem_test_instance): - actual_data = system_info_telem_test_instance.get_data() +def test_system_info_telem_send(system_info_telem_test_instance, spy_send_telemetry): + system_info_telem_test_instance.send() expected_data = SYSTEM_INFO - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "system_info" @pytest.fixture @@ -117,14 +106,11 @@ def trace_telem_test_instance(): return TraceTelem(MSG) -def test_trace_telem_category(trace_telem_test_instance): - assert trace_telem_test_instance.telem_category == 'trace' - - -def test_trace_telem_get_data(trace_telem_test_instance): - actual_data = trace_telem_test_instance.get_data() +def test_trace_telem_send(trace_telem_test_instance, spy_send_telemetry): + trace_telem_test_instance.send() expected_data = {'msg': MSG} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "trace" @pytest.fixture @@ -132,11 +118,8 @@ def tunnel_telem_test_instance(): return TunnelTelem() -def test_tunnel_telem_category(tunnel_telem_test_instance): - assert tunnel_telem_test_instance.telem_category == 'tunnel' - - -def test_tunnel_telem_get_data(tunnel_telem_test_instance): - actual_data = tunnel_telem_test_instance.get_data() +def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry): + tunnel_telem_test_instance.send() expected_data = {'proxy': None} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "tunnel" From 0ac9ce949c8d9fce0fa91e285e135826c27fdca3 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 18 Feb 2021 09:38:37 -0500 Subject: [PATCH 06/13] agent: reformat test_base_telem_classes.py with black --- .../tests/test_base_telem_classes.py | 52 ++++++++++--------- 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py index 1bc117b30..bbff9641c 100644 --- a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py +++ b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py @@ -12,14 +12,14 @@ from infection_monkey.telemetry.trace_telem import TraceTelem from infection_monkey.telemetry.tunnel_telem import TunnelTelem -DOMAIN_NAME = 'domain-name' -HOSTNAME = 'hostname' -IP = '0.0.0.0' +DOMAIN_NAME = "domain-name" +HOSTNAME = "hostname" +IP = "0.0.0.0" IS_DONE = True -MSG = 'message' +MSG = "message" RESULT = False SYSTEM_INFO = {} -VERSION = 'version' +VERSION = "version" HOST = VictimHost(IP, DOMAIN_NAME) EXPLOITER = WmiExploiter(HOST) PBA = ScheduleJobs() @@ -32,33 +32,39 @@ def exploit_telem_test_instance(): def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): exploit_telem_test_instance.send() - expected_data = {'result': RESULT, - 'machine': HOST.as_dict(), - 'exploiter': EXPLOITER.__class__.__name__, - 'info': EXPLOITER.exploit_info, - 'attempts': EXPLOITER.exploit_attempts} + expected_data = { + "result": RESULT, + "machine": HOST.as_dict(), + "exploiter": EXPLOITER.__class__.__name__, + "info": EXPLOITER.exploit_info, + "attempts": EXPLOITER.exploit_attempts, + } assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "exploit" @pytest.fixture def post_breach_telem_test_instance(mocker): - mocker.patch('infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip', - return_value=(HOSTNAME, IP)) + mocker.patch( + "infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip", + return_value=(HOSTNAME, IP), + ) return PostBreachTelem(PBA, RESULT) def test_post_breach_telem_category(post_breach_telem_test_instance): - assert post_breach_telem_test_instance.telem_category == 'post_breach' + assert post_breach_telem_test_instance.telem_category == "post_breach" def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry): post_breach_telem_test_instance.send() - expected_data = {'command': PBA.command, - 'result': RESULT, - 'name': PBA.name, - 'hostname': HOSTNAME, - 'ip': IP} + expected_data = { + "command": PBA.command, + "result": RESULT, + "name": PBA.name, + "hostname": HOSTNAME, + "ip": IP, + } assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "post_breach" @@ -70,8 +76,7 @@ def scan_telem_test_instance(): def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry): scan_telem_test_instance.send() - expected_data = {'machine': HOST.as_dict(), - 'service_count': len(HOST.services)} + expected_data = {"machine": HOST.as_dict(), "service_count": len(HOST.services)} assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "scan" @@ -83,8 +88,7 @@ def state_telem_test_instance(): def test_state_telem_send(state_telem_test_instance, spy_send_telemetry): state_telem_test_instance.send() - expected_data = {'done': IS_DONE, - 'version': VERSION} + expected_data = {"done": IS_DONE, "version": VERSION} assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "state" @@ -108,7 +112,7 @@ def trace_telem_test_instance(): def test_trace_telem_send(trace_telem_test_instance, spy_send_telemetry): trace_telem_test_instance.send() - expected_data = {'msg': MSG} + expected_data = {"msg": MSG} assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "trace" @@ -120,6 +124,6 @@ def tunnel_telem_test_instance(): def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry): tunnel_telem_test_instance.send() - expected_data = {'proxy': None} + expected_data = {"proxy": None} assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "tunnel" From 86ffaf358fadd5e729d496d005dce818139d5e64 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 18 Feb 2021 09:53:55 -0500 Subject: [PATCH 07/13] agent: break test_base_telem_classes into discrete test files --- .../tests/test_base_telem_classes.py | 129 ------------------ .../telemetry/tests/test_exploit_telem.py | 31 +++++ .../telemetry/tests/test_post_breach_telem.py | 32 +++++ .../telemetry/tests/test_scan_telem.py | 20 +++ .../telemetry/tests/test_state_telem.py | 18 +++ .../telemetry/tests/test_system_info_telem.py | 17 +++ .../telemetry/tests/test_trace_telem.py | 17 +++ .../telemetry/tests/test_tunnel_telem.py | 15 ++ 8 files changed, 150 insertions(+), 129 deletions(-) delete mode 100644 monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_exploit_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_scan_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_state_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_system_info_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_trace_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py diff --git a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py deleted file mode 100644 index bbff9641c..000000000 --- a/monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py +++ /dev/null @@ -1,129 +0,0 @@ -import pytest - -from infection_monkey.exploit.wmiexec import WmiExploiter -from infection_monkey.model.host import VictimHost -from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs -from infection_monkey.telemetry.exploit_telem import ExploitTelem -from infection_monkey.telemetry.post_breach_telem import PostBreachTelem -from infection_monkey.telemetry.scan_telem import ScanTelem -from infection_monkey.telemetry.state_telem import StateTelem -from infection_monkey.telemetry.system_info_telem import SystemInfoTelem -from infection_monkey.telemetry.trace_telem import TraceTelem -from infection_monkey.telemetry.tunnel_telem import TunnelTelem - - -DOMAIN_NAME = "domain-name" -HOSTNAME = "hostname" -IP = "0.0.0.0" -IS_DONE = True -MSG = "message" -RESULT = False -SYSTEM_INFO = {} -VERSION = "version" -HOST = VictimHost(IP, DOMAIN_NAME) -EXPLOITER = WmiExploiter(HOST) -PBA = ScheduleJobs() - - -@pytest.fixture -def exploit_telem_test_instance(): - return ExploitTelem(EXPLOITER, RESULT) - - -def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): - exploit_telem_test_instance.send() - expected_data = { - "result": RESULT, - "machine": HOST.as_dict(), - "exploiter": EXPLOITER.__class__.__name__, - "info": EXPLOITER.exploit_info, - "attempts": EXPLOITER.exploit_attempts, - } - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "exploit" - - -@pytest.fixture -def post_breach_telem_test_instance(mocker): - mocker.patch( - "infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip", - return_value=(HOSTNAME, IP), - ) - return PostBreachTelem(PBA, RESULT) - - -def test_post_breach_telem_category(post_breach_telem_test_instance): - assert post_breach_telem_test_instance.telem_category == "post_breach" - - -def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry): - post_breach_telem_test_instance.send() - expected_data = { - "command": PBA.command, - "result": RESULT, - "name": PBA.name, - "hostname": HOSTNAME, - "ip": IP, - } - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "post_breach" - - -@pytest.fixture -def scan_telem_test_instance(): - return ScanTelem(HOST) - - -def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry): - scan_telem_test_instance.send() - expected_data = {"machine": HOST.as_dict(), "service_count": len(HOST.services)} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "scan" - - -@pytest.fixture -def state_telem_test_instance(): - return StateTelem(IS_DONE, VERSION) - - -def test_state_telem_send(state_telem_test_instance, spy_send_telemetry): - state_telem_test_instance.send() - expected_data = {"done": IS_DONE, "version": VERSION} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "state" - - -@pytest.fixture -def system_info_telem_test_instance(): - return SystemInfoTelem(SYSTEM_INFO) - - -def test_system_info_telem_send(system_info_telem_test_instance, spy_send_telemetry): - system_info_telem_test_instance.send() - expected_data = SYSTEM_INFO - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "system_info" - - -@pytest.fixture -def trace_telem_test_instance(): - return TraceTelem(MSG) - - -def test_trace_telem_send(trace_telem_test_instance, spy_send_telemetry): - trace_telem_test_instance.send() - expected_data = {"msg": MSG} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "trace" - - -@pytest.fixture -def tunnel_telem_test_instance(): - return TunnelTelem() - - -def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry): - tunnel_telem_test_instance.send() - expected_data = {"proxy": None} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == "tunnel" diff --git a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py new file mode 100644 index 000000000..1002a3cb3 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py @@ -0,0 +1,31 @@ +import pytest + +from infection_monkey.exploit.wmiexec import WmiExploiter +from infection_monkey.model.host import VictimHost +from infection_monkey.telemetry.exploit_telem import ExploitTelem + + +HOSTNAME = "hostname" +DOMAIN_NAME = "domain-name" +IP = "0.0.0.0" +HOST = VictimHost(IP, DOMAIN_NAME) +EXPLOITER = WmiExploiter(HOST) +RESULT = False + + +@pytest.fixture +def exploit_telem_test_instance(): + return ExploitTelem(EXPLOITER, RESULT) + + +def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): + exploit_telem_test_instance.send() + expected_data = { + "result": RESULT, + "machine": HOST.as_dict(), + "exploiter": EXPLOITER.__class__.__name__, + "info": EXPLOITER.exploit_info, + "attempts": EXPLOITER.exploit_attempts, + } + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "exploit" diff --git a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py new file mode 100644 index 000000000..e6cbc45b2 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py @@ -0,0 +1,32 @@ +import pytest + +from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs +from infection_monkey.telemetry.post_breach_telem import PostBreachTelem + + +HOSTNAME = "hostname" +IP = "0.0.0.0" +PBA = ScheduleJobs() +RESULT = False + + +@pytest.fixture +def post_breach_telem_test_instance(mocker): + mocker.patch( + "infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip", + return_value=(HOSTNAME, IP), + ) + return PostBreachTelem(PBA, RESULT) + + +def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry): + post_breach_telem_test_instance.send() + expected_data = { + "command": PBA.command, + "result": RESULT, + "name": PBA.name, + "hostname": HOSTNAME, + "ip": IP, + } + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "post_breach" diff --git a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py new file mode 100644 index 000000000..d75aecdb1 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py @@ -0,0 +1,20 @@ +import pytest + +from infection_monkey.telemetry.scan_telem import ScanTelem +from infection_monkey.model.host import VictimHost + +DOMAIN_NAME = "domain-name" +IP = "0.0.0.0" +HOST = VictimHost(IP, DOMAIN_NAME) + + +@pytest.fixture +def scan_telem_test_instance(): + return ScanTelem(HOST) + + +def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry): + scan_telem_test_instance.send() + expected_data = {"machine": HOST.as_dict(), "service_count": len(HOST.services)} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "scan" diff --git a/monkey/infection_monkey/telemetry/tests/test_state_telem.py b/monkey/infection_monkey/telemetry/tests/test_state_telem.py new file mode 100644 index 000000000..a8beaf5ad --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_state_telem.py @@ -0,0 +1,18 @@ +import pytest + +from infection_monkey.telemetry.state_telem import StateTelem + +IS_DONE = True +VERSION = "version" + + +@pytest.fixture +def state_telem_test_instance(): + return StateTelem(IS_DONE, VERSION) + + +def test_state_telem_send(state_telem_test_instance, spy_send_telemetry): + state_telem_test_instance.send() + expected_data = {"done": IS_DONE, "version": VERSION} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "state" diff --git a/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py b/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py new file mode 100644 index 000000000..11692d4bb --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py @@ -0,0 +1,17 @@ +import pytest + +from infection_monkey.telemetry.system_info_telem import SystemInfoTelem + +SYSTEM_INFO = {} + + +@pytest.fixture +def system_info_telem_test_instance(): + return SystemInfoTelem(SYSTEM_INFO) + + +def test_system_info_telem_send(system_info_telem_test_instance, spy_send_telemetry): + system_info_telem_test_instance.send() + expected_data = SYSTEM_INFO + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "system_info" diff --git a/monkey/infection_monkey/telemetry/tests/test_trace_telem.py b/monkey/infection_monkey/telemetry/tests/test_trace_telem.py new file mode 100644 index 000000000..e4bb06a6b --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_trace_telem.py @@ -0,0 +1,17 @@ +import pytest + +from infection_monkey.telemetry.trace_telem import TraceTelem + +MSG = "message" + + +@pytest.fixture +def trace_telem_test_instance(): + return TraceTelem(MSG) + + +def test_trace_telem_send(trace_telem_test_instance, spy_send_telemetry): + trace_telem_test_instance.send() + expected_data = {"msg": MSG} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "trace" diff --git a/monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py b/monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py new file mode 100644 index 000000000..81a32bb44 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py @@ -0,0 +1,15 @@ +import pytest + +from infection_monkey.telemetry.tunnel_telem import TunnelTelem + + +@pytest.fixture +def tunnel_telem_test_instance(): + return TunnelTelem() + + +def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry): + tunnel_telem_test_instance.send() + expected_data = {"proxy": None} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "tunnel" From 4efdeeacc39562ed6841ead0d3ba6618fca56aaa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 18 Feb 2021 09:59:52 -0500 Subject: [PATCH 08/13] agent: remove dependency on pytest-mock --- monkey/infection_monkey/requirements.txt | 1 - .../telemetry/tests/test_post_breach_telem.py | 7 ++----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/requirements.txt b/monkey/infection_monkey/requirements.txt index b81018d95..0a1dbd282 100644 --- a/monkey/infection_monkey/requirements.txt +++ b/monkey/infection_monkey/requirements.txt @@ -13,6 +13,5 @@ pyftpdlib==1.5.6 pymssql<3.0 pypykatz==0.3.12 pysmb==1.2.5 -pytest-mock==3.5.1 requests>=2.24 wmi==1.5.1 ; sys_platform == 'win32' diff --git a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py index e6cbc45b2..d38cfdbde 100644 --- a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py @@ -11,11 +11,8 @@ RESULT = False @pytest.fixture -def post_breach_telem_test_instance(mocker): - mocker.patch( - "infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip", - return_value=(HOSTNAME, IP), - ) +def post_breach_telem_test_instance(monkeypatch): + monkeypatch.setattr(PostBreachTelem, "_get_hostname_and_ip", lambda: (HOSTNAME, IP)) return PostBreachTelem(PBA, RESULT) From 08addff8c55e974b5bcb82c88c563a2790803fbb Mon Sep 17 00:00:00 2001 From: Shreya Date: Thu, 18 Feb 2021 20:13:33 +0530 Subject: [PATCH 09/13] Modify tests for attack telem classes and technique telems - test `send()` instead of `get_data()` using fixture `spy_send_telemetry` --- .../telemetry/attack/t1064_telem.py | 1 + .../telemetry/attack/t1197_telem.py | 1 + .../tests/test_attack_telem_classes.py | 33 ++--- .../telemetry/tests/test_technique_telems.py | 118 +++++++----------- 4 files changed, 60 insertions(+), 93 deletions(-) diff --git a/monkey/infection_monkey/telemetry/attack/t1064_telem.py b/monkey/infection_monkey/telemetry/attack/t1064_telem.py index efea27063..94be44a79 100644 --- a/monkey/infection_monkey/telemetry/attack/t1064_telem.py +++ b/monkey/infection_monkey/telemetry/attack/t1064_telem.py @@ -3,6 +3,7 @@ from infection_monkey.telemetry.attack.usage_telem import AttackTelem class T1064Telem(AttackTelem): def __init__(self, status, usage): + # TODO: rename parameter "usage" to avoid confusion with parameter "usage" in UsageTelem techniques """ T1064 telemetry. :param status: ScanStatus of technique diff --git a/monkey/infection_monkey/telemetry/attack/t1197_telem.py b/monkey/infection_monkey/telemetry/attack/t1197_telem.py index 387c3aa13..769f93823 100644 --- a/monkey/infection_monkey/telemetry/attack/t1197_telem.py +++ b/monkey/infection_monkey/telemetry/attack/t1197_telem.py @@ -5,6 +5,7 @@ __author__ = "itay.mizeretz" class T1197Telem(VictimHostTelem): def __init__(self, status, machine, usage): + # TODO: rename parameter "usage" to avoid confusion with parameter "usage" in UsageTelem techniques """ T1197 telemetry. :param status: ScanStatus of technique diff --git a/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py index de77c6ec1..13dc02322 100644 --- a/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py +++ b/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py @@ -18,15 +18,12 @@ def attack_telem_test_instance(): return AttackTelem(TECHNIQUE, STATUS) -def test_attack_telem_category(attack_telem_test_instance): - assert attack_telem_test_instance.telem_category == 'attack' - - -def test_attack_telem_get_data(attack_telem_test_instance): - actual_data = attack_telem_test_instance.get_data() +def test_attack_telem_send(attack_telem_test_instance, spy_send_telemetry): + attack_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': TECHNIQUE} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -34,16 +31,13 @@ def usage_telem_test_instance(): return UsageTelem(TECHNIQUE, STATUS, USAGE) -def test_usage_telem_category(usage_telem_test_instance): - assert usage_telem_test_instance.telem_category == 'attack' - - -def test_usage_telem_get_data(usage_telem_test_instance): - actual_data = usage_telem_test_instance.get_data() +def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry): + usage_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': TECHNIQUE, 'usage': USAGE.name} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -51,14 +45,11 @@ def victim_host_telem_test_instance(): return VictimHostTelem(TECHNIQUE, STATUS, MACHINE) -def test_victim_host_telem_category(victim_host_telem_test_instance): - assert victim_host_telem_test_instance.telem_category == 'attack' - - -def test_victim_host_telem_get_data(victim_host_telem_test_instance): - actual_data = victim_host_telem_test_instance.get_data() +def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry): + victim_host_telem_test_instance.send() expected_data = {'machine': {'domain_name': MACHINE.domain_name, 'ip_addr': MACHINE.ip_addr}, 'status': STATUS.value, 'technique': TECHNIQUE} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/test_technique_telems.py b/monkey/infection_monkey/telemetry/tests/test_technique_telems.py index 907205a0f..b2c73867d 100644 --- a/monkey/infection_monkey/telemetry/tests/test_technique_telems.py +++ b/monkey/infection_monkey/telemetry/tests/test_technique_telems.py @@ -13,16 +13,17 @@ from infection_monkey.telemetry.attack.t1197_telem import T1197Telem from infection_monkey.telemetry.attack.t1222_telem import T1222Telem +COMMAND = 'echo hi' +DST_IP = '0.0.0.1' +FILENAME = 'virus.exe' GATHERED_DATA_TYPE = '[Type of data collected]' INFO = '[Additional info]' MACHINE = VictimHost('127.0.0.1') +PATH = 'path/to/file.txt' +SRC_IP = '0.0.0.0' STATUS = ScanStatus.USED USAGE = UsageEnum.SMB -SRC_IP = '0.0.0.0' -DST_IP = '0.0.0.1' -FILENAME = 'virus.exe' -PATH = 'path/to/file.txt' -COMMAND = 'echo hi' +USAGE_STR = '[Usage info]' @pytest.fixture @@ -30,17 +31,14 @@ def T1005_telem_test_instance(): return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO) -def test_T1005_telem_category(T1005_telem_test_instance): - assert T1005_telem_test_instance.telem_category == 'attack' - - -def test_T1005_get_data(T1005_telem_test_instance): - actual_data = T1005_telem_test_instance.get_data() +def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry): + T1005_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1005', 'gathered_data_type': GATHERED_DATA_TYPE, 'info': INFO} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -48,33 +46,27 @@ def T1035_telem_test_instance(): return T1035Telem(STATUS, USAGE) -def test_T1035_telem_category(T1035_telem_test_instance): - assert T1035_telem_test_instance.telem_category == 'attack' - - -def test_T1035_get_data(T1035_telem_test_instance): - actual_data = T1035_telem_test_instance.get_data() +def test_T1035_send(T1035_telem_test_instance, spy_send_telemetry): + T1035_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1035', 'usage': USAGE.name} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture def T1064_telem_test_instance(): - return T1064Telem(STATUS, USAGE) + return T1064Telem(STATUS, USAGE_STR) -def test_T1064_telem_category(T1064_telem_test_instance): - assert T1064_telem_test_instance.telem_category == 'attack' - - -def test_T1064_get_data(T1064_telem_test_instance): - actual_data = T1064_telem_test_instance.get_data() +def test_T1064_send(T1064_telem_test_instance, spy_send_telemetry): + T1064_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1064', - 'usage': USAGE} - assert actual_data == expected_data + 'usage': USAGE_STR} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -82,18 +74,15 @@ def T1105_telem_test_instance(): return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME) -def test_T1105_telem_category(T1105_telem_test_instance): - assert T1105_telem_test_instance.telem_category == 'attack' - - -def test_T1105_get_data(T1105_telem_test_instance): - actual_data = T1105_telem_test_instance.get_data() +def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry): + T1105_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1105', 'filename': FILENAME, 'src': SRC_IP, 'dst': DST_IP} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -101,16 +90,13 @@ def T1106_telem_test_instance(): return T1106Telem(STATUS, USAGE) -def test_T1106_telem_category(T1106_telem_test_instance): - assert T1106_telem_test_instance.telem_category == 'attack' - - -def test_T1106_get_data(T1106_telem_test_instance): - actual_data = T1106_telem_test_instance.get_data() +def test_T1106_send(T1106_telem_test_instance, spy_send_telemetry): + T1106_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1106', 'usage': USAGE.name} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -118,16 +104,13 @@ def T1107_telem_test_instance(): return T1107Telem(STATUS, PATH) -def test_T1107_telem_category(T1107_telem_test_instance): - assert T1107_telem_test_instance.telem_category == 'attack' - - -def test_T1107_get_data(T1107_telem_test_instance): - actual_data = T1107_telem_test_instance.get_data() +def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry): + T1107_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1107', 'path': PATH} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -135,35 +118,29 @@ def T1129_telem_test_instance(): return T1129Telem(STATUS, USAGE) -def test_T1129_telem_category(T1129_telem_test_instance): - assert T1129_telem_test_instance.telem_category == 'attack' - - -def test_T1129_get_data(T1129_telem_test_instance): - actual_data = T1129_telem_test_instance.get_data() +def test_T1129_send(T1129_telem_test_instance, spy_send_telemetry): + T1129_telem_test_instance.send() expected_data = {'status': STATUS.value, 'technique': 'T1129', 'usage': USAGE.name} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture def T1197_telem_test_instance(): - return T1197Telem(STATUS, MACHINE, USAGE) + return T1197Telem(STATUS, MACHINE, USAGE_STR) -def test_T1197_telem_category(T1197_telem_test_instance): - assert T1197_telem_test_instance.telem_category == 'attack' - - -def test_T1197_get_data(T1197_telem_test_instance): - actual_data = T1197_telem_test_instance.get_data() +def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry): + T1197_telem_test_instance.send() expected_data = {'machine': {'domain_name': MACHINE.domain_name, 'ip_addr': MACHINE.ip_addr}, 'status': STATUS.value, 'technique': 'T1197', - 'usage': USAGE} - assert actual_data == expected_data + 'usage': USAGE_STR} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' @pytest.fixture @@ -171,15 +148,12 @@ def T1222_telem_test_instance(): return T1222Telem(STATUS, COMMAND, MACHINE) -def test_T1222_telem_category(T1222_telem_test_instance): - assert T1222_telem_test_instance.telem_category == 'attack' - - -def test_T1222_get_data(T1222_telem_test_instance): - actual_data = T1222_telem_test_instance.get_data() +def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry): + T1222_telem_test_instance.send() expected_data = {'machine': {'domain_name': MACHINE.domain_name, 'ip_addr': MACHINE.ip_addr}, 'status': STATUS.value, 'technique': 'T1222', 'command': COMMAND} - assert actual_data == expected_data + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' From a4603853a92d7c6848893e674b1fb6c040405faf Mon Sep 17 00:00:00 2001 From: Shreya Date: Thu, 18 Feb 2021 22:35:26 +0530 Subject: [PATCH 10/13] Split test_attack_telem_classes.py and test_technique_telems.py into separate test files --- .../tests/attack/test_attack_telem.py | 21 +++ .../tests/attack/test_t1005_telem.py | 24 +++ .../tests/attack/test_t1035_telem.py | 22 +++ .../tests/attack/test_t1064_telem.py | 22 +++ .../tests/attack/test_t1105_telem.py | 26 +++ .../tests/attack/test_t1106_telem.py | 22 +++ .../tests/attack/test_t1107_telem.py | 22 +++ .../tests/attack/test_t1129_telem.py | 22 +++ .../tests/attack/test_t1197_telem.py | 26 +++ .../tests/attack/test_t1222_telem.py | 26 +++ .../tests/attack/test_usage_telem.py | 23 +++ .../tests/attack/test_victim_host_telem.py | 25 +++ .../tests/test_attack_telem_classes.py | 55 ------ .../telemetry/tests/test_technique_telems.py | 159 ------------------ 14 files changed, 281 insertions(+), 214 deletions(-) create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py create mode 100644 monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py delete mode 100644 monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py delete mode 100644 monkey/infection_monkey/telemetry/tests/test_technique_telems.py diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py new file mode 100644 index 000000000..750075fb2 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py @@ -0,0 +1,21 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.telemetry.attack.attack_telem import AttackTelem + + +STATUS = ScanStatus.USED +TECHNIQUE = 'T9999' + + +@pytest.fixture +def attack_telem_test_instance(): + return AttackTelem(TECHNIQUE, STATUS) + + +def test_attack_telem_send(attack_telem_test_instance, spy_send_telemetry): + attack_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': TECHNIQUE} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py new file mode 100644 index 000000000..757f0de13 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py @@ -0,0 +1,24 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.telemetry.attack.t1005_telem import T1005Telem + + +GATHERED_DATA_TYPE = '[Type of data collected]' +INFO = '[Additional info]' +STATUS = ScanStatus.USED + + +@pytest.fixture +def T1005_telem_test_instance(): + return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO) + + +def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry): + T1005_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1005', + 'gathered_data_type': GATHERED_DATA_TYPE, + 'info': INFO} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py new file mode 100644 index 000000000..a3133fcb9 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py @@ -0,0 +1,22 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.telemetry.attack.t1035_telem import T1035Telem + + +STATUS = ScanStatus.USED +USAGE = UsageEnum.SMB + + +@pytest.fixture +def T1035_telem_test_instance(): + return T1035Telem(STATUS, USAGE) + + +def test_T1035_send(T1035_telem_test_instance, spy_send_telemetry): + T1035_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1035', + 'usage': USAGE.name} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py new file mode 100644 index 000000000..575b57540 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py @@ -0,0 +1,22 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.telemetry.attack.t1064_telem import T1064Telem + + +STATUS = ScanStatus.USED +USAGE_STR = '[Usage info]' + + +@pytest.fixture +def T1064_telem_test_instance(): + return T1064Telem(STATUS, USAGE_STR) + + +def test_T1064_send(T1064_telem_test_instance, spy_send_telemetry): + T1064_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1064', + 'usage': USAGE_STR} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py new file mode 100644 index 000000000..050003e55 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py @@ -0,0 +1,26 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.telemetry.attack.t1105_telem import T1105Telem + + +DST_IP = '0.0.0.1' +FILENAME = 'virus.exe' +SRC_IP = '0.0.0.0' +STATUS = ScanStatus.USED + + +@pytest.fixture +def T1105_telem_test_instance(): + return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME) + + +def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry): + T1105_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1105', + 'filename': FILENAME, + 'src': SRC_IP, + 'dst': DST_IP} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py new file mode 100644 index 000000000..e47568c03 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py @@ -0,0 +1,22 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.telemetry.attack.t1106_telem import T1106Telem + + +STATUS = ScanStatus.USED +USAGE = UsageEnum.SMB + + +@pytest.fixture +def T1106_telem_test_instance(): + return T1106Telem(STATUS, USAGE) + + +def test_T1106_send(T1106_telem_test_instance, spy_send_telemetry): + T1106_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1106', + 'usage': USAGE.name} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py new file mode 100644 index 000000000..2635f429b --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py @@ -0,0 +1,22 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.telemetry.attack.t1107_telem import T1107Telem + + +PATH = 'path/to/file.txt' +STATUS = ScanStatus.USED + + +@pytest.fixture +def T1107_telem_test_instance(): + return T1107Telem(STATUS, PATH) + + +def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry): + T1107_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1107', + 'path': PATH} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py new file mode 100644 index 000000000..ce5562e1a --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py @@ -0,0 +1,22 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.telemetry.attack.t1129_telem import T1129Telem + + +STATUS = ScanStatus.USED +USAGE = UsageEnum.SMB + + +@pytest.fixture +def T1129_telem_test_instance(): + return T1129Telem(STATUS, USAGE) + + +def test_T1129_send(T1129_telem_test_instance, spy_send_telemetry): + T1129_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': 'T1129', + 'usage': USAGE.name} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py new file mode 100644 index 000000000..3c620e854 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py @@ -0,0 +1,26 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.model import VictimHost +from infection_monkey.telemetry.attack.t1197_telem import T1197Telem + + +MACHINE = VictimHost('127.0.0.1') +STATUS = ScanStatus.USED +USAGE_STR = '[Usage info]' + + +@pytest.fixture +def T1197_telem_test_instance(): + return T1197Telem(STATUS, MACHINE, USAGE_STR) + + +def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry): + T1197_telem_test_instance.send() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': 'T1197', + 'usage': USAGE_STR} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py new file mode 100644 index 000000000..da87dfe7d --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py @@ -0,0 +1,26 @@ +import pytest + +from common.utils.attack_utils import ScanStatus +from infection_monkey.model import VictimHost +from infection_monkey.telemetry.attack.t1222_telem import T1222Telem + + +COMMAND = 'echo hi' +MACHINE = VictimHost('127.0.0.1') +STATUS = ScanStatus.USED + + +@pytest.fixture +def T1222_telem_test_instance(): + return T1222Telem(STATUS, COMMAND, MACHINE) + + +def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry): + T1222_telem_test_instance.send() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': 'T1222', + 'command': COMMAND} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py new file mode 100644 index 000000000..b707242a3 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py @@ -0,0 +1,23 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.telemetry.attack.usage_telem import UsageTelem + + +STATUS = ScanStatus.USED +TECHNIQUE = 'T9999' +USAGE = UsageEnum.SMB + + +@pytest.fixture +def usage_telem_test_instance(): + return UsageTelem(TECHNIQUE, STATUS, USAGE) + + +def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry): + usage_telem_test_instance.send() + expected_data = {'status': STATUS.value, + 'technique': TECHNIQUE, + 'usage': USAGE.name} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py new file mode 100644 index 000000000..3743b7d76 --- /dev/null +++ b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py @@ -0,0 +1,25 @@ +import pytest + +from common.utils.attack_utils import ScanStatus, UsageEnum +from infection_monkey.model import VictimHost +from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem + + +MACHINE = VictimHost('127.0.0.1') +STATUS = ScanStatus.USED +TECHNIQUE = 'T9999' + + +@pytest.fixture +def victim_host_telem_test_instance(): + return VictimHostTelem(TECHNIQUE, STATUS, MACHINE) + + +def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry): + victim_host_telem_test_instance.send() + expected_data = {'machine': {'domain_name': MACHINE.domain_name, + 'ip_addr': MACHINE.ip_addr}, + 'status': STATUS.value, + 'technique': TECHNIQUE} + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py b/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py deleted file mode 100644 index 13dc02322..000000000 --- a/monkey/infection_monkey/telemetry/tests/test_attack_telem_classes.py +++ /dev/null @@ -1,55 +0,0 @@ -import pytest - -from common.utils.attack_utils import ScanStatus, UsageEnum -from infection_monkey.model import VictimHost -from infection_monkey.telemetry.attack.attack_telem import AttackTelem -from infection_monkey.telemetry.attack.usage_telem import UsageTelem -from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem - - -MACHINE = VictimHost('127.0.0.1') -STATUS = ScanStatus.USED -TECHNIQUE = 'T9999' -USAGE = UsageEnum.SMB - - -@pytest.fixture -def attack_telem_test_instance(): - return AttackTelem(TECHNIQUE, STATUS) - - -def test_attack_telem_send(attack_telem_test_instance, spy_send_telemetry): - attack_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': TECHNIQUE} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def usage_telem_test_instance(): - return UsageTelem(TECHNIQUE, STATUS, USAGE) - - -def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry): - usage_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': TECHNIQUE, - 'usage': USAGE.name} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def victim_host_telem_test_instance(): - return VictimHostTelem(TECHNIQUE, STATUS, MACHINE) - - -def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry): - victim_host_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, - 'status': STATUS.value, - 'technique': TECHNIQUE} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' diff --git a/monkey/infection_monkey/telemetry/tests/test_technique_telems.py b/monkey/infection_monkey/telemetry/tests/test_technique_telems.py deleted file mode 100644 index b2c73867d..000000000 --- a/monkey/infection_monkey/telemetry/tests/test_technique_telems.py +++ /dev/null @@ -1,159 +0,0 @@ -import pytest - -from common.utils.attack_utils import ScanStatus, UsageEnum -from infection_monkey.model import VictimHost -from infection_monkey.telemetry.attack.t1005_telem import T1005Telem -from infection_monkey.telemetry.attack.t1035_telem import T1035Telem -from infection_monkey.telemetry.attack.t1064_telem import T1064Telem -from infection_monkey.telemetry.attack.t1105_telem import T1105Telem -from infection_monkey.telemetry.attack.t1106_telem import T1106Telem -from infection_monkey.telemetry.attack.t1107_telem import T1107Telem -from infection_monkey.telemetry.attack.t1129_telem import T1129Telem -from infection_monkey.telemetry.attack.t1197_telem import T1197Telem -from infection_monkey.telemetry.attack.t1222_telem import T1222Telem - - -COMMAND = 'echo hi' -DST_IP = '0.0.0.1' -FILENAME = 'virus.exe' -GATHERED_DATA_TYPE = '[Type of data collected]' -INFO = '[Additional info]' -MACHINE = VictimHost('127.0.0.1') -PATH = 'path/to/file.txt' -SRC_IP = '0.0.0.0' -STATUS = ScanStatus.USED -USAGE = UsageEnum.SMB -USAGE_STR = '[Usage info]' - - -@pytest.fixture -def T1005_telem_test_instance(): - return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO) - - -def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry): - T1005_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1005', - 'gathered_data_type': GATHERED_DATA_TYPE, - 'info': INFO} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1035_telem_test_instance(): - return T1035Telem(STATUS, USAGE) - - -def test_T1035_send(T1035_telem_test_instance, spy_send_telemetry): - T1035_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1035', - 'usage': USAGE.name} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1064_telem_test_instance(): - return T1064Telem(STATUS, USAGE_STR) - - -def test_T1064_send(T1064_telem_test_instance, spy_send_telemetry): - T1064_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1064', - 'usage': USAGE_STR} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1105_telem_test_instance(): - return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME) - - -def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry): - T1105_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1105', - 'filename': FILENAME, - 'src': SRC_IP, - 'dst': DST_IP} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1106_telem_test_instance(): - return T1106Telem(STATUS, USAGE) - - -def test_T1106_send(T1106_telem_test_instance, spy_send_telemetry): - T1106_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1106', - 'usage': USAGE.name} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1107_telem_test_instance(): - return T1107Telem(STATUS, PATH) - - -def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry): - T1107_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1107', - 'path': PATH} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1129_telem_test_instance(): - return T1129Telem(STATUS, USAGE) - - -def test_T1129_send(T1129_telem_test_instance, spy_send_telemetry): - T1129_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1129', - 'usage': USAGE.name} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1197_telem_test_instance(): - return T1197Telem(STATUS, MACHINE, USAGE_STR) - - -def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry): - T1197_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, - 'status': STATUS.value, - 'technique': 'T1197', - 'usage': USAGE_STR} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' - - -@pytest.fixture -def T1222_telem_test_instance(): - return T1222Telem(STATUS, COMMAND, MACHINE) - - -def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry): - T1222_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, - 'status': STATUS.value, - 'technique': 'T1222', - 'command': COMMAND} - assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' From 15107eeea3d7be46b292a6af92bef2bf578a9f96 Mon Sep 17 00:00:00 2001 From: Shreya Date: Thu, 18 Feb 2021 22:40:30 +0530 Subject: [PATCH 11/13] Use constants/literals for tests --- .../tests/attack/test_t1197_telem.py | 8 +++--- .../tests/attack/test_t1222_telem.py | 8 +++--- .../tests/attack/test_victim_host_telem.py | 8 +++--- .../telemetry/tests/test_exploit_telem.py | 25 +++++++++++++++---- .../telemetry/tests/test_post_breach_telem.py | 13 ++++++++-- .../telemetry/tests/test_scan_telem.py | 11 +++++++- 6 files changed, 56 insertions(+), 17 deletions(-) diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py index 3c620e854..89d174090 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py @@ -5,7 +5,9 @@ from infection_monkey.model import VictimHost from infection_monkey.telemetry.attack.t1197_telem import T1197Telem -MACHINE = VictimHost('127.0.0.1') +DOMAIN_NAME = 'domain-name' +IP = '127.0.0.1' +MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED USAGE_STR = '[Usage info]' @@ -17,8 +19,8 @@ def T1197_telem_test_instance(): def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry): T1197_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, + expected_data = {'machine': {'domain_name': DOMAIN_NAME, + 'ip_addr': IP}, 'status': STATUS.value, 'technique': 'T1197', 'usage': USAGE_STR} diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py index da87dfe7d..7a8f88a75 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py @@ -6,7 +6,9 @@ from infection_monkey.telemetry.attack.t1222_telem import T1222Telem COMMAND = 'echo hi' -MACHINE = VictimHost('127.0.0.1') +DOMAIN_NAME = 'domain-name' +IP = '127.0.0.1' +MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED @@ -17,8 +19,8 @@ def T1222_telem_test_instance(): def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry): T1222_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, + expected_data = {'machine': {'domain_name': DOMAIN_NAME, + 'ip_addr': IP}, 'status': STATUS.value, 'technique': 'T1222', 'command': COMMAND} diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py index 3743b7d76..6a102983b 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py @@ -5,7 +5,9 @@ from infection_monkey.model import VictimHost from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem -MACHINE = VictimHost('127.0.0.1') +DOMAIN_NAME = 'domain-name' +IP = '127.0.0.1' +MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED TECHNIQUE = 'T9999' @@ -17,8 +19,8 @@ def victim_host_telem_test_instance(): def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry): victim_host_telem_test_instance.send() - expected_data = {'machine': {'domain_name': MACHINE.domain_name, - 'ip_addr': MACHINE.ip_addr}, + expected_data = {'machine': {'domain_name': DOMAIN_NAME, + 'ip_addr': IP}, 'status': STATUS.value, 'technique': TECHNIQUE} assert spy_send_telemetry.data == expected_data diff --git a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py index 1002a3cb3..92ca51ac3 100644 --- a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py @@ -5,11 +5,26 @@ from infection_monkey.model.host import VictimHost from infection_monkey.telemetry.exploit_telem import ExploitTelem -HOSTNAME = "hostname" DOMAIN_NAME = "domain-name" IP = "0.0.0.0" HOST = VictimHost(IP, DOMAIN_NAME) +HOST_AS_DICT = {'ip_addr': IP, + 'domain_name': DOMAIN_NAME, + 'os': {}, + 'services': {}, + 'icmp': False, + 'monkey_exe': None, + 'default_tunnel': None, + 'default_server': None} EXPLOITER = WmiExploiter(HOST) +EXPLOITER_NAME = 'WmiExploiter' +EXPLOITER_INFO = {'display_name': WmiExploiter._EXPLOITED_SERVICE, + 'started': '', + 'finished': '', + 'vulnerable_urls': [], + 'vulnerable_ports': [], + 'executed_cmds': []} +EXPLOITER_ATTEMPTS = [] RESULT = False @@ -22,10 +37,10 @@ def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): exploit_telem_test_instance.send() expected_data = { "result": RESULT, - "machine": HOST.as_dict(), - "exploiter": EXPLOITER.__class__.__name__, - "info": EXPLOITER.exploit_info, - "attempts": EXPLOITER.exploit_attempts, + "machine": HOST_AS_DICT, + "exploiter": EXPLOITER_NAME, + "info": EXPLOITER_INFO, + "attempts": EXPLOITER_ATTEMPTS, } assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "exploit" diff --git a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py index d38cfdbde..5dd1123ab 100644 --- a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py @@ -1,12 +1,21 @@ import pytest +from common.data.post_breach_consts import POST_BREACH_JOB_SCHEDULING from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs +from infection_monkey.post_breach.job_scheduling.linux_job_scheduling import \ + get_linux_commands_to_schedule_jobs +from infection_monkey.post_breach.job_scheduling.windows_job_scheduling import \ + get_windows_commands_to_schedule_jobs from infection_monkey.telemetry.post_breach_telem import PostBreachTelem +from infection_monkey.utils.environment import is_windows_os HOSTNAME = "hostname" IP = "0.0.0.0" PBA = ScheduleJobs() +PBA_COMMAND = get_windows_commands_to_schedule_jobs() if is_windows_os() else\ + ' '.join(get_linux_commands_to_schedule_jobs()) +PBA_NAME = POST_BREACH_JOB_SCHEDULING RESULT = False @@ -19,9 +28,9 @@ def post_breach_telem_test_instance(monkeypatch): def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry): post_breach_telem_test_instance.send() expected_data = { - "command": PBA.command, + "command": PBA_COMMAND, "result": RESULT, - "name": PBA.name, + "name": PBA_NAME, "hostname": HOSTNAME, "ip": IP, } diff --git a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py index d75aecdb1..f35d84289 100644 --- a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py @@ -6,6 +6,15 @@ from infection_monkey.model.host import VictimHost DOMAIN_NAME = "domain-name" IP = "0.0.0.0" HOST = VictimHost(IP, DOMAIN_NAME) +HOST_AS_DICT = {'ip_addr': IP, + 'domain_name': DOMAIN_NAME, + 'os': {}, + 'services': {}, + 'icmp': False, + 'monkey_exe': None, + 'default_tunnel': None, + 'default_server': None} +HOST_SERVICES = {} @pytest.fixture @@ -15,6 +24,6 @@ def scan_telem_test_instance(): def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry): scan_telem_test_instance.send() - expected_data = {"machine": HOST.as_dict(), "service_count": len(HOST.services)} + expected_data = {"machine": HOST_AS_DICT, "service_count": len(HOST_SERVICES)} assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.telem_category == "scan" From 8bd30ceb4c53af49461c9b0cd39d02e9b747ddd2 Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 19 Feb 2021 00:07:03 +0530 Subject: [PATCH 12/13] Format code using black --- .../tests/attack/test_attack_telem.py | 7 ++-- .../tests/attack/test_t1005_telem.py | 16 +++++---- .../tests/attack/test_t1035_telem.py | 6 ++-- .../tests/attack/test_t1064_telem.py | 8 ++--- .../tests/attack/test_t1105_telem.py | 20 ++++++----- .../tests/attack/test_t1106_telem.py | 6 ++-- .../tests/attack/test_t1107_telem.py | 8 ++--- .../tests/attack/test_t1129_telem.py | 6 ++-- .../tests/attack/test_t1197_telem.py | 19 ++++++----- .../tests/attack/test_t1222_telem.py | 19 ++++++----- .../tests/attack/test_usage_telem.py | 12 ++++--- .../tests/attack/test_victim_host_telem.py | 17 +++++----- .../telemetry/tests/conftest.py | 2 +- .../telemetry/tests/test_exploit_telem.py | 34 +++++++++++-------- .../telemetry/tests/test_post_breach_telem.py | 17 ++++++---- .../telemetry/tests/test_scan_telem.py | 19 ++++++----- .../telemetry/tests/test_state_telem.py | 1 + .../telemetry/tests/test_system_info_telem.py | 1 + .../telemetry/tests/test_trace_telem.py | 1 + 19 files changed, 116 insertions(+), 103 deletions(-) diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py index 750075fb2..5d14d0aad 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py @@ -5,7 +5,7 @@ from infection_monkey.telemetry.attack.attack_telem import AttackTelem STATUS = ScanStatus.USED -TECHNIQUE = 'T9999' +TECHNIQUE = "T9999" @pytest.fixture @@ -15,7 +15,6 @@ def attack_telem_test_instance(): def test_attack_telem_send(attack_telem_test_instance, spy_send_telemetry): attack_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': TECHNIQUE} + expected_data = {"status": STATUS.value, "technique": TECHNIQUE} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py index 757f0de13..528d6dca8 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py @@ -4,8 +4,8 @@ from common.utils.attack_utils import ScanStatus from infection_monkey.telemetry.attack.t1005_telem import T1005Telem -GATHERED_DATA_TYPE = '[Type of data collected]' -INFO = '[Additional info]' +GATHERED_DATA_TYPE = "[Type of data collected]" +INFO = "[Additional info]" STATUS = ScanStatus.USED @@ -16,9 +16,11 @@ def T1005_telem_test_instance(): def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry): T1005_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1005', - 'gathered_data_type': GATHERED_DATA_TYPE, - 'info': INFO} + expected_data = { + "status": STATUS.value, + "technique": "T1005", + "gathered_data_type": GATHERED_DATA_TYPE, + "info": INFO, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py index a3133fcb9..6c4e704bf 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py @@ -15,8 +15,6 @@ def T1035_telem_test_instance(): def test_T1035_send(T1035_telem_test_instance, spy_send_telemetry): T1035_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1035', - 'usage': USAGE.name} + expected_data = {"status": STATUS.value, "technique": "T1035", "usage": USAGE.name} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py index 575b57540..fce3107ff 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py @@ -5,7 +5,7 @@ from infection_monkey.telemetry.attack.t1064_telem import T1064Telem STATUS = ScanStatus.USED -USAGE_STR = '[Usage info]' +USAGE_STR = "[Usage info]" @pytest.fixture @@ -15,8 +15,6 @@ def T1064_telem_test_instance(): def test_T1064_send(T1064_telem_test_instance, spy_send_telemetry): T1064_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1064', - 'usage': USAGE_STR} + expected_data = {"status": STATUS.value, "technique": "T1064", "usage": USAGE_STR} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py index 050003e55..3b71bd56e 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py @@ -4,9 +4,9 @@ from common.utils.attack_utils import ScanStatus from infection_monkey.telemetry.attack.t1105_telem import T1105Telem -DST_IP = '0.0.0.1' -FILENAME = 'virus.exe' -SRC_IP = '0.0.0.0' +DST_IP = "0.0.0.1" +FILENAME = "virus.exe" +SRC_IP = "0.0.0.0" STATUS = ScanStatus.USED @@ -17,10 +17,12 @@ def T1105_telem_test_instance(): def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry): T1105_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1105', - 'filename': FILENAME, - 'src': SRC_IP, - 'dst': DST_IP} + expected_data = { + "status": STATUS.value, + "technique": "T1105", + "filename": FILENAME, + "src": SRC_IP, + "dst": DST_IP, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py index e47568c03..f51d124d0 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py @@ -15,8 +15,6 @@ def T1106_telem_test_instance(): def test_T1106_send(T1106_telem_test_instance, spy_send_telemetry): T1106_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1106', - 'usage': USAGE.name} + expected_data = {"status": STATUS.value, "technique": "T1106", "usage": USAGE.name} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py index 2635f429b..2e519a934 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py @@ -4,7 +4,7 @@ from common.utils.attack_utils import ScanStatus from infection_monkey.telemetry.attack.t1107_telem import T1107Telem -PATH = 'path/to/file.txt' +PATH = "path/to/file.txt" STATUS = ScanStatus.USED @@ -15,8 +15,6 @@ def T1107_telem_test_instance(): def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry): T1107_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1107', - 'path': PATH} + expected_data = {"status": STATUS.value, "technique": "T1107", "path": PATH} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py index ce5562e1a..f07e83ae7 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py @@ -15,8 +15,6 @@ def T1129_telem_test_instance(): def test_T1129_send(T1129_telem_test_instance, spy_send_telemetry): T1129_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': 'T1129', - 'usage': USAGE.name} + expected_data = {"status": STATUS.value, "technique": "T1129", "usage": USAGE.name} assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py index 89d174090..c67832281 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py @@ -5,11 +5,11 @@ from infection_monkey.model import VictimHost from infection_monkey.telemetry.attack.t1197_telem import T1197Telem -DOMAIN_NAME = 'domain-name' -IP = '127.0.0.1' +DOMAIN_NAME = "domain-name" +IP = "127.0.0.1" MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED -USAGE_STR = '[Usage info]' +USAGE_STR = "[Usage info]" @pytest.fixture @@ -19,10 +19,11 @@ def T1197_telem_test_instance(): def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry): T1197_telem_test_instance.send() - expected_data = {'machine': {'domain_name': DOMAIN_NAME, - 'ip_addr': IP}, - 'status': STATUS.value, - 'technique': 'T1197', - 'usage': USAGE_STR} + expected_data = { + "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP}, + "status": STATUS.value, + "technique": "T1197", + "usage": USAGE_STR, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py index 7a8f88a75..f053b9ca4 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py @@ -5,9 +5,9 @@ from infection_monkey.model import VictimHost from infection_monkey.telemetry.attack.t1222_telem import T1222Telem -COMMAND = 'echo hi' -DOMAIN_NAME = 'domain-name' -IP = '127.0.0.1' +COMMAND = "echo hi" +DOMAIN_NAME = "domain-name" +IP = "127.0.0.1" MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED @@ -19,10 +19,11 @@ def T1222_telem_test_instance(): def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry): T1222_telem_test_instance.send() - expected_data = {'machine': {'domain_name': DOMAIN_NAME, - 'ip_addr': IP}, - 'status': STATUS.value, - 'technique': 'T1222', - 'command': COMMAND} + expected_data = { + "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP}, + "status": STATUS.value, + "technique": "T1222", + "command": COMMAND, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py index b707242a3..1a4009be9 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py @@ -5,7 +5,7 @@ from infection_monkey.telemetry.attack.usage_telem import UsageTelem STATUS = ScanStatus.USED -TECHNIQUE = 'T9999' +TECHNIQUE = "T9999" USAGE = UsageEnum.SMB @@ -16,8 +16,10 @@ def usage_telem_test_instance(): def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry): usage_telem_test_instance.send() - expected_data = {'status': STATUS.value, - 'technique': TECHNIQUE, - 'usage': USAGE.name} + expected_data = { + "status": STATUS.value, + "technique": TECHNIQUE, + "usage": USAGE.name, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py index 6a102983b..98d62f05b 100644 --- a/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py +++ b/monkey/infection_monkey/telemetry/tests/attack/test_victim_host_telem.py @@ -5,11 +5,11 @@ from infection_monkey.model import VictimHost from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem -DOMAIN_NAME = 'domain-name' -IP = '127.0.0.1' +DOMAIN_NAME = "domain-name" +IP = "127.0.0.1" MACHINE = VictimHost(IP, DOMAIN_NAME) STATUS = ScanStatus.USED -TECHNIQUE = 'T9999' +TECHNIQUE = "T9999" @pytest.fixture @@ -19,9 +19,10 @@ def victim_host_telem_test_instance(): def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry): victim_host_telem_test_instance.send() - expected_data = {'machine': {'domain_name': DOMAIN_NAME, - 'ip_addr': IP}, - 'status': STATUS.value, - 'technique': TECHNIQUE} + expected_data = { + "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP}, + "status": STATUS.value, + "technique": TECHNIQUE, + } assert spy_send_telemetry.data == expected_data - assert spy_send_telemetry.telem_category == 'attack' + assert spy_send_telemetry.telem_category == "attack" diff --git a/monkey/infection_monkey/telemetry/tests/conftest.py b/monkey/infection_monkey/telemetry/tests/conftest.py index dab650174..cbb1b8074 100644 --- a/monkey/infection_monkey/telemetry/tests/conftest.py +++ b/monkey/infection_monkey/telemetry/tests/conftest.py @@ -11,5 +11,5 @@ def spy_send_telemetry(monkeypatch): _spy_send_telemetry.telem_category = None _spy_send_telemetry.data = None - monkeypatch.setattr(ControlClient, 'send_telemetry', _spy_send_telemetry) + monkeypatch.setattr(ControlClient, "send_telemetry", _spy_send_telemetry) return _spy_send_telemetry diff --git a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py index 92ca51ac3..a1d79ef64 100644 --- a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py @@ -8,22 +8,26 @@ from infection_monkey.telemetry.exploit_telem import ExploitTelem DOMAIN_NAME = "domain-name" IP = "0.0.0.0" HOST = VictimHost(IP, DOMAIN_NAME) -HOST_AS_DICT = {'ip_addr': IP, - 'domain_name': DOMAIN_NAME, - 'os': {}, - 'services': {}, - 'icmp': False, - 'monkey_exe': None, - 'default_tunnel': None, - 'default_server': None} +HOST_AS_DICT = { + "ip_addr": IP, + "domain_name": DOMAIN_NAME, + "os": {}, + "services": {}, + "icmp": False, + "monkey_exe": None, + "default_tunnel": None, + "default_server": None, +} EXPLOITER = WmiExploiter(HOST) -EXPLOITER_NAME = 'WmiExploiter' -EXPLOITER_INFO = {'display_name': WmiExploiter._EXPLOITED_SERVICE, - 'started': '', - 'finished': '', - 'vulnerable_urls': [], - 'vulnerable_ports': [], - 'executed_cmds': []} +EXPLOITER_NAME = "WmiExploiter" +EXPLOITER_INFO = { + "display_name": WmiExploiter._EXPLOITED_SERVICE, + "started": "", + "finished": "", + "vulnerable_urls": [], + "vulnerable_ports": [], + "executed_cmds": [], +} EXPLOITER_ATTEMPTS = [] RESULT = False diff --git a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py index 5dd1123ab..adff1ca62 100644 --- a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py @@ -2,10 +2,12 @@ import pytest from common.data.post_breach_consts import POST_BREACH_JOB_SCHEDULING from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs -from infection_monkey.post_breach.job_scheduling.linux_job_scheduling import \ - get_linux_commands_to_schedule_jobs -from infection_monkey.post_breach.job_scheduling.windows_job_scheduling import \ - get_windows_commands_to_schedule_jobs +from infection_monkey.post_breach.job_scheduling.linux_job_scheduling import ( + get_linux_commands_to_schedule_jobs, +) +from infection_monkey.post_breach.job_scheduling.windows_job_scheduling import ( + get_windows_commands_to_schedule_jobs, +) from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.environment import is_windows_os @@ -13,8 +15,11 @@ from infection_monkey.utils.environment import is_windows_os HOSTNAME = "hostname" IP = "0.0.0.0" PBA = ScheduleJobs() -PBA_COMMAND = get_windows_commands_to_schedule_jobs() if is_windows_os() else\ - ' '.join(get_linux_commands_to_schedule_jobs()) +PBA_COMMAND = ( + get_windows_commands_to_schedule_jobs() + if is_windows_os() + else " ".join(get_linux_commands_to_schedule_jobs()) +) PBA_NAME = POST_BREACH_JOB_SCHEDULING RESULT = False diff --git a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py index f35d84289..645cbbaf7 100644 --- a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py @@ -3,17 +3,20 @@ import pytest from infection_monkey.telemetry.scan_telem import ScanTelem from infection_monkey.model.host import VictimHost + DOMAIN_NAME = "domain-name" IP = "0.0.0.0" HOST = VictimHost(IP, DOMAIN_NAME) -HOST_AS_DICT = {'ip_addr': IP, - 'domain_name': DOMAIN_NAME, - 'os': {}, - 'services': {}, - 'icmp': False, - 'monkey_exe': None, - 'default_tunnel': None, - 'default_server': None} +HOST_AS_DICT = { + "ip_addr": IP, + "domain_name": DOMAIN_NAME, + "os": {}, + "services": {}, + "icmp": False, + "monkey_exe": None, + "default_tunnel": None, + "default_server": None, +} HOST_SERVICES = {} diff --git a/monkey/infection_monkey/telemetry/tests/test_state_telem.py b/monkey/infection_monkey/telemetry/tests/test_state_telem.py index a8beaf5ad..5d0eeabce 100644 --- a/monkey/infection_monkey/telemetry/tests/test_state_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_state_telem.py @@ -2,6 +2,7 @@ import pytest from infection_monkey.telemetry.state_telem import StateTelem + IS_DONE = True VERSION = "version" diff --git a/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py b/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py index 11692d4bb..dc362f7a7 100644 --- a/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py @@ -2,6 +2,7 @@ import pytest from infection_monkey.telemetry.system_info_telem import SystemInfoTelem + SYSTEM_INFO = {} diff --git a/monkey/infection_monkey/telemetry/tests/test_trace_telem.py b/monkey/infection_monkey/telemetry/tests/test_trace_telem.py index e4bb06a6b..9b297c4ea 100644 --- a/monkey/infection_monkey/telemetry/tests/test_trace_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_trace_telem.py @@ -2,6 +2,7 @@ import pytest from infection_monkey.telemetry.trace_telem import TraceTelem + MSG = "message" From 2bc27b48de9f9cd2f3937ce3bf80518c023e856b Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 19 Feb 2021 00:44:28 +0530 Subject: [PATCH 13/13] Use stub for PBA --- .../telemetry/tests/test_post_breach_telem.py | 25 +++++++------------ 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py index adff1ca62..ebd085a8d 100644 --- a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py +++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py @@ -1,31 +1,24 @@ import pytest -from common.data.post_breach_consts import POST_BREACH_JOB_SCHEDULING -from infection_monkey.post_breach.actions.schedule_jobs import ScheduleJobs -from infection_monkey.post_breach.job_scheduling.linux_job_scheduling import ( - get_linux_commands_to_schedule_jobs, -) -from infection_monkey.post_breach.job_scheduling.windows_job_scheduling import ( - get_windows_commands_to_schedule_jobs, -) from infection_monkey.telemetry.post_breach_telem import PostBreachTelem -from infection_monkey.utils.environment import is_windows_os HOSTNAME = "hostname" IP = "0.0.0.0" -PBA = ScheduleJobs() -PBA_COMMAND = ( - get_windows_commands_to_schedule_jobs() - if is_windows_os() - else " ".join(get_linux_commands_to_schedule_jobs()) -) -PBA_NAME = POST_BREACH_JOB_SCHEDULING +PBA_COMMAND = "run some pba" +PBA_NAME = "some pba" RESULT = False +class StubSomePBA: + def __init__(self): + self.name = PBA_NAME + self.command = PBA_COMMAND + + @pytest.fixture def post_breach_telem_test_instance(monkeypatch): + PBA = StubSomePBA() monkeypatch.setattr(PostBreachTelem, "_get_hostname_and_ip", lambda: (HOSTNAME, IP)) return PostBreachTelem(PBA, RESULT)