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 0000000..49cc8ef Binary files /dev/null and b/out/.gitkeep differ 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