v0.4.0

CHANGELOG.md

@@ -1,4 +1,13 @@
 
+## [v0.4.0] - 2024-01-10
+
+### Changed
+
+- Added a new optional `on_new_callback` function to the PayloadType class
+  - This allows you to take additional actions on a new callback based on your payload type
+- Added new MythicRPC* functions for searching edges associated with a callback and for creating new tasks for a callback
+- Needs Mythic v3.2.12+
+
 ## [v0.3.6] - 2023-12-22
 
 ### Changed

mythic_container/MythicCommandBase.py

@@ -1658,6 +1658,87 @@ def __str__(self):
         return json.dumps(self.to_json(), sort_keys=True, indent=2)
 
 
+class PTOnNewCallbackAllData:
+    """A container for all information about a callback including the callback, build parameters, commands, etc.
+
+    Note: Lowercase names are used in the __init__ function to auto-populate from JSON, but attributes are upper case.
+
+    Attributes:
+        Callback (PTTaskMessageCallbackData): The information about this callback
+        Payload (PTTaskMessagePayloadData): The information about this callback's associated payload
+        Commands (list[str]): The names of all the commands currently loaded into this callback
+        PayloadType (str): The name of the payload type
+        BuildParameters (list[MythicRPCPayloadConfigurationBuildParameter]): Information about the build parameters used to generate the payload for this callback
+        C2Profiles (list[MythicRPCPayloadConfigurationC2Profile]): Information about the c2 profiles associated with this callback and their values
+
+    Functions:
+        to_json(self): return dictionary form of class
+    """
+    args: TaskArguments
+
+    def __init__(self,
+                 callback: dict = {},
+                 build_parameters: list[dict] = [],
+                 commands: list[str] = [],
+                 payload: dict = {},
+                 c2info: list[dict] = [],
+                 payload_type: str = "",
+                 **kwargs):
+        self.Callback = PTTaskMessageCallbackData(**callback)
+        self.Payload = PTTaskMessagePayloadData(**payload)
+        self.Commands = commands
+        self.PayloadType = payload_type
+        self.BuildParameters = [MythicRPCPayloadConfigurationBuildParameter(**x) for x in build_parameters]
+        self.C2Profiles = [MythicRPCPayloadConfigurationC2Profile(**x) for x in c2info]
+
+        for k, v in kwargs.items():
+            logger.info(f"unknown kwarg {k} with value {v}")
+
+    def to_json(self):
+        return {
+            "callback": self.Callback.to_json(),
+            "build_parameters": [x.to_json() for x in self.BuildParameters],
+            "commands": self.Commands,
+            "payload": self.Payload.to_json(),
+            "c2info": [x.to_json() for x in self.C2Profiles],
+            "payload_type": self.PayloadType
+        }
+
+    def __str__(self):
+        return json.dumps(self.to_json(), sort_keys=True, indent=2)
+
+
+class PTOnNewCallbackResponse:
+    """The result of executing the on_new_callback function for a payload type
+
+    Attributes:
+        AgentCallbackID (str): The Agent Callback UUID of the new callback
+        Success (bool): Did the function execute successfully or not
+        Error (str): If the function failed to execute, return an error message here
+
+
+    Functions:
+        to_json(self): return dictionary form of class
+    """
+    def __init__(self,
+                 AgentCallbackID: str,
+                 Success: bool = True,
+                 Error: str = ""):
+        self.AgentCallbackID = AgentCallbackID
+        self.Success = Success
+        self.Error = Error
+
+    def to_json(self):
+        return {
+            "agent_callback_id": self.AgentCallbackID,
+            "success": self.Success,
+            "error": self.Error
+        }
+
+    def __str__(self):
+        return json.dumps(self.to_json(), sort_keys=True, indent=2)
+
+
 class PTTaskCompletionFunctionMessage:
     """A request to execute the completion function for a task or subtask
 

mythic_container/MythicGoRPC/__init__.py

@@ -5,6 +5,7 @@
 from .send_mythic_rpc_artifact_search import *
 from .send_mythic_rpc_callback_add_command import *
 from .send_mythic_rpc_callback_create import *
+from .send_mythic_rpc_callback_edge_search import *
 from .send_mythic_rpc_callback_decrypt_bytes import *
 from .send_mythic_rpc_callback_encrypt_bytes import *
 from .send_mythic_rpc_callback_remove_command import *
@@ -41,6 +42,7 @@
 from .send_mythic_rpc_proxy_stop import *
 from .send_mythic_rpc_response_create import *
 from .send_mythic_rpc_response_search import *
+from .send_mythic_rpc_task_create import *
 from .send_mythic_rpc_task_create_subtask import *
 from .send_mythic_rpc_task_create_subtask_group import *
 from .send_mythic_rpc_task_display_to_real_id_search import *

mythic_container/MythicGoRPC/send_mythic_rpc_task_create.py

@@ -0,0 +1,60 @@
+import mythic_container
+from mythic_container.logging import logger
+
+MYTHIC_RPC_TASK_CREATE = "mythic_rpc_task_create"
+
+
+class MythicRPCTaskCreateMessage:
+    def __init__(self,
+                 AgentCallbackID: str,
+                 CommandName: str = None,
+                 Params: str = None,
+                 ParameterGroupName: str = None,
+                 Token: int = None,
+                 **kwargs):
+        self.AgentCallbackID = AgentCallbackID
+        self.CommandName = CommandName
+        self.Params = Params
+        self.ParameterGroupName = ParameterGroupName
+        self.Token = Token
+        for k, v in kwargs.items():
+            logger.info(f"Unknown kwarg {k} - {v}")
+
+    def to_json(self):
+        return {
+            "agent_callback_id": self.AgentCallbackID,
+            "command_name": self.CommandName,
+            "params": self.Params,
+            "parameter_group_name": self.ParameterGroupName,
+            "token": self.Token
+        }
+
+
+class MythicRPCTaskCreateMessageResponse:
+    def __init__(self,
+                 success: bool = False,
+                 error: str = "",
+                 task_id: int = None,
+                 task_display_id: int = None,
+                 **kwargs):
+        self.Success = success
+        self.Error = error
+        self.TaskID = task_id
+        self.TaskDisplayID = task_display_id
+        for k, v in kwargs.items():
+            logger.info(f"Unknown kwarg {k} - {v}")
+
+    def to_json(self):
+        return {
+            "success": self.Success,
+            "error": self.Error,
+            "task_id": self.TaskID,
+            "task_display_id": self.TaskDisplayID,
+        }
+
+
+async def SendMythicRPCTaskCreate(
+        msg: MythicRPCTaskCreateMessage) -> MythicRPCTaskCreateMessageResponse:
+    response = await mythic_container.RabbitmqConnection.SendRPCDictMessage(queue=MYTHIC_RPC_TASK_CREATE,
+                                                                            body=msg.to_json())
+    return MythicRPCTaskCreateMessageResponse(**response)

mythic_container/PayloadBuilder.py

@@ -1,6 +1,8 @@
 from enum import Enum
 from abc import abstractmethod
 import base64
+
+from .MythicCommandBase import PTOnNewCallbackAllData, PTOnNewCallbackResponse
 from .logging import logger
 import json
 from collections.abc import Callable, Awaitable
@@ -84,6 +86,7 @@ class DictionaryChoice:
         default_value (str): Default value to fill in
 
     """
+
     def __init__(self,
                  name: str,
                  default_value: str = "",
@@ -131,6 +134,7 @@ class BuildParameter:
             Indicate if this value should be used to generate a crypto key or not
 
     """
+
     def __init__(
             self,
             name: str,
@@ -195,6 +199,7 @@ class C2ProfileParameters:
         get_parameters_dict:
             Get a dictionary of the parameters for the profile
     """
+
     def __init__(self, c2profile: dict, parameters: dict = None):
         self.parameters = {}
         self.c2profile = c2profile
@@ -221,6 +226,7 @@ class CommandList:
        get_commands:
            Get a list of the command names
    """
+
     def __init__(self, commands: [str] = None):
         self.commands = []
         if commands is not None:
@@ -264,6 +270,7 @@ class BuildResponse:
        get_commands:
            Get a list of the command names
    """
+
     def __init__(self, status: BuildStatus, payload: bytes = None, build_message: str = None, build_stderr: str = None,
                  build_stdout: str = None, updated_command_list: [str] = None, updated_filename: str = None):
         self.status = status
@@ -327,6 +334,7 @@ class BuildStep:
             The description of the step to display to users
 
    """
+
     def __init__(self,
                  step_name: str,
                  step_description: str):
@@ -354,6 +362,7 @@ class PTOtherServiceRPCMessage:
     Functions:
         to_json(self): return dictionary form of class
     """
+
     def __init__(self,
                  ServiceName: str = None,
                  service_name: str = None,
@@ -491,6 +500,8 @@ class PayloadType:
             Given an instance of a bare payload and all the configuration options that the user selected (build parameters, c2 profile parameters), build the payload
         get_parameter:
             Get the value of a build parameter
+        on_new_callback(self, newCallback):
+            Given the context of a new callback, perform some initial analysis and tasking
     """
     uuid: str = None
     c2info: [C2ProfileParameters] = None
@@ -590,6 +601,9 @@ def build_parameters(self):
     async def build(self) -> BuildResponse:
         pass
 
+    async def on_new_callback(self, newCallback: PTOnNewCallbackAllData) -> PTOnNewCallbackResponse:
+        return PTOnNewCallbackResponse(AgentCallbackID=newCallback.Callback.AgentCallbackID, Success=True)
+
     def get_parameter(self, name):
         for arg in self.build_parameters:
             if arg.name == name:
@@ -647,4 +661,5 @@ def to_json(self):
     def __str__(self):
         return json.dumps(self.to_json(), sort_keys=True, indent=2)
 
+
 payloadTypes: dict[str, PayloadType] = {}

mythic_container/__init__.py

@@ -1,9 +1,9 @@
 from .rabbitmq import rabbitmqConnectionClass
 from .mythic_service import start_and_run_forever, test_command
 
-containerVersion = "v1.1.1"
+containerVersion = "v1.2.0"
 
-PyPi_version = "0.3.7"
+PyPi_version = "0.4.0"
 
 RabbitmqConnection = rabbitmqConnectionClass()
 
@@ -16,6 +16,8 @@
 PT_RPC_COMMAND_TYPEDARRAY_PARSE_FUNCTION = "pt_command_typedarray_parse"
 PAYLOAD_BUILD_ROUTING_KEY = "payload_build"
 PT_BUILD_RESPONSE_ROUTING_KEY = "pt_build_response"
+PT_ON_NEW_CALLBACK_ROUTING_KEY = "pt_on_new_callback"
+PT_ON_NEW_CALLBACK_RESPONSE_ROUTING_KEY = "pt_on_new_callback_response"
 PT_BUILD_C2_RESPONSE_ROUTING_KEY = "pt_c2_build_response"
 PT_TASK_OPSEC_PRE_CHECK = "pt_task_opsec_pre_check"
 PT_TASK_OPSEC_PRE_CHECK_RESPONSE = "pt_task_opsec_pre_check_response"

mythic_container/agent_utils.py

@@ -89,6 +89,34 @@ async def buildWrapper(msg: bytes) -> None:
         )
 
 
+async def onNewCallback(msg: bytes) -> None:
+    try:
+        msgDict = ujson.loads(msg)
+        for name, pt in PayloadBuilder.payloadTypes.items():
+            if pt.name == msgDict["payload_type"]:
+                agent_builder = pt.__class__()
+                try:
+                    callbackData = mythic_container.MythicCommandBase.PTOnNewCallbackAllData(**msgDict)
+                    callback_response = await agent_builder.on_new_callback(callbackData)
+                    await mythic_container.RabbitmqConnection.SendDictDirectMessage(
+                        queue=mythic_container.PT_ON_NEW_CALLBACK_RESPONSE_ROUTING_KEY,
+                        body=callback_response.to_json()
+                    )
+                except Exception as b:
+                    logger.exception(f"[-] Failed to process on_new_callback for agent {pt.name}")
+                    await mythic_container.RabbitmqConnection.SendDictDirectMessage(
+                        queue=mythic_container.PT_ON_NEW_CALLBACK_RESPONSE_ROUTING_KEY,
+                        body={"success": False, "error": f"{traceback.format_exc()}\n{b}",
+                              "agent_callback_id": msgDict["callback"]["agent_callback_id"]}
+                    )
+    except Exception as e:
+        logger.exception(f"[-] Failed to process on_new_callback request")
+        await mythic_container.RabbitmqConnection.SendDictDirectMessage(
+            queue=mythic_container.PT_ON_NEW_CALLBACK_RESPONSE_ROUTING_KEY,
+            body={"success": False, "build_stderr": f"{traceback.format_exc()}\n{e}"}
+        )
+
+
 async def initialize_task(
         command_class: MythicCommandBase.CommandBase,
         message_json: dict,

mythic_container/mythic_service.py

@@ -53,6 +53,11 @@ async def startPayloadRabbitMQ(pt: PayloadBuilder.PayloadType) -> None:
         routing_key=getRoutingKey(pt.name, mythic_container.PAYLOAD_BUILD_ROUTING_KEY),
         handler=agent_utils.buildWrapper
     )))
+    payloadQueueTasks.append(asyncio.create_task(mythic_container.RabbitmqConnection.ReceiveFromMythicDirectExchange(
+        queue=getRoutingKey(pt.name, mythic_container.PT_ON_NEW_CALLBACK_ROUTING_KEY),
+        routing_key=getRoutingKey(pt.name, mythic_container.PT_ON_NEW_CALLBACK_ROUTING_KEY),
+        handler=agent_utils.onNewCallback
+    )))
     payloadQueueTasks.append(asyncio.create_task(mythic_container.RabbitmqConnection.ReceiveFromMythicDirectExchange(
         queue=getRoutingKey(pt.name, mythic_container.PT_TASK_OPSEC_PRE_CHECK),
         routing_key=getRoutingKey(pt.name, mythic_container.PT_TASK_OPSEC_PRE_CHECK),

setup.py

@@ -10,7 +10,7 @@
 # This call to setup() does all the work
 setup(
     name="mythic_container",
-    version="0.3.7",
+    version="0.4.0",
     description="Functionality for Mythic Services",
     long_description=README,
     long_description_content_type="text/markdown",