From 66829f586cd9742fdb8df0c822dabffa468ff724 Mon Sep 17 00:00:00 2001
From: XiaoliChan <2209553467@qq.com>
Date: Wed, 16 Aug 2023 12:16:48 +0800
Subject: [PATCH] [smb] SMBexec: avoid looping by adding numbers of tires.

Signed-off-by: XiaoliChan <2209553467@qq.com>
---
 cme/protocols/smb.py            |  7 ++++---
 cme/protocols/smb/atexec.py     |  2 +-
 cme/protocols/smb/proto_args.py |  2 +-
 cme/protocols/smb/smbexec.py    | 23 +++++++++++++++--------
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/cme/protocols/smb.py b/cme/protocols/smb.py
index 8153ecb6..3d10a87a 100755
--- a/cme/protocols/smb.py
+++ b/cme/protocols/smb.py
@@ -729,7 +729,7 @@ class smb(connection):
                         self.kdcHost,
                         self.hash,
                         self.logger,
-                        self.args.atexec_tires
+                        self.args.get_output_tires
                     )  # self.args.share)
                     self.logger.info("Executed command via atexec")
                     break
@@ -753,7 +753,8 @@ class smb(connection):
                         self.hash,
                         self.args.share,
                         self.args.port,
-                        self.logger
+                        self.logger,
+                        self.args.get_output_tires
                     )
                     self.logger.info("Executed command via smbexec")
                     break
@@ -784,7 +785,7 @@ class smb(connection):
                     self.logger.highlight(line.strip())
             return output
         else:
-            self.logger.fail(f"Execute command failed {currnet_method}")
+            self.logger.fail(f"Execute command failed with {currnet_method}")
             return False
  
     @requires_admin
diff --git a/cme/protocols/smb/atexec.py b/cme/protocols/smb/atexec.py
index c5a1695c..9128d2c7 100755
--- a/cme/protocols/smb/atexec.py
+++ b/cme/protocols/smb/atexec.py
@@ -204,7 +204,7 @@ class TSCH_EXEC:
                         break
                     except Exception as e:
                         if tries >= self.__tires:
-                            self.logger.fail(f'ATEXEC: Get output file error, maybe go detection by AV software, please try "--atexec-tires" option. If it\'s still failing maybe something is blocking the schedule job, try another exec method')
+                            self.logger.fail(f'ATEXEC: Get output file error, maybe go detection by AV software, please try "--get-output-tires" option. If it\'s still failing maybe something is blocking the schedule job, try another exec method')
                             break
                         if str(e).find("SHARING") > 0 or str(e).find("STATUS_OBJECT_NAME_NOT_FOUND") >= 0:
                             sleep(3)
diff --git a/cme/protocols/smb/proto_args.py b/cme/protocols/smb/proto_args.py
index ffd81768..34fbf903 100644
--- a/cme/protocols/smb/proto_args.py
+++ b/cme/protocols/smb/proto_args.py
@@ -80,7 +80,7 @@ def proto_args(parser, std_parser, module_parser):
         cgroup.add_argument("--exec-method", choices={"wmiexec", "mmcexec", "smbexec", "atexec"}, default=None,
                             help="method to execute the command. Ignored if in MSSQL mode (default: wmiexec)")
         cgroup.add_argument("--wmiexec-timeout", help="WMIEXEC connection timeout, default is 5 secondes", type=int, default=5)
-        cgroup.add_argument("--atexec-tires", help="Number of times atexec tries to get results, default is 5", type=int, default=5)
+        cgroup.add_argument("--get-output-tires", help="Number of times atexec/smbexec tries to get results, default is 5", type=int, default=5)
         cgroup.add_argument("--codec", default="utf-8",
                             help="Set encoding used (codec) from the target's output (default "
                                  "\"utf-8\"). If errors are detected, run chcp.com at the target, "
diff --git a/cme/protocols/smb/smbexec.py b/cme/protocols/smb/smbexec.py
index cc0d7d9e..fa7bb826 100755
--- a/cme/protocols/smb/smbexec.py
+++ b/cme/protocols/smb/smbexec.py
@@ -26,7 +26,8 @@ class SMBEXEC:
         hashes=None,
         share=None,
         port=445,
-        logger=cme_logger
+        logger=cme_logger,
+        tires=None
     ):
         self.__host = host
         self.__share_name = "C$"
@@ -51,6 +52,7 @@ class SMBEXEC:
         self.__aesKey = aesKey
         self.__doKerberos = doKerberos
         self.__kdcHost = kdcHost
+        self.__tires = tires
         self.logger = logger
 
         if hashes is not None:
@@ -150,21 +152,26 @@ class SMBEXEC:
         if self.__retOutput is False:
             self.__outputBuffer = ""
             return
+        tires = 0
         while True:
             try:
+                cme_logger.info(f"Attempting to read {self.__share}\\{self.__output}")
                 self.__smbconnection.getFile(self.__share, self.__output, self.output_callback)
                 break
             except Exception as e:
-                print(e)
-                if str(e).find("STATUS_SHARING_VIOLATION") >= 0:
+                if tires >= self.__tires:
+                    self.logger.fail(f'SMBEXEC: Get output file error, maybe go detection by AV software, please try "--get-output-tires" option. If it\'s still failing maybe something is blocking the schedule job, try another exec method')
+                    break
+                if str(e).find("STATUS_SHARING_VIOLATION") >= 0 or str(e).find("STATUS_OBJECT_NAME_NOT_FOUND") >= 0:
                     # Output not finished, let's wait
                     sleep(2)
-                    pass
+                    tires += 1
                 else:
-                    self.logger.debug(e)
-                    pass
-
-        self.__smbconnection.deleteFile(self.__share, self.__output)
+                    raise
+        
+        if self.__outputBuffer:
+            cme_logger.debug(f"Deleting file {self.__share}\\{self.__output}")
+            self.__smbconnection.deleteFile(self.__share, self.__output)
 
     def execute_fileless(self, data):
         self.__output = gen_random_string(6)