Agent: Fix LocketHTTPServer mypy error in MSSQLExploiter

guardicore · Oct 6, 2022 · ec617df · ec617df
1 parent 4784662
commit ec617df
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/monkey/infection_monkey/exploit/mssqlexec.py b/monkey/infection_monkey/exploit/mssqlexec.py
@@ -1,7 +1,7 @@
 import logging
 from pathlib import PureWindowsPath
 from time import sleep, time
-from typing import Iterable, Tuple
+from typing import Iterable, Optional, Tuple
 
 import pymssql
 
@@ -171,9 +171,12 @@ def _upload_agent(self, agent_path_on_victim: PureWindowsPath):
 
         self._run_agent_download_command(agent_path_on_victim)
 
-        MSSQLExploiter._stop_agent_server(http_thread)
+        if http_thread:
+            MSSQLExploiter._stop_agent_server(http_thread)
 
-    def _start_agent_server(self, agent_path_on_victim: PureWindowsPath) -> LockedHTTPServer:
+    def _start_agent_server(
+        self, agent_path_on_victim: PureWindowsPath
+    ) -> Optional[LockedHTTPServer]:
         self.agent_http_path, http_thread = HTTPTools.create_locked_transfer(
             self.host, str(agent_path_on_victim), self.agent_binary_repository
         )