From 31c5d637107d0f22e99c0382b7ef54f39d0efc4b Mon Sep 17 00:00:00 2001
From: Marius Ciepluch <11855163+norandom@users.noreply.github.com>
Date: Mon, 29 Jul 2024 10:29:15 +0200
Subject: [PATCH] Improved simulator

---
 .../Simulator/simulator.py                    | 47 +++++++++++++++----
 1 file changed, 39 insertions(+), 8 deletions(-)

diff --git a/2-1-initial-access-malware/Simulator/simulator.py b/2-1-initial-access-malware/Simulator/simulator.py
index aa2d3c2..2f46396 100755
--- a/2-1-initial-access-malware/Simulator/simulator.py
+++ b/2-1-initial-access-malware/Simulator/simulator.py
@@ -2,6 +2,33 @@ import os
 import sys
 import time
 import pyautogui
+import win32gui
+import win32con
+
+def close_cmd_window():
+    def enum_windows_callback(hwnd, result):
+        window_title = win32gui.GetWindowText(hwnd).lower()
+        if (win32gui.IsWindowVisible(hwnd) and 
+            'cmd.exe' in window_title and 
+            'Anaconda' not in window_title):
+            win32gui.PostMessage(hwnd, win32con.WM_CLOSE, 0, 0)
+    
+    win32gui.EnumWindows(enum_windows_callback, None)
+
+def close_excel_without_saving():
+    pyautogui.hotkey('alt', 'f4')
+    time.sleep(1)
+    
+    # Try to locate and click the "Don't Save" button
+    try:
+        dont_save_button = pyautogui.locateOnScreen('dont_save_button.png', confidence=0.8)
+        if dont_save_button:
+            pyautogui.click(dont_save_button)
+            print("Clicked 'Don't Save' button")
+        else:
+            print("Save dialogue not found, Excel may have closed without prompting")
+    except pyautogui.ImageNotFoundException:
+        print("Save dialogue not found, Excel may have closed without prompting")
 
 def open_excel_with_macros(file_path):
     # Get the directory of the current script/executable
@@ -18,10 +45,13 @@ def open_excel_with_macros(file_path):
     enable_button = None
 
     while (time.time() - start_time) < timeout:
-        enable_button = pyautogui.locateCenterOnScreen(enable_button_image, confidence=0.8)
-        if enable_button:
-            pyautogui.click(enable_button)
-            break
+        try:
+            enable_button = pyautogui.locateCenterOnScreen(enable_button_image, confidence=0.8)
+            if enable_button:
+                pyautogui.click(enable_button)
+                break
+        except pyautogui.ImageNotFoundException:
+            pass
         time.sleep(1)  # Check every 1 second
 
     if not enable_button:
@@ -31,12 +61,13 @@ def open_excel_with_macros(file_path):
     time.sleep(10)  # Adjust time based on expected macro execution time
     
     # Close Excel without saving
-    pyautogui.hotkey('alt', 'f4')
-    time.sleep(1)
-    pyautogui.press('n')  # Press 'n' in response to Excel's save prompt
+    close_excel_without_saving()
+
+    # Close any cmd.exe windows that might have opened, except Anaconda prompt
+    close_cmd_window()
 
 def main():
-    directory = r'C:\Users\mariu\Desktop\Corpus'  # Adjust the path to your files
+    directory = r'C:\Users\student\Desktop\Corpus'  # Adjust the path to your files
     files = os.listdir(directory)
     excel_files = [file for file in files if file.endswith(('.xlsx', '.xlsm'))]