Custom error messages on module execution (#118)

* Allow the custom generate function to send an error message back to the clients.

* fix check for csharpserver plugin status on

* import optimize

* docs for helper method

* missed some conversions

* safety check for tuples coming back from custom generate

empire/client/src/menus/UseModuleMenu.py

@@ -74,7 +74,11 @@ def execute(self):
             shell_return.daemon = True
             shell_return.start()
         elif 'error' in response.keys():
-            print(print_util.color('[!] Error: ' + response['error']))
+            if response['error'].startswith('[!]'):
+                msg = response['error']
+            else:
+                msg = f"[!] Error: {response['error']}"
+            print(print_util.color(msg))
 
     @command
     def generate(self):

empire/server/common/modules.py

@@ -59,8 +59,14 @@ def execute_module(self, module: PydanticModule, params: Dict, user_id: int) \
             return None, err
 
         module_data = self._generate_script(module, cleaned_options, self.main_menu.obfuscate, self.main_menu.obfuscateCommand)
+        if isinstance(module_data, tuple):
+            (module_data, err) = module_data
+        else:
+            # Not all modules return a tuple. If they just return a single value,
+            # we don't want to throw an unpacking error.
+            err = None
         if not module_data or module_data == "":
-            return None, 'module produced an empty script'
+            return None, err or 'module produced an empty script'
         if not module_data.isascii():
             return None, 'module source contains non-ascii characters'
 
@@ -179,14 +185,15 @@ def _set_default_values(module: PydanticModule, params: Dict):
             if params.get(option.name):
                 option.value = params[option.name]
 
-    def _generate_script(self, module: PydanticModule, params: Dict, obfuscate=False, obfuscate_command='') -> str:
+    def _generate_script(self, module: PydanticModule, params: Dict, obfuscate=False, obfuscate_command='') -> \
+            Tuple[Optional[str], Optional[str]]:
         """
         Generate the script to execute
         :param module: the execution parameters (already validated)
         :param params: the execution parameters
         :param obfuscate:
         :param obfuscate_command:
-        :return: the generated script
+        :return: tuple containing the generated script and an error if it exists
         """
         if module.advanced.custom_generate:
             return module.advanced.generate_class.generate(self.main_menu, module, params, obfuscate, obfuscate_command)
@@ -198,7 +205,7 @@ def _generate_script(self, module: PydanticModule, params: Dict, obfuscate=False
             return self._generate_script_csharp(module, params)
 
     @staticmethod
-    def _generate_script_python(module: PydanticModule, params: Dict):
+    def _generate_script_python(module: PydanticModule, params: Dict) -> Tuple[Optional[str], Optional[str]]:
         if module.script_path:
             with open(module.script_path, 'r') as stream:
                 script = stream.read()
@@ -209,9 +216,10 @@ def _generate_script_python(module: PydanticModule, params: Dict):
             if key.lower() != "agent" and key.lower() != "computername":
                 script = script.replace('{{ ' + key + ' }}', value).replace('{{' + key + '}}', value)
 
-        return script
+        return script, None
 
-    def _generate_script_powershell(self, module: PydanticModule, params: Dict, obfuscate=False, obfuscate_command=''):
+    def _generate_script_powershell(self, module: PydanticModule, params: Dict, obfuscate=False, obfuscate_command='') \
+            -> Tuple[Optional[str], Optional[str]]:
         if module.script_path:
             # Get preobfuscated module code
             if obfuscate:
@@ -256,11 +264,13 @@ def _generate_script_powershell(self, module: PydanticModule, params: Dict, obfu
             script = helpers.obfuscate(self.main_menu.installPath, psScript=script, obfuscationCommand=obfuscate_command)
         script = data_util.keyword_obfuscation(script)
 
-        return script
+        return script, None
 
-    def _generate_script_csharp(self, module: PydanticModule, params: Dict):
+    def _generate_script_csharp(self, module: PydanticModule, params: Dict) -> Tuple[Optional[str], Optional[str]]:
         try:
-            compiler = self.main_menu.loadedPlugins["csharpserver"]
+            compiler = self.main_menu.loadedPlugins.get("csharpserver")
+            if not compiler.status == 'ON':
+                return None, 'csharpserver plugin not running'
             compiler.do_send_message(module.compiler_yaml, module.name)
             dll_payload = open(self.main_menu.installPath + "/csharp/Covenant/Data/Tasks/CSharp/Compiled/net40/" + module.name + ".compiled" , "rb").read()
             dll_encoded = base64.b64encode(dll_payload).decode("UTF-8")
@@ -270,11 +280,13 @@ def _generate_script_csharp(self, module: PydanticModule, params: Dict):
                     if value and value != '':
                         script += "," + value
 
-            return script
+            return script, None
 
         except Exception as e:
             print(e)
-            print(helpers.color(f"[!] Compile Error"))
+            msg = f"[!] Compile Error"
+            print(helpers.color(msg))
+            return None, msg
 
     def _load_modules(self, root_path=''):
         """

empire/server/modules/exfiltration/Invoke_ExfilDataToGitHub.py

@@ -3,8 +3,9 @@
 from builtins import object
 from builtins import str
 
-from empire.server.utils import data_util
 from empire.server.common import helpers
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -136,8 +137,7 @@ def generate(self, obfuscate=False, obfuscationCommand=""):
         try:
             f = open(moduleSource, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(moduleSource)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(moduleSource))
 
         moduleCode = f.read()
         f.close()

empire/server/modules/external/generate_agent.py

@@ -2,12 +2,12 @@
 
 import os
 import string
-
 # Empire imports
 from builtins import object
 
-from empire.server.utils import data_util
 from empire.server.common import helpers
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -78,8 +78,7 @@ def execute(self):
         out_file = self.options['OutFile']['Value']
 
         if listener_name not in self.mainMenu.listeners.activeListeners:
-            print(helpers.color("[!] Error: %s not an active listener"))
-            return None
+            return handle_error_message("[!] Error: %s not an active listener")
 
         active_listener = self.mainMenu.listeners.activeListeners[listener_name]
 
@@ -113,8 +112,7 @@ def execute(self):
             agent_code += "\nInvoke-Empire -Servers @('%s') -StagingKey '%s' -SessionKey '%s' -SessionID '%s';" % (
                 host, staging_key, session_key, session_id)
         else:
-            print(helpers.color('[!] Only PowerShell agent generation is supported at this time.'))
-            return ''
+            return handle_error_message('[!] Only PowerShell agent generation is supported at this time.')
 
             # Get the random function name generated at install and patch the stager with the proper function name
         agent_code = data_util.keyword_obfuscation(agent_code)

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

@@ -1,14 +1,14 @@
 from __future__ import print_function
 
-from builtins import str
-from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
 import base64
+from builtins import object
+from builtins import str
 from typing import Dict
 
+from empire.server.common import helpers
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -53,16 +53,14 @@ def parse_assembly_args(args):
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
         module_code = f.read()
         f.close()
 
         try:
             f = open(params['Assembly'], 'rb')
         except:
-            print(helpers.color("[!] Could not read .NET assembly path at: " + str(params['Arguments'])))
-            return ""
+            return handle_error_message("[!] Could not read .NET assembly path at: " + str(params['Arguments']))
 
         assembly_data = f.read()
         f.close()

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

@@ -1,13 +1,12 @@
 from __future__ import print_function
 
-from builtins import str
 from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
+from builtins import str
 from typing import Dict
 
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -41,18 +40,16 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
-
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
+
         module_code = f.read()
         f.close()
 
         script = module_code
         script_end = ""
         if not main_menu.listeners.is_listener_valid(listener_name):
             # not a valid listener, return nothing for the script
-            print(helpers.color("[!] Invalid listener: %s" % (listener_name)))
-            return ''
+            return handle_error_message("[!] Invalid listener: %s" % (listener_name))
         else:
 
             l = main_menu.stagers.stagers['multi/launcher']
@@ -66,8 +63,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
             launcher = l.generate()
 
             if launcher == '':
-                print(helpers.color('[!] Error in launcher generation.'))
-                return ''
+                return handle_error_message('[!] Error in launcher generation.')
             else:
                 launcher_code = launcher.split(' ')[-1]
 

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

@@ -1,14 +1,14 @@
 from __future__ import print_function
 
-from builtins import str
-from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
 import base64
+from builtins import object
+from builtins import str
 from typing import Dict
 
+from empire.server.common import helpers
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -23,8 +23,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
 
         module_code = f.read()
         f.close()
@@ -35,8 +34,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
 
         #check if dllpath or PEUrl is set. Both are required params in their respective parameter sets.
         if params['DllPath'] == "" and params['PEUrl'] == "":
-            print(helpers.color("[!] Please provide a PEUrl or DllPath"))
-            return ""
+            return handle_error_message("[!] Please provide a PEUrl or DllPath")
         for option,values in params.items():
             if option.lower() != "agent":
                 if option.lower() == "dllpath":

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

@@ -1,13 +1,13 @@
 from __future__ import print_function
 
-from builtins import str
 from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
+from builtins import str
 from typing import Dict
 
+from empire.server.common import helpers
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -22,8 +22,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
 
         module_code = f.read()
         f.close()
@@ -35,8 +34,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         listener_name = params['Listener']
         if listener_name != "":
             if not main_menu.listeners.is_listener_valid(listener_name):
-                print(helpers.color("[!] Invalid listener: " + listener_name))
-                return ""
+                return handle_error_message("[!] Invalid listener: " + listener_name)
             else:
                 # TODO: redo pulling these listener configs...
                 #Old method no longer working

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

@@ -1,13 +1,13 @@
 from __future__ import print_function
 
-from builtins import str
 from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
+from builtins import str
 from typing import Dict
 
+from empire.server.common import helpers
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -22,8 +22,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
 
         module_code = f.read()
         f.close()

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

@@ -1,13 +1,13 @@
 from __future__ import print_function
 
-from builtins import str
 from builtins import object
-
-from empire.server.utils import data_util
-from empire.server.common import helpers
+from builtins import str
 from typing import Dict
 
+from empire.server.common import helpers
 from empire.server.common.module_models import PydanticModule
+from empire.server.utils import data_util
+from empire.server.utils.module_util import handle_error_message
 
 
 class Module(object):
@@ -20,8 +20,7 @@ def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool =
         try:
             f = open(module_source, 'r')
         except:
-            print(helpers.color("[!] Could not read module source path at: " + str(module_source)))
-            return ""
+            return handle_error_message("[!] Could not read module source path at: " + str(module_source))
 
         module_code = f.read()
         f.close()