From ce5475812b064b23a7b91c85284066d0a7177313 Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Fri, 16 Feb 2024 09:30:08 +0000
Subject: [PATCH] refactor: injector

---
 phases/assembler.py |  6 +++---
 phases/injector.py  | 24 +++++++++++++-----------
 supermega.py        | 10 +++++++---
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/phases/assembler.py b/phases/assembler.py
index bb73fb7..ca732bb 100644
--- a/phases/assembler.py
+++ b/phases/assembler.py
@@ -19,12 +19,11 @@ def asm_to_shellcode(asm_in: FilePath, build_exe: FilePath, shellcode_out: FileP
         "/entry:AlignRSP"
     ])
     if not os.path.isfile(build_exe):
-        logger.error("Error")
-        return
+        raise Exception("Compiling failed")
     code = extract_code_from_exe(build_exe)
+    observer.add_code("generate_shc_from_asm", code) 
     with open(shellcode_out, 'wb') as f:
         f.write(code)
-    return code
 
 
 def merge_loader_payload(shellcode_in: FilePath, shellcode_out: FilePath, payload: FilePath, decoder_style: DecoderStyle):
@@ -36,6 +35,7 @@ def merge_loader_payload(shellcode_in: FilePath, shellcode_out: FilePath, payloa
         data_payload = input2.read()
 
     if project.decoder_style == DecoderStyle.PLAIN_1:
+        # Nothing to do
         pass
     elif project.decoder_style == DecoderStyle.XOR_1:
         xor_key = 0x42
diff --git a/phases/injector.py b/phases/injector.py
index b5db46c..eb92793 100644
--- a/phases/injector.py
+++ b/phases/injector.py
@@ -11,13 +11,14 @@ from project import project
 logger = logging.getLogger("Injector")
 
 
-def inject_exe(shc_file: FilePath):
-    exe_in: FilePath = project.inject_exe_in
-    exe_out: FilePath = project.inject_exe_out
-    exe_capabilities: ExeCapabilities = project.exe_capabilities
-
+def inject_exe(
+    shellcode_in: FilePath,
+    exe_in: FilePath,
+    exe_out: FilePath,
+    exe_capabilities: ExeCapabilities,
+):
     logger.info("--[ Injecting: {} into: {} -> {} ".format(
-        shc_file, exe_in, exe_out
+        shellcode_in, exe_in, exe_out
     ))
 
     # create copy of file exe_in to exe_out
@@ -29,7 +30,7 @@ def inject_exe(shc_file: FilePath):
         "python3.exe",
         "redbackdoorer.py",
         project.inject_mode,
-        shc_file,
+        shellcode_in,
         exe_out
     ])
 
@@ -40,8 +41,7 @@ def inject_exe(shc_file: FilePath):
         code = extract_code_from_exe(exe_out)
         for cap in exe_capabilities.get_all().values():
             if not cap.id in code:
-                logger.error("Capability ID {} not found, abort".format(cap.id))
-                raise Exception()
+                raise Exception("Capability ID {} not found, abort".format(cap.id))
             
             off = code.index(cap.id)
             current_address = off + exe_capabilities.image_base + exe_capabilities.text_virtaddr
@@ -53,10 +53,12 @@ def inject_exe(shc_file: FilePath):
                 current_address, destination_address
             )
             code = code.replace(cap.id, jmp)
-            write_code_section(exe_out, code)
+
+        # write back our patched code into the exe
+        write_code_section(exe_file=exe_out, new_data=code)
 
      
-def verify_injected_exe(exefile):
+def verify_injected_exe(exefile: FilePath):
     logger.info("---[ Verify infected exe: {} ".format(exefile))
     # remove indicator file
     pathlib.Path(project.verify_filename).unlink(missing_ok=True)
diff --git a/supermega.py b/supermega.py
index af8ad8e..fb5dca4 100644
--- a/supermega.py
+++ b/supermega.py
@@ -179,11 +179,10 @@ def start():
 
     # Convert: ASM -> Shellcode
     if project.generate_shc_from_asm:
-        code = phases.assembler.asm_to_shellcode(
+        phases.assembler.asm_to_shellcode(
             asm_in = main_asm_file, 
             build_exe = main_exe_file, 
             shellcode_out = main_shc_file)
-        observer.add_code("generate_shc_from_asm", code) 
     
     # Try: Starting the shellcode (rarely useful)
     if project.try_start_loader_shellcode:
@@ -230,7 +229,12 @@ def start():
     if project.inject:
         #debug_data["original_exe"] = file_readall_binary(options["inject_exe_in"])
 
-        phases.injector.inject_exe(main_shc_file)
+        phases.injector.inject_exe(
+            shellcode_in = main_shc_file,
+            exe_in = project.inject_exe_in,
+            exe_out = project.inject_exe_out,
+            exe_capabilities = project.exe_capabilities
+        )
         if project.verify:
             logger.info("--[ Verify final exe")
             if phases.injector.verify_injected_exe(project.inject_exe_out):