Add/modify tests for attack telems

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

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

monkey/infection_monkey/telemetry/attack/victim_host_telem_test.py

@@ -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)