Agent: Use Timer in Log4ShellExploiter

monkey/infection_monkey/exploit/log4shell.py

@@ -18,6 +18,7 @@ from infection_monkey.network.info import get_free_tcp_port
 from infection_monkey.network.tools import get_interface_to_target
 from infection_monkey.utils.commands import build_monkey_commandline
 from infection_monkey.utils.monkey_dir import get_monkey_dir_path
+from infection_monkey.utils.timer import Timer
 
 logger = logging.getLogger(__name__)
 
@@ -173,11 +174,9 @@ class Log4ShellExploiter(WebRCE):
         return victim_called_back
 
     def _wait_for_victim_to_download_java_bytecode(self) -> bool:
-        start_time = time.time()
+        timer = Timer(Log4ShellExploiter.REQUEST_TO_VICTIM_TIMEOUT)
 
-        while not self._victim_timeout_expired(
-            start_time, Log4ShellExploiter.REQUEST_TO_VICTIM_TIMEOUT
-        ):
+        while not timer.is_expired():
             if self._exploit_class_http_server.exploit_class_downloaded():
                 self.exploit_result.exploitation_success = True
                 return True
@@ -191,9 +190,9 @@ class Log4ShellExploiter(WebRCE):
         return False
 
     def _wait_for_victim_to_download_agent(self):
-        start_time = time.time()
+        timer = Timer(LONG_REQUEST_TIMEOUT)
 
-        while not self._victim_timeout_expired(start_time, LONG_REQUEST_TIMEOUT):
+        while not timer.is_expired():
             if self._agent_http_server_thread.downloads > 0:
                 break
 
@@ -202,7 +201,3 @@ class Log4ShellExploiter(WebRCE):
 
             # TODO: if the http server got an error we're waiting for nothing here
             time.sleep(1)
-
-    @classmethod
-    def _victim_timeout_expired(cls, start_time: float, timeout: int) -> bool:
-        return timeout < (time.time() - start_time)