From 050a84d890ffff819c86b05111bc4d6f596a3199 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Fri, 7 Oct 2022 11:13:44 +0530
Subject: [PATCH] Agent: Move password restoration event publishing in
 zerologon

---
 monkey/infection_monkey/exploit/zerologon.py | 22 ++++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/monkey/infection_monkey/exploit/zerologon.py b/monkey/infection_monkey/exploit/zerologon.py
index 9110968e1..68761eff7 100644
--- a/monkey/infection_monkey/exploit/zerologon.py
+++ b/monkey/infection_monkey/exploit/zerologon.py
@@ -107,8 +107,6 @@ class ZerologonExploiter(HostExploiter):
                 self.exploit_info["password_restored"] = False
                 logger.info("System exploited but couldn't restore password!")
 
-            self._publish_password_restoration_event(success=self.exploit_info["password_restored"])
-
             self.store_extracted_creds_for_exploitation()
         else:
             logger.info("System was not exploited.")
@@ -317,15 +315,6 @@ class ZerologonExploiter(HostExploiter):
                     "nt_hash": nthash,
                 }
 
-    def _publish_password_restoration_event(self, success: bool):
-        password_restoration_event = PasswordRestorationEvent(
-            source=get_agent_id(),
-            target=self.host.ip_addr,
-            tags=PASSWORD_RESTORATION_EVENT_TAGS,
-            success=success,
-        )
-        self.agent_event_queue.publish(password_restoration_event)
-
     def store_extracted_creds_for_exploitation(self) -> None:
         for user in self._extracted_creds.keys():
             self.send_extracted_creds_as_credential_telemetry(
@@ -518,11 +507,22 @@ class ZerologonExploiter(HostExploiter):
 
     def assess_restoration_attempt_result(self, restoration_attempt_result) -> bool:
         if restoration_attempt_result:
+            self._publish_password_restoration_event(success=True)
             logger.debug("DC machine account password should be restored to its original value.")
             return True
 
+        self._publish_password_restoration_event(success=False)
         return False
 
+    def _publish_password_restoration_event(self, success: bool):
+        password_restoration_event = PasswordRestorationEvent(
+            source=get_agent_id(),
+            target=self.host.ip_addr,
+            tags=PASSWORD_RESTORATION_EVENT_TAGS,
+            success=success,
+        )
+        self.agent_event_queue.publish(password_restoration_event)
+
 
 class NetrServerPasswordSet(nrpc.NDRCALL):
     opnum = 6