Add ModuleValidationException, ModuleExecutionException, PluginExecut…

…ionException, module decorators, and deprecate tuple returns (#746)

* add module exceptions

* add plugin execution exception for feature parity. deprecate returning tuple from plugin for errors

* fix tests

* add additional tests

* update deprecation waring to link to wiki

* fix last test

CHANGELOG.md

@@ -14,6 +14,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+-   Added validation and execution exceptions for modules to raise (@Vinnybod)
+-   Added decorators for module generate functions to automatically get the module_source and call finalize_module (@Vinnybod)
+-   Added execution exception to plugins (@Vinnybod)
+
+### Deprecated
+
+-   Returning tuples from module generate functions is deprecated
+    -   To return a 400, raise a `ModuleValidationException`
+    -   To return a 500, raise a `ModuleExecutionException`
+    -   Stop using `handle_error_message`
+-   Returning tuples from plugin execution functions is deprecated
+    -   To return a 400, raise a `PluginValidationException`
+    -   To return a 500, raise a `PluginExecutionException`
+
 ### Changed
 
 -   Migrated some Pydantic and FastAPI usage away from deprecated features (@Vinnybod)

docs/module-development/powershell-modules.md

@@ -63,8 +63,13 @@ The generate function **should** treat these parameters as read only, to not cau
 ```python
 class Module(object):
     @staticmethod
-    def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool = False, obfuscation_command: str = "") -> Tuple[Optiona[str], Optional[str]]:
-        pass
+    def generate(
+        main_menu: MainMenu,
+        module: EmpireModule,
+        params: dict,
+        obfuscate: bool = False,
+        obfuscation_command: str = "",
+    ):
 ```
 
 Examples of modules that use this custom generate function:
@@ -73,12 +78,119 @@ Examples of modules that use this custom generate function:
 * [invoke\_assembly](https://github.com/BC-SECURITY/Empire/blob/master/empire/server/modules/powershell/code\_execution/invoke\_assembly.py)
 * [seatbelt](https://github.com/BC-SECURITY/Empire/blob/master/empire/server/modules/powershell/situational\_awareness/host/seatbelt.py)
 
-If an error occurs during the execution of the generate function, return the error message using `handle_error_message`, which will ensure that the client receives the error message in the REST response.
+#### Error Handling
+
+If an error occurs during the execution of the generate function and it goes unchecked,
+the client will receive a 500 error.
+
+There are two Exceptions that can be raised by the generate function:
+**ModuleValidationException**: This exception should be raised if the module fails validation. This will return a 400 error to the client with the error message.
+**ModuleExecutionException**: This exception should be raised if the module fails execution. This will return a 500 error to the client with the error message.
+
+```python
+raise ModuleValidationException("Error Message")
+raise ModuleExecutionException("Error Message")
+```
+
+##### Deprecated
+
+Previously, it was recommended that the generate function return a tuple of the script and the error.
+`handle_error_message` was provided as a helper function to handle this tuple.
+
+This is no longer recommended, but is still supported. Please migrate away from the tuple return type
+to raising exceptions. The tuple return type will be removed in a future major release.
+
+#### Functions
 
 `get_module_source` is used pull the script from the yaml file defined in **script\_path**. Once the script has been loaded, it will determine if obfuscation is enabled and obfuscate it.
 
 `finialize_module` will combine the `script` and `script_end` into a single script and then will apply obfuscation, if it is enabled.
 
+
+#### Decorators
+
+`@auto_get_source` is a decorator that will automatically call `get_module_source` and pass the script to the decorated function.
+To use this decorator, the function must have a `script` kwarg and the `script_path` must be set in the yaml config.
+
+```python
+@staticmethod
+@auto_get_source
+def generate(
+    main_menu: MainMenu,
+    module: EmpireModule,
+    params: dict,
+    obfuscate: bool = False,
+    obfuscation_command: str = "",
+    script: str = "",
+):
+    # do stuff
+    ...
+
+# The above is the equivalent of:
+@staticmethod
+def generate(
+    main_menu: MainMenu,
+    module: EmpireModule,
+    params: dict,
+    obfuscate: bool = False,
+    obfuscation_command: 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)
+
+    # do stuff
+    ...
+```
+
+`@auto_finalize` is a decorator that will automatically call `finalize_module` on the returned script from the decorated function.
+
+To use this decorator, the function must not utilize the deprecated tuple return type or the
+`handle_error_message` function. First migrate the function to raise exceptions before using this decorator.
+
+```python
+@staticmethod
+@auto_finalize
+def generate(
+    main_menu: MainMenu,
+    module: EmpireModule,
+    params: dict,
+    obfuscate: bool = False,
+    obfuscation_command: str = "",
+):
+    # Do stuff
+
+    return script, script_end
+
+# The above is the equivalent of:
+@staticmethod
+def generate(
+    main_menu: MainMenu,
+    module: EmpireModule,
+    params: dict,
+    obfuscate: bool = False,
+    obfuscation_command: str = "",
+):
+    # Do stuff
+
+    script, script_end = main_menu.modulesv2.finalize_module(
+        script=script,
+        script_end=script_end,
+        obfuscate=obfuscate,
+        obfuscate_command=obfuscation_command,
+    )
+
+    return script
+```
+
+
+
 ### String Formatting
 
 **option\_format\_string:** This tells Empire how to format all of the options before injecting them into the `script_end`. In most cases, the default option format string will be fine: `-{{ KEY }} "{{ VALUE }}"`.

docs/plugins/plugin-development.md

@@ -10,6 +10,20 @@ The execute function is the entry point for the plugin. It is called when the pl
 
 If the plugin doesn't have `**kwargs`, then no kwargs will be sent. This is to ensure backwards compatibility with plugin pre-5.2.
 
+### Error Handling
+
+If an error occurs during the execution of the plugin and it goes unchecked,
+the client will receive a 500 error.
+
+There are two Exceptions that can be raised by the plugin execution function:
+**PluginValidationException**: This exception should be raised if the plugin fails validation. This will return a 400 error to the client with the error message.
+**PluginExecutionException**: This exception should be raised if the plugin fails execution. This will return a 500 error to the client with the error message.
+
+```python
+raise PluginValidationException("Error Message")
+raise PluginExecutionException("Error Message")
+```
+
 ### Response
 
 Before the plugin's execute function is called, the core Empire code will validate the command arguments. If the arguments are invalid, the API will return a 400 error with the error message.
@@ -19,8 +33,15 @@ The execute function can return a String, a Boolean, or a Tuple of (Any, String)
 * None - The execution will be considered successful.
 * String - The string will be displayed to the user executing the plugin and the execution will be considered successful.
 * Boolean - If the boolean is True, the execution will be considered successful. If the boolean is False, the execution will be considered failed.
+
+#### Deprecated
+
 * Tuple - The tuple must be a tuple of (Any, String). The second value in the tuple represents an error message. The string will be displayed to the user executing the plugin and the execution will be considered failed.
 
+This is deprecated.
+Instead of returning an error message in a tuple, raise a `PluginValidationException` or `PluginExecutionException`.
+
+
 ```python
 def execute(self, command, **kwargs):
     ...
@@ -31,22 +52,11 @@ def execute(self, command, **kwargs):
     # return True
 
     # Failed execution
+    # raise PluginValidationException("Error Message")
+    # raise PluginExecutionException("Error Message")
     # return False, "Execution failed"
 ```
 
-### Custom Exceptions
-
-If the plugin raises a `PluginValidationException`, the API will return a 400 error with the error message.
-
-```python
-from empire.server.core.exceptions import PluginValidationException
-
-def execute(self, command, **kwargs):
-    ...
-
-    raise PluginValidationException("This is a validation error")
-```
-
 ## Plugin Tasks
 Plugins can store tasks. The data model looks pretty close to Agent tasks. This is for agent executions that:
 

empire/server/api/v2/agent/agent_task_api.py

@@ -46,6 +46,10 @@
 from empire.server.core.db import models
 from empire.server.core.db.models import AgentTaskStatus
 from empire.server.core.download_service import DownloadService
+from empire.server.core.exceptions import (
+    ModuleExecutionException,
+    ModuleValidationException,
+)
 from empire.server.server import main
 from empire.server.utils.data_util import is_port_in_use
 
@@ -256,12 +260,25 @@ async def create_task_module(
     current_user: CurrentUser,
     db_agent: models.Agent = Depends(get_agent),
 ):
-    resp, err = agent_task_service.create_task_module(
-        db, db_agent, module_request, current_user.id
-    )
+    try:
+        resp, err = agent_task_service.create_task_module(
+            db, db_agent, module_request, current_user.id
+        )
 
-    if err:
-        raise HTTPException(status_code=400, detail=err)
+        # This is for backwards compatibility with modules returning
+        # tuples for exceptions. All modules should remove returning
+        # tuples in favor of raising exceptions by Empire 6.0
+        if err:
+            raise HTTPException(status_code=400, detail=err)
+    except HTTPException as e:
+        # Propagate the HTTPException from above
+        raise e from None
+    except ModuleValidationException as e:
+        raise HTTPException(status_code=400, detail=str(e)) from e
+    except ModuleExecutionException as e:
+        raise HTTPException(status_code=500, detail=str(e)) from e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
     return domain_to_dto_task(resp)
 

empire/server/api/v2/plugin/plugin_api.py

@@ -14,7 +14,10 @@
 )
 from empire.server.api.v2.shared_dependencies import CurrentSession
 from empire.server.api.v2.shared_dto import BadRequestResponse, NotFoundResponse
-from empire.server.core.exceptions import PluginValidationException
+from empire.server.core.exceptions import (
+    PluginExecutionException,
+    PluginValidationException,
+)
 from empire.server.server import main
 
 plugin_service = main.pluginsv2
@@ -67,6 +70,8 @@ async def execute_plugin(
         )
     except PluginValidationException as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
+    except PluginExecutionException as e:
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
     if results is False or err:
         raise HTTPException(500, err or "internal plugin error")

empire/server/core/exceptions.py

@@ -1,2 +1,14 @@
 class PluginValidationException(Exception):
     pass
+
+
+class PluginExecutionException(Exception):
+    pass
+
+
+class ModuleValidationException(Exception):
+    pass
+
+
+class ModuleExecutionException(Exception):
+    pass