From efe2039984edd15276de9d3caec9bae410543dd6 Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Sun, 4 Feb 2024 08:40:17 +0000
Subject: [PATCH] fix: better file cleanup

---
 helper.py    |  24 +++++++++++++++---------
 out/.gitkeep | Bin 0 -> 6 bytes
 supermega.py |   9 ++++++++-
 3 files changed, 23 insertions(+), 10 deletions(-)
 create mode 100644 out/.gitkeep

diff --git a/helper.py b/helper.py
index 30ed44a..a7cffc5 100644
--- a/helper.py
+++ b/helper.py
@@ -3,6 +3,7 @@ import os
 import pefile
 import time
 import shutil
+import pathlib
 
 SHC_VERIFY_SLEEP = 0.1
 
@@ -17,15 +18,20 @@ verify_filename = r'C:\Temp\a'
 
 
 def clean_files():
-    try:
-        os.remove("main.asm")       # generated from compiling source/main.c
-        os.remove("main-clean.asm") # cleaned for being a shellcode
-        os.remove("main-clean.exe") # assembled
-        os.remove("main-clean.bin")
-        os.remove("main-clean-append.bin")
-        os.remove(verify_filename)
-    except OSError:
-        pass
+    print("--[ Cleanup files ]")
+    files_to_clean = [
+        "main.asm",
+        "main.obj",
+        "main-clean.asm",
+        "main-clean.bin",
+        "main-clean-append.bin",
+        "main-clean.obj",
+        "mllink$.lnk",
+        verify_filename,
+        #"main-clean.exe",  # at the end as it may still shutdown?
+    ]
+    for file in files_to_clean:
+        pathlib.Path(file).unlink(missing_ok=True)
 
 
 def make_c_to_asm(c_file, asm_file, asm_clean_file, payload_len):
diff --git a/out/.gitkeep b/out/.gitkeep
new file mode 100644
index 0000000000000000000000000000000000000000..49cc8ef0e116cef009fe0bd72473a964bbd07f9b
GIT binary patch
literal 6
NcmezWkC%aq0RRg=0u=xN

literal 0
HcmV?d00001

diff --git a/supermega.py b/supermega.py
index c350008..4b4b8ef 100644
--- a/supermega.py
+++ b/supermega.py
@@ -25,6 +25,9 @@ options_default = {
     "verify": False,
 
     "cleanup_files_on_start": True,
+    "cleanup_files_on_exit": True, # all is just in out/
+
+
     "generate_asm_from_c": True,
     "generate_shc_from_asm": True,
     "test_loader_shellcode": False,
@@ -49,6 +52,8 @@ options_verify = {
     "verify": True,
 
     "cleanup_files_on_start": True,
+    "cleanup_files_on_exit": True, # all is just in out/
+
     "generate_asm_from_c": True,
     "generate_shc_from_asm": True,
     "test_loader_shellcode": False,
@@ -124,12 +129,14 @@ def main():
         # copy it to out
         shutil.copyfile("main-clean-append.bin", os.path.join("out/", "main-clean-append.bin"))
 
-
     if options["inject_exe"]:
         inject_exe("main-clean-append.bin", options["inject_exe_in"], options["inject_exe_out"])
         if options["verify"]:
             print("--[ Verify final exe ]")
             verify_injected_exe(options["inject_exe_out"])
 
+    if options["cleanup_files_on_exit"]:
+        clean_files()
+
 if __name__ == "__main__":
     main()
\ No newline at end of file