p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

agent: reformat test_base_telem_classes.py with black

This commit is contained in:
Mike Salvatore 2021-02-18 09:38:37 -05:00
parent c2ed31bde8
commit 0ac9ce949c
1 changed files with 28 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
monkey/infection_monkey/telemetry/tests/test_base_telem_classes.py
View File

@ -12,14 +12,14 @@ from infection_monkey.telemetry.trace_telem import TraceTelem
from infection_monkey.telemetry.tunnel_telem import TunnelTelem from infection_monkey.telemetry.tunnel_telem import TunnelTelem
DOMAIN_NAME = 'domain-name' DOMAIN_NAME = "domain-name"
HOSTNAME = 'hostname' HOSTNAME = "hostname"
IP = '0.0.0.0' IP = "0.0.0.0"
IS_DONE = True IS_DONE = True
MSG = 'message' MSG = "message"
RESULT = False RESULT = False
SYSTEM_INFO = {} SYSTEM_INFO = {}
VERSION = 'version' VERSION = "version"
HOST = VictimHost(IP, DOMAIN_NAME) HOST = VictimHost(IP, DOMAIN_NAME)
EXPLOITER = WmiExploiter(HOST) EXPLOITER = WmiExploiter(HOST)
PBA = ScheduleJobs() PBA = ScheduleJobs()
@ -32,33 +32,39 @@ def exploit_telem_test_instance():
def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry): def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry):
exploit_telem_test_instance.send() exploit_telem_test_instance.send()
expected_data = {'result': RESULT, expected_data = {
'machine': HOST.as_dict(), "result": RESULT,
'exploiter': EXPLOITER.__class__.__name__, "machine": HOST.as_dict(),
'info': EXPLOITER.exploit_info, "exploiter": EXPLOITER.__class__.__name__,
'attempts': EXPLOITER.exploit_attempts} "info": EXPLOITER.exploit_info,
"attempts": EXPLOITER.exploit_attempts,
}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "exploit" assert spy_send_telemetry.telem_category == "exploit"
@pytest.fixture @pytest.fixture
def post_breach_telem_test_instance(mocker): def post_breach_telem_test_instance(mocker):
mocker.patch('infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip', mocker.patch(
return_value=(HOSTNAME, IP)) "infection_monkey.telemetry.post_breach_telem.PostBreachTelem._get_hostname_and_ip",
return_value=(HOSTNAME, IP),
)
return PostBreachTelem(PBA, RESULT) return PostBreachTelem(PBA, RESULT)
def test_post_breach_telem_category(post_breach_telem_test_instance): 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): def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_telemetry):
post_breach_telem_test_instance.send() post_breach_telem_test_instance.send()
expected_data = {'command': PBA.command, expected_data = {
'result': RESULT, "command": PBA.command,
'name': PBA.name, "result": RESULT,
'hostname': HOSTNAME, "name": PBA.name,
'ip': IP} "hostname": HOSTNAME,
"ip": IP,
}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "post_breach" 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): def test_scan_telem_send(scan_telem_test_instance, spy_send_telemetry):
scan_telem_test_instance.send() scan_telem_test_instance.send()
expected_data = {'machine': HOST.as_dict(), expected_data = {"machine": HOST.as_dict(), "service_count": len(HOST.services)}
'service_count': len(HOST.services)}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "scan" 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): def test_state_telem_send(state_telem_test_instance, spy_send_telemetry):
state_telem_test_instance.send() state_telem_test_instance.send()
expected_data = {'done': IS_DONE, expected_data = {"done": IS_DONE, "version": VERSION}
'version': VERSION}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "state" 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): def test_trace_telem_send(trace_telem_test_instance, spy_send_telemetry):
trace_telem_test_instance.send() trace_telem_test_instance.send()
expected_data = {'msg': MSG} expected_data = {"msg": MSG}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "trace" 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): def test_tunnel_telem_send(tunnel_telem_test_instance, spy_send_telemetry):
tunnel_telem_test_instance.send() tunnel_telem_test_instance.send()
expected_data = {'proxy': None} expected_data = {"proxy": None}
assert spy_send_telemetry.data == expected_data assert spy_send_telemetry.data == expected_data
assert spy_send_telemetry.telem_category == "tunnel" assert spy_send_telemetry.telem_category == "tunnel"