From 4158ed802b0afc5945f412b982073f95287d6d15 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Fri, 19 Feb 2021 17:15:49 +0200
Subject: [PATCH] Refactored telemetry unit tests to json encode data the same
 way telemetries do.

---
 .../telemetry/tests/attack/test_attack_telem.py              | 4 ++++
 .../telemetry/tests/attack/test_t1005_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1035_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1064_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1105_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1106_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1107_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1129_telem.py               | 3 +++
 .../telemetry/tests/attack/test_t1197_telem.py               | 5 ++++-
 .../telemetry/tests/attack/test_t1222_telem.py               | 5 ++++-
 .../telemetry/tests/attack/test_usage_telem.py               | 3 +++
 .../telemetry/tests/attack/test_victim_host_telem.py         | 5 ++++-
 .../infection_monkey/telemetry/tests/test_exploit_telem.py   | 3 +++
 .../telemetry/tests/test_post_breach_telem.py                | 3 +++
 monkey/infection_monkey/telemetry/tests/test_scan_telem.py   | 4 ++++
 monkey/infection_monkey/telemetry/tests/test_state_telem.py  | 4 ++++
 .../telemetry/tests/test_system_info_telem.py                | 3 +++
 monkey/infection_monkey/telemetry/tests/test_trace_telem.py  | 4 ++++
 monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py | 4 ++++
 19 files changed, 65 insertions(+), 3 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 5d14d0aad..0812b1ea6 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_attack_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -16,5 +18,7 @@ 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 = json.dumps(expected_data, cls=attack_telem_test_instance.json_encoder)
+
     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
index 528d6dca8..6464d1121 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1005_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -22,5 +24,6 @@ def test_T1005_send(T1005_telem_test_instance, spy_send_telemetry):
         "gathered_data_type": GATHERED_DATA_TYPE,
         "info": INFO,
     }
+    expected_data = json.dumps(expected_data, cls=T1005_telem_test_instance.json_encoder)
     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
index 6c4e704bf..6313278ff 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1035_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -16,5 +18,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 = json.dumps(expected_data, cls=T1035_telem_test_instance.json_encoder)
     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
index fce3107ff..08031316b 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1064_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -16,5 +18,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 = json.dumps(expected_data, cls=T1064_telem_test_instance.json_encoder)
     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
index 3b71bd56e..4c3947141 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1105_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -24,5 +26,6 @@ def test_T1105_send(T1105_telem_test_instance, spy_send_telemetry):
         "src": SRC_IP,
         "dst": DST_IP,
     }
+    expected_data = json.dumps(expected_data, cls=T1105_telem_test_instance.json_encoder)
     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
index f51d124d0..db537cc51 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1106_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -16,5 +18,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 = json.dumps(expected_data, cls=T1106_telem_test_instance.json_encoder)
     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
index 2e519a934..993040244 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1107_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -16,5 +18,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 = json.dumps(expected_data, cls=T1107_telem_test_instance.json_encoder)
     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
index f07e83ae7..fa619f148 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1129_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -16,5 +18,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 = json.dumps(expected_data, cls=T1129_telem_test_instance.json_encoder)
     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
index c67832281..c5aa8874a 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1197_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -20,10 +22,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",
+        "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP},
         "usage": USAGE_STR,
     }
+    expected_data = json.dumps(expected_data, cls=T1197_telem_test_instance.json_encoder)
     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
index f053b9ca4..d3aeaddd6 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_t1222_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus
@@ -20,10 +22,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",
+        "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP},
         "command": COMMAND,
     }
+    expected_data = json.dumps(expected_data, cls=T1222_telem_test_instance.json_encoder)
     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
index 1a4009be9..983c1961d 100644
--- a/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/attack/test_usage_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -21,5 +23,6 @@ def test_usage_telem_send(usage_telem_test_instance, spy_send_telemetry):
         "technique": TECHNIQUE,
         "usage": USAGE.name,
     }
+    expected_data = json.dumps(expected_data, cls=usage_telem_test_instance.json_encoder)
     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
index 98d62f05b..014aadb8f 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
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from common.utils.attack_utils import ScanStatus, UsageEnum
@@ -20,9 +22,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,
+        "machine": {"domain_name": DOMAIN_NAME, "ip_addr": IP}
     }
+    expected_data = json.dumps(expected_data, cls=victim_host_telem_test_instance.json_encoder)
     assert spy_send_telemetry.data == expected_data
     assert spy_send_telemetry.telem_category == "attack"
diff --git a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py
index a1d79ef64..56d39fe06 100644
--- a/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_exploit_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.exploit.wmiexec import WmiExploiter
@@ -46,5 +48,6 @@ def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry):
         "info": EXPLOITER_INFO,
         "attempts": EXPLOITER_ATTEMPTS,
     }
+    expected_data = json.dumps(expected_data, cls=exploit_telem_test_instance.json_encoder)
     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 ebd085a8d..4aaaedb08 100644
--- a/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_post_breach_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
@@ -32,5 +34,6 @@ def test_post_breach_telem_send(post_breach_telem_test_instance, spy_send_teleme
         "hostname": HOSTNAME,
         "ip": IP,
     }
+    expected_data = json.dumps(expected_data, cls=post_breach_telem_test_instance.json_encoder)
     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
index 645cbbaf7..017a7d062 100644
--- a/monkey/infection_monkey/telemetry/tests/test_scan_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_scan_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.scan_telem import ScanTelem
@@ -28,5 +30,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 = json.dumps(expected_data, cls=scan_telem_test_instance.json_encoder)
+
     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
index 5d0eeabce..fe7bb3293 100644
--- a/monkey/infection_monkey/telemetry/tests/test_state_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_state_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.state_telem import StateTelem
@@ -15,5 +17,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 = json.dumps(expected_data, cls=state_telem_test_instance.json_encoder)
+
     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
index dc362f7a7..0caba8967 100644
--- a/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_system_info_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
@@ -14,5 +16,6 @@ def system_info_telem_test_instance():
 def test_system_info_telem_send(system_info_telem_test_instance, spy_send_telemetry):
     system_info_telem_test_instance.send()
     expected_data = SYSTEM_INFO
+    expected_data = json.dumps(expected_data, cls=system_info_telem_test_instance.json_encoder)
     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
index 9b297c4ea..567750e96 100644
--- a/monkey/infection_monkey/telemetry/tests/test_trace_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_trace_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.trace_telem import TraceTelem
@@ -14,5 +16,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 = json.dumps(expected_data, cls=trace_telem_test_instance.json_encoder)
+
     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
index 81a32bb44..eab763790 100644
--- a/monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py
+++ b/monkey/infection_monkey/telemetry/tests/test_tunnel_telem.py
@@ -1,3 +1,5 @@
+import json
+
 import pytest
 
 from infection_monkey.telemetry.tunnel_telem import TunnelTelem
@@ -11,5 +13,7 @@ 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 = json.dumps(expected_data, cls=tunnel_telem_test_instance.json_encoder)
+
     assert spy_send_telemetry.data == expected_data
     assert spy_send_telemetry.telem_category == "tunnel"