From d2a81bd191b4a0b2bb613425344fbad1bbf34638 Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Fri, 9 Feb 2024 20:41:56 +0000
Subject: [PATCH] refactor: syntax and comment updates

---
 pehelper.py      | 10 ++--------
 phases/ctoasm.py | 10 +++++-----
 supermega.py     | 15 ++++-----------
 3 files changed, 11 insertions(+), 24 deletions(-)

diff --git a/pehelper.py b/pehelper.py
index da29bbd..3212f2c 100644
--- a/pehelper.py
+++ b/pehelper.py
@@ -5,6 +5,8 @@ from keystone import Ks, KS_ARCH_X86, KS_MODE_64
 from capstone import Cs, CS_ARCH_X86, CS_MODE_64
 
 
+# keystone/capstone stuff
+
 def assemble_and_disassemble_jump(current_address, destination_address):
     #print("    Make jmp from 0x{:X} to 0x{:X}".format(
     #    current_address, destination_address
@@ -42,14 +44,6 @@ def extract_iat(pe):
                 continue
             imp_name = imp.name.decode('utf-8')
             imp_addr = imp.address
-            #pprint.pprint(imp.keys())
-            #print(type(imp))
-
-            #print("{}  {} - 0x{:08X}".format(
-            #    dll_name,
-            #    imp_name,
-            #    imp_addr
-            #))
 
             if not dll_name in iat:
                  iat[dll_name] = []
diff --git a/phases/ctoasm.py b/phases/ctoasm.py
index 0e95904..1a388e3 100644
--- a/phases/ctoasm.py
+++ b/phases/ctoasm.py
@@ -15,7 +15,7 @@ def make_c_to_asm(c_file, asm_file, payload_len, capabilities: ExeCapabilities):
         "fixup": "",
     }
 
-    # Phase 1: Compile
+    # Phase 1: C To Assembly
     print("---[ Compile: {} ]".format(c_file))
     run_process_checkret([
             config.get("path_cl"),
@@ -30,7 +30,7 @@ def make_c_to_asm(c_file, asm_file, payload_len, capabilities: ExeCapabilities):
         return
     asm["initial"] = file_readall_text(asm_file)
 
-    # Phase 2: Assembly cleanup
+    # Phase 1.1: Assembly cleanup
     asm_clean_file = asm_file + ".clean"
     print("---[ Cleanup: {} ]".format(asm_file))
     run_process_checkret([
@@ -45,7 +45,7 @@ def make_c_to_asm(c_file, asm_file, payload_len, capabilities: ExeCapabilities):
         shutil.move(asm_clean_file, asm_file)
         asm["cleanup"] = file_readall_text(asm_file)
 
-    # Phase 2: Assembly fixup
+    # Phase 1.2: Assembly fixup
     print("---[ Fixup  : {} ]".format(asm_file))
     if not fixup_asm_file(asm_file, payload_len, capabilities):
         print("Error: Fixup failed")
@@ -57,9 +57,9 @@ def make_c_to_asm(c_file, asm_file, payload_len, capabilities: ExeCapabilities):
 
 
 def bytes_to_asm_db(byte_data):
-    # Convert each byte to a string in hexadecimal format suffixed with 'h'
+    # Convert each byte to a string in hexadecimal format 
+    # prefixed with '0' and suffixed with 'h'
     hex_values = [f"0{byte:02x}H" for byte in byte_data]
-    # Join the hex values into a single string with ', ' as separator
     formatted_string = ', '.join(hex_values)
     return "\tDB " + formatted_string
 
diff --git a/supermega.py b/supermega.py
index caf675c..d21bdb5 100644
--- a/supermega.py
+++ b/supermega.py
@@ -140,14 +140,6 @@ main_exe_file = os.path.join(build_dir, "main.exe")
 main_shc_file = os.path.join(build_dir, "main.bin")
 
 debug_data = {
-    "loader_shellcode": b"",
-    "payload_shellcode": b"",
-    "final_shellcode": b"",
-
-    "asm_initial": "",
-    "asm_cleanup": "",
-    "asm_fixup": "",
-
     "original_exe": b"",
     "infected_exe": b"",
 }
@@ -248,14 +240,13 @@ def start(options):
     #        if not verify_shellcode("main-clean.bin"):
     #            return
 
+    # Merge shellcode/loader with payload
     if options["dataref_style"] == DataRefStyle.APPEND:
         print("--[ Merge stager: {} + {} -> {} ] ".format(main_shc_file, options["payload"], main_shc_file))
         with open(main_shc_file, 'rb') as input1:
             data_stager = input1.read()
-
         with open(options["payload"], 'rb') as input2:
             data_payload = input2.read()
-
         print("---[ Size: Stager: {} and Payload: {}  Sum: {} ]".format(
             len(data_stager), len(data_payload), len(data_stager)+len(data_payload)))
 
@@ -277,6 +268,7 @@ def start(options):
         # copy it to out
         shutil.copyfile(main_shc_file, os.path.join("out/", os.path.basename(main_shc_file)))
 
+    # inject merged loader into an exe
     if options["inject_exe"]:
         debug_data["original_exe"] = file_readall_binary(options["inject_exe_in"])
 
@@ -297,11 +289,12 @@ def start(options):
                 options["inject_exe_out"],
             ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
-    # dump
+    # dump the info i gathered
     file = open('latest.pickle', 'wb')
     pickle.dump(data, file)
     file.close()
 
+    # delete files
     if options["cleanup_files_on_exit"]:
         clean_files()