From 08addff8c55e974b5bcb82c88c563a2790803fbb Mon Sep 17 00:00:00 2001 From: Shreya Date: Thu, 18 Feb 2021 20:13:33 +0530 Subject: [PATCH] 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'