Modify tests for attack telem classes and technique telems

- test `send()` instead of `get_data()` using fixture `spy_send_telemetry`
This commit is contained in:
Shreya 2021-02-18 20:13:33 +05:30
parent 4efdeeacc3
commit 08addff8c5
4 changed files with 60 additions and 93 deletions

View File

@ -3,6 +3,7 @@ from infection_monkey.telemetry.attack.usage_telem import AttackTelem
class T1064Telem(AttackTelem): class T1064Telem(AttackTelem):
def __init__(self, status, usage): def __init__(self, status, usage):
# TODO: rename parameter "usage" to avoid confusion with parameter "usage" in UsageTelem techniques
""" """
T1064 telemetry. T1064 telemetry.
:param status: ScanStatus of technique :param status: ScanStatus of technique

View File

@ -5,6 +5,7 @@ __author__ = "itay.mizeretz"
class T1197Telem(VictimHostTelem): class T1197Telem(VictimHostTelem):
def __init__(self, status, machine, usage): def __init__(self, status, machine, usage):
# TODO: rename parameter "usage" to avoid confusion with parameter "usage" in UsageTelem techniques
""" """
T1197 telemetry. T1197 telemetry.
:param status: ScanStatus of technique :param status: ScanStatus of technique

View File

@ -18,15 +18,12 @@ def attack_telem_test_instance():
return AttackTelem(TECHNIQUE, STATUS) return AttackTelem(TECHNIQUE, STATUS)
def test_attack_telem_category(attack_telem_test_instance): def test_attack_telem_send(attack_telem_test_instance, spy_send_telemetry):
assert attack_telem_test_instance.telem_category == 'attack' attack_telem_test_instance.send()
def test_attack_telem_get_data(attack_telem_test_instance):
actual_data = attack_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': TECHNIQUE} 'technique': TECHNIQUE}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -34,16 +31,13 @@ def usage_telem_test_instance():
return UsageTelem(TECHNIQUE, STATUS, USAGE) return UsageTelem(TECHNIQUE, STATUS, USAGE)
def test_usage_telem_category(usage_telem_test_instance): def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry):
assert usage_telem_test_instance.telem_category == 'attack' usage_telem_test_instance.send()
def test_usage_telem_get_data(usage_telem_test_instance):
actual_data = usage_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': TECHNIQUE, 'technique': TECHNIQUE,
'usage': USAGE.name} 'usage': USAGE.name}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -51,14 +45,11 @@ def victim_host_telem_test_instance():
return VictimHostTelem(TECHNIQUE, STATUS, MACHINE) return VictimHostTelem(TECHNIQUE, STATUS, MACHINE)
def test_victim_host_telem_category(victim_host_telem_test_instance): def test_victim_host_telem_send(victim_host_telem_test_instance, spy_send_telemetry):
assert victim_host_telem_test_instance.telem_category == 'attack' victim_host_telem_test_instance.send()
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, expected_data = {'machine': {'domain_name': MACHINE.domain_name,
'ip_addr': MACHINE.ip_addr}, 'ip_addr': MACHINE.ip_addr},
'status': STATUS.value, 'status': STATUS.value,
'technique': TECHNIQUE} 'technique': TECHNIQUE}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'

View File

@ -13,16 +13,17 @@ from infection_monkey.telemetry.attack.t1197_telem import T1197Telem
from infection_monkey.telemetry.attack.t1222_telem import T1222Telem 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]' GATHERED_DATA_TYPE = '[Type of data collected]'
INFO = '[Additional info]' INFO = '[Additional info]'
MACHINE = VictimHost('127.0.0.1') MACHINE = VictimHost('127.0.0.1')
PATH = 'path/to/file.txt'
SRC_IP = '0.0.0.0'
STATUS = ScanStatus.USED STATUS = ScanStatus.USED
USAGE = UsageEnum.SMB USAGE = UsageEnum.SMB
SRC_IP = '0.0.0.0' USAGE_STR = '[Usage info]'
DST_IP = '0.0.0.1'
FILENAME = 'virus.exe'
PATH = 'path/to/file.txt'
COMMAND = 'echo hi'
@pytest.fixture @pytest.fixture
@ -30,17 +31,14 @@ def T1005_telem_test_instance():
return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO) return T1005Telem(STATUS, GATHERED_DATA_TYPE, INFO)
def test_T1005_telem_category(T1005_telem_test_instance): def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry):
assert T1005_telem_test_instance.telem_category == 'attack' T1005_telem_test_instance.send()
def test_T1005_get_data(T1005_telem_test_instance):
actual_data = T1005_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1005', 'technique': 'T1005',
'gathered_data_type': GATHERED_DATA_TYPE, 'gathered_data_type': GATHERED_DATA_TYPE,
'info': INFO} 'info': INFO}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -48,33 +46,27 @@ def T1035_telem_test_instance():
return T1035Telem(STATUS, USAGE) return T1035Telem(STATUS, USAGE)
def test_T1035_telem_category(T1035_telem_test_instance): def test_T1035_send(T1035_telem_test_instance, spy_send_telemetry):
assert T1035_telem_test_instance.telem_category == 'attack' T1035_telem_test_instance.send()
def test_T1035_get_data(T1035_telem_test_instance):
actual_data = T1035_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1035', 'technique': 'T1035',
'usage': USAGE.name} 'usage': USAGE.name}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
def T1064_telem_test_instance(): def T1064_telem_test_instance():
return T1064Telem(STATUS, USAGE) return T1064Telem(STATUS, USAGE_STR)
def test_T1064_telem_category(T1064_telem_test_instance): def test_T1064_send(T1064_telem_test_instance, spy_send_telemetry):
assert T1064_telem_test_instance.telem_category == 'attack' T1064_telem_test_instance.send()
def test_T1064_get_data(T1064_telem_test_instance):
actual_data = T1064_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1064', 'technique': 'T1064',
'usage': USAGE} 'usage': USAGE_STR}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -82,18 +74,15 @@ def T1105_telem_test_instance():
return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME) return T1105Telem(STATUS, SRC_IP, DST_IP, FILENAME)
def test_T1105_telem_category(T1105_telem_test_instance): def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry):
assert T1105_telem_test_instance.telem_category == 'attack' T1105_telem_test_instance.send()
def test_T1105_get_data(T1105_telem_test_instance):
actual_data = T1105_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1105', 'technique': 'T1105',
'filename': FILENAME, 'filename': FILENAME,
'src': SRC_IP, 'src': SRC_IP,
'dst': DST_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 @pytest.fixture
@ -101,16 +90,13 @@ def T1106_telem_test_instance():
return T1106Telem(STATUS, USAGE) return T1106Telem(STATUS, USAGE)
def test_T1106_telem_category(T1106_telem_test_instance): def test_T1106_send(T1106_telem_test_instance, spy_send_telemetry):
assert T1106_telem_test_instance.telem_category == 'attack' T1106_telem_test_instance.send()
def test_T1106_get_data(T1106_telem_test_instance):
actual_data = T1106_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1106', 'technique': 'T1106',
'usage': USAGE.name} 'usage': USAGE.name}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -118,16 +104,13 @@ def T1107_telem_test_instance():
return T1107Telem(STATUS, PATH) return T1107Telem(STATUS, PATH)
def test_T1107_telem_category(T1107_telem_test_instance): def test_T1107_send(T1107_telem_test_instance, spy_send_telemetry):
assert T1107_telem_test_instance.telem_category == 'attack' T1107_telem_test_instance.send()
def test_T1107_get_data(T1107_telem_test_instance):
actual_data = T1107_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1107', 'technique': 'T1107',
'path': PATH} 'path': PATH}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -135,35 +118,29 @@ def T1129_telem_test_instance():
return T1129Telem(STATUS, USAGE) return T1129Telem(STATUS, USAGE)
def test_T1129_telem_category(T1129_telem_test_instance): def test_T1129_send(T1129_telem_test_instance, spy_send_telemetry):
assert T1129_telem_test_instance.telem_category == 'attack' T1129_telem_test_instance.send()
def test_T1129_get_data(T1129_telem_test_instance):
actual_data = T1129_telem_test_instance.get_data()
expected_data = {'status': STATUS.value, expected_data = {'status': STATUS.value,
'technique': 'T1129', 'technique': 'T1129',
'usage': USAGE.name} 'usage': USAGE.name}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
def T1197_telem_test_instance(): 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): def test_T1197_send(T1197_telem_test_instance, spy_send_telemetry):
assert T1197_telem_test_instance.telem_category == 'attack' T1197_telem_test_instance.send()
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, expected_data = {'machine': {'domain_name': MACHINE.domain_name,
'ip_addr': MACHINE.ip_addr}, 'ip_addr': MACHINE.ip_addr},
'status': STATUS.value, 'status': STATUS.value,
'technique': 'T1197', 'technique': 'T1197',
'usage': USAGE} 'usage': USAGE_STR}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'
@pytest.fixture @pytest.fixture
@ -171,15 +148,12 @@ def T1222_telem_test_instance():
return T1222Telem(STATUS, COMMAND, MACHINE) return T1222Telem(STATUS, COMMAND, MACHINE)
def test_T1222_telem_category(T1222_telem_test_instance): def test_T1222_send(T1222_telem_test_instance, spy_send_telemetry):
assert T1222_telem_test_instance.telem_category == 'attack' T1222_telem_test_instance.send()
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, expected_data = {'machine': {'domain_name': MACHINE.domain_name,
'ip_addr': MACHINE.ip_addr}, 'ip_addr': MACHINE.ip_addr},
'status': STATUS.value, 'status': STATUS.value,
'technique': 'T1222', 'technique': 'T1222',
'command': COMMAND} 'command': COMMAND}
assert actual_data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == 'attack'