From a2c11759a4d5061c20592e20c42e363b525c4612 Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Mon, 22 Feb 2021 22:33:39 +0530
Subject: [PATCH] Add unit tests

---
 .../exploit/tests/test_zerologon.py           | 94 +++++++++++++++++++
 .../zerologon_utils/test_vuln_assessment.py   | 41 ++++++++
 monkey/infection_monkey/exploit/zerologon.py  |  2 +-
 3 files changed, 136 insertions(+), 1 deletion(-)
 create mode 100644 monkey/infection_monkey/exploit/tests/test_zerologon.py
 create mode 100644 monkey/infection_monkey/exploit/tests/zerologon_utils/test_vuln_assessment.py

diff --git a/monkey/infection_monkey/exploit/tests/test_zerologon.py b/monkey/infection_monkey/exploit/tests/test_zerologon.py
new file mode 100644
index 000000000..efc8a75e2
--- /dev/null
+++ b/monkey/infection_monkey/exploit/tests/test_zerologon.py
@@ -0,0 +1,94 @@
+import pytest
+
+from infection_monkey.exploit.zerologon import ZerologonExploiter
+from infection_monkey.model.host import VictimHost
+
+
+DOMAIN_NAME = "domain-name"
+IP = "0.0.0.0"
+NETBIOS_NAME = "NetBIOS Name"
+
+USERS = ["Administrator", "Bob"]
+RIDS = ["500", "1024"]
+LM_HASHES = ["abc123", "098zyx"]
+NT_HASHES = ["def456", "765vut"]
+
+
+@pytest.fixture
+def zerologon_exploiter_object(monkeypatch):
+    def mock_report_login_attempt(**kwargs):
+        return None
+
+    host = VictimHost(IP, DOMAIN_NAME)
+    obj = ZerologonExploiter(host)
+    monkeypatch.setattr(obj, "dc_name", NETBIOS_NAME, raising=False)
+    monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt)
+    return obj
+
+
+def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object):
+    dummy_exploit_attempt_result = {"ErrorCode": 0}
+    assert zerologon_exploiter_object.assess_exploit_attempt_result(
+        dummy_exploit_attempt_result
+    )
+
+
+def test_assess_exploit_attempt_result_with_error(zerologon_exploiter_object):
+    dummy_exploit_attempt_result = {"ErrorCode": 1}
+    assert not zerologon_exploiter_object.assess_exploit_attempt_result(
+        dummy_exploit_attempt_result
+    )
+
+
+def test_assess_restoration_attempt_result_restored(zerologon_exploiter_object):
+    dummy_restoration_attempt_result = object()
+    assert zerologon_exploiter_object.assess_restoration_attempt_result(
+        dummy_restoration_attempt_result
+    )
+
+
+def test_assess_restoration_attempt_result_not_restored(zerologon_exploiter_object):
+    dummy_restoration_attempt_result = False
+    assert not zerologon_exploiter_object.assess_restoration_attempt_result(
+        dummy_restoration_attempt_result
+    )
+
+
+def test__extract_user_creds_from_secrets_good_data(zerologon_exploiter_object):
+    mock_dumped_secrets = [
+        f"{USERS[i]}:{RIDS[i]}:{LM_HASHES[i]}:{NT_HASHES[i]}:::"
+        for i in range(len(USERS))
+    ]
+    expected_extracted_creds = {
+        USERS[0]: {
+            "RID": int(RIDS[0]),
+            "lm_hash": LM_HASHES[0],
+            "nt_hash": NT_HASHES[0],
+        },
+        USERS[1]: {
+            "RID": int(RIDS[1]),
+            "lm_hash": LM_HASHES[1],
+            "nt_hash": NT_HASHES[1],
+        },
+    }
+    assert (
+        zerologon_exploiter_object._extract_user_creds_from_secrets(mock_dumped_secrets)
+        is None
+    )
+    assert zerologon_exploiter_object._extracted_creds == expected_extracted_creds
+
+
+def test__extract_user_creds_from_secrets_bad_data(zerologon_exploiter_object):
+    mock_dumped_secrets = [
+        f"{USERS[i]}:{RIDS[i]}:::{LM_HASHES[i]}:{NT_HASHES[i]}:::"
+        for i in range(len(USERS))
+    ]
+    expected_extracted_creds = {
+        USERS[0]: {"RID": int(RIDS[0]), "lm_hash": "", "nt_hash": ""},
+        USERS[1]: {"RID": int(RIDS[1]), "lm_hash": "", "nt_hash": ""},
+    }
+    assert (
+        zerologon_exploiter_object._extract_user_creds_from_secrets(mock_dumped_secrets)
+        is None
+    )
+    assert zerologon_exploiter_object._extracted_creds == expected_extracted_creds
diff --git a/monkey/infection_monkey/exploit/tests/zerologon_utils/test_vuln_assessment.py b/monkey/infection_monkey/exploit/tests/zerologon_utils/test_vuln_assessment.py
new file mode 100644
index 000000000..cddb7522f
--- /dev/null
+++ b/monkey/infection_monkey/exploit/tests/zerologon_utils/test_vuln_assessment.py
@@ -0,0 +1,41 @@
+import pytest
+from nmb.NetBIOS import NetBIOS
+
+from infection_monkey.exploit.zerologon_utils.vuln_assessment import \
+    get_dc_details
+from infection_monkey.model.host import VictimHost
+
+
+DOMAIN_NAME = "domain-name"
+IP = "0.0.0.0"
+
+
+@pytest.fixture
+def host():
+    return VictimHost(IP, DOMAIN_NAME)
+
+
+def test_get_dc_details_multiple_netbios_names(host, monkeypatch):
+    def mock_queryIPForName(*args, **kwargs):
+        return NETBIOS_NAMES
+
+    monkeypatch.setattr(NetBIOS, "queryIPForName", mock_queryIPForName)
+
+    NETBIOS_NAMES = ["Name1", "Name2", "Name3"]
+    dc_ip, dc_name, dc_handle = get_dc_details(host)
+    assert dc_ip == IP
+    assert dc_name == NETBIOS_NAMES[0]
+    assert dc_handle == f"\\\\{NETBIOS_NAMES[0]}"
+
+
+def test_get_dc_details_no_netbios_names(host, monkeypatch):
+    def mock_queryIPForName(*args, **kwargs):
+        return NETBIOS_NAMES
+
+    monkeypatch.setattr(NetBIOS, "queryIPForName", mock_queryIPForName)
+
+    NETBIOS_NAMES = []
+    dc_ip, dc_name, dc_handle = get_dc_details(host)
+    assert dc_ip == IP
+    assert dc_name == ""
+    assert dc_handle == "\\\\"
diff --git a/monkey/infection_monkey/exploit/zerologon.py b/monkey/infection_monkey/exploit/zerologon.py
index 9f539e67f..5cf589fcd 100644
--- a/monkey/infection_monkey/exploit/zerologon.py
+++ b/monkey/infection_monkey/exploit/zerologon.py
@@ -407,7 +407,7 @@ class ZerologonExploiter(HostExploiter):
 
     def try_restoration_attempt(
         self, rpc_con: rpcrt.DCERPC_v5, original_pwd_nthash: str
-    ) -> bool:
+    ) -> Optional[object]:
         try:
             restoration_attempt_result = self.attempt_restoration(
                 rpc_con, original_pwd_nthash