fix: better file cleanup

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):

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()