Payload Type Definition

from mythic_container.PayloadBuilder import *
from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *
import json

class Apollo(PayloadType):
   name = "apollo"
   file_extension = "exe"
   agent_type = AgentType.Agent
   author = "@djhohnstein, @its_a_feature_"
   mythic_encrypts = True
   supported_os = [
       SupportedOS.Windows
   ]
   semver = "2.3.51"
   wrapped_payloads = ["scarecrow_wrapper", "service_wrapper"]
   note = """
A fully featured .NET 4.0 compatible training agent. Version: {}.
NOTE: P2P Not compatible with v2.2 agents!
NOTE: v2.3.2+ has a different bof loader than 2.3.1 and are incompatible since their arguments are different
   """.format(semver)
   supports_dynamic_loading = True
   shellcode_format_options = ["Binary", "Base64", "C", "Ruby", "Python", "Powershell", "C#", "Hex"]
   shellcode_bypass_options = ["None", "Abort on fail", "Continue on fail"]
   supports_multiple_c2_instances_in_build = False
   supports_multiple_c2_in_build = False
   c2_parameter_deviations = {
       "http": {
           "get_uri": C2ParameterDeviation(supported=False),
           "query_path_name": C2ParameterDeviation(supported=False),
           #"headers": C2ParameterDeviation(supported=True, dictionary_choices=[
           #    DictionaryChoice(name="User-Agent", default_value="Hello", default_show=True),
           #    DictionaryChoice(name="HostyHost", default_show=False, default_value=""),
           #])
       }
   }
   build_parameters = [
       BuildParameter(
           name="output_type",
           parameter_type=BuildParameterType.ChooseOne,
           choices=["WinExe", "Shellcode", "Service", "Source"],
           default_value="WinExe",
           description="Output as shellcode, executable, sourcecode, or service.",
       ),
       BuildParameter(
           name="shellcode_format",
           parameter_type=BuildParameterType.ChooseOne,
           choices=shellcode_format_options,
           default_value="Binary",
           description="Donut shellcode format options.",
           group_name="Shellcode Options",
           hide_conditions=[
               HideCondition(name="output_type", operand=HideConditionOperand.NotEQ, value="Shellcode")
           ]
       ),
       BuildParameter(
           name="shellcode_bypass",
           parameter_type=BuildParameterType.ChooseOne,
           choices=shellcode_bypass_options,
           default_value="Continue on fail",
           description="Donut shellcode AMSI/WLDP/ETW Bypass options.",
           group_name="Shellcode Options",
           hide_conditions=[
               HideCondition(name="output_type", operand=HideConditionOperand.NotEQ, value="Shellcode")
           ]
       ),
       BuildParameter(
           name="adjust_filename",
           parameter_type=BuildParameterType.Boolean,
           default_value=False,
           description="Automatically adjust payload extension based on selected choices.",
       ),
       BuildParameter(
           name="debug",
           parameter_type=BuildParameterType.Boolean,
           default_value=False,
           description="Create a DEBUG version.",
       )
   ]
   c2_profiles = ["http", "smb", "tcp", "websocket"]
   agent_path = pathlib.Path(".") / "apollo" / "mythic"
   agent_code_path = pathlib.Path(".") / "apollo" / "agent_code"
   agent_icon_path = agent_path / "agent_functions" / "apollo.svg"
   build_steps = [
       BuildStep(step_name="Gathering Files", step_description="Copying files to temp location"),
       BuildStep(step_name="Compiling", step_description="Compiling with nuget and dotnet"),
       BuildStep(step_name="Donut", step_description="Converting to Shellcode"),
       BuildStep(step_name="Creating Service", step_description="Creating Service EXE from Shellcode")
   ]
   # here you can optionally define your own help function for getting help with your payload type's commands
   #async def command_help_function(self, msg: HelpFunctionMessage) -> HelpFunctionMessageResponse:
       #    return HelpFunctionMessageResponse(output=f"we did it!\nInput: {msg}", success=False)

    async def build(self) -> BuildResponse:
        # this function gets called to create an instance of your payload
        resp = BuildResponse(status=BuildStatus.Success)
        return resp

class BuildParameterType(str, Enum):
    """Types of parameters available for building payloads

    Attributes:
        String:
            A string value
        ChooseOne:
            A list of choices for the user to select exactly one
        ChooseMultiple:
            A list of choices for the user to select 0 or more
        Array:
            The user can supply multiple values in an Array format
        Date:
            The user can select a Date in YYYY-MM-DD format
        Dictionary:
            The user can supply a dictionary of values
        Boolean:
            The user can toggle a switch for True/False
        File:
            The user can select a file that gets uploaded - a file UUID gets passed in during build
        TypedArray:
            The user can supply an array where each element also has a drop-down option of choices
    """
    String = "String"
    ChooseOne = "ChooseOne"
    ChooseMultiple = "ChooseMultiple"
    Array = "Array"
    Date = "Date"
    Dictionary = "Dictionary"
    Boolean = "Boolean"
    File = "File"
    TypedArray = "TypedArray"

class BuildParameter:
    """Build Parameter Definition for use when generating payloads

    Attributes:
        name (str):
            Name of the parameter for scripting and for when building payloads
        description (str):
            Informative description displayed when building a payload
        default_value (any):
            Default value to pre-populate
        randomize (bool):
            Should this value be randomized (requires format_string)
        format_string (str):
            A regex used for randomizing values if randomize is true
        parameter_type (BuildParameterType):
            The type of parameter this is
        required (bool):
            Is this parameter required to have a non-empty value or not
        verifier_regex (str):
            Regex used to verify that the user typed something appropriate
        choices (list[str]):
            Choices for ChooseOne parameter type
        dictionary_choices (list[DictionaryChoice]):
            Configuration options for the Dictionary parameter type
        crypto_type (bool):
            Indicate if this value should be used to generate a crypto key or not
        group_name (str):
            An optional name used to group parameters together in the UI
        supported_os (list[str]):
            An optional list of supported operating systems where this parameter applies
        hide_conditions (list[HideCondition]):
            An optional list of conditions where this parameter should be hidden from view (this is a series of OR not AND conditions)
    """

    def __init__(
            self,
            name: str,
            parameter_type: BuildParameterType = None,
            description: str = None,
            required: bool = None,
            randomize: bool = None,
            format_string: str = "",
            crypto_type: bool = False,
            verifier_regex: str = None,
            default_value: any = None,
            choices: list[str] = None,
            dictionary_choices: list[DictionaryChoice] = None,
            value: any = None,
            verifier_func: callable = None,
            group_name: str = None,
            supported_os: list[str] = None,
            hide_conditions: list[HideCondition] = None,
    ):
        self.name = name
        self.verifier_func = verifier_func
        self.parameter_type = (
            parameter_type if parameter_type is not None else BuildParameterType.String
        )
        self.description = description if description is not None else ""
        self.required = required if required is not None else True
        self.verifier_regex = verifier_regex if verifier_regex is not None else ""
        self.default_value = default_value
        if value is None:
            self.value = default_value
        else:
            self.value = value
        self.choices = choices
        self.dictionary_choices = dictionary_choices
        self.crypto_type = crypto_type
        self.randomize = randomize
        self.format_string = format_string
        self.group_name = group_name
        self.supported_os = supported_os
        self.hide_conditions = hide_conditions

    def to_json(self):
        return {
            "name": self.name,
            "description": self.description,
            "default_value": self.default_value,
            "randomize": self.randomize,
            "format_string": self.format_string,
            "required": self.required,
            "parameter_type": self.parameter_type.value,
            "verifier_regex": self.verifier_regex,
            "crypto_type": self.crypto_type,
            "choices": self.choices,
            "dictionary_choices": [x.to_json() for x in
                                   self.dictionary_choices] if self.dictionary_choices is not None else None,
            "group_name": self.group_name,
            "supported_os": self.supported_os,
            "hide_conditions": [x.to_json() for x in self.hide_conditions] if self.hide_conditions is not None else None
        }

    def __str__(self):
        return json.dumps(self.to_json(), sort_keys=True, indent=2)

async def build(self) -> BuildResponse:
        # this function gets called to create an instance of your payload
        return BuildResponse(status=BuildStatus.Success)

for cmd in self.commands.get_commands():
    command_code += open(self.agent_code_path / "{}.js".format(cmd), 'r').read() + "\n"

for c2 in self.c2info:
    c2_code = ""
    try:
        profile = c2.get_c2profile()
        c2_code = open(
            self.agent_code_path
            / "c2_profiles"
            / "{}.js".format(profile["name"]),
            "r",
        ).read()
        for key, val in c2.get_parameters_dict().items():
            if key == "AESPSK":
                c2_code = c2_code.replace(key, val["enc_key"] if val["enc_key"] is not None else "")
            elif not isinstance(val, str):
                c2_code = c2_code.replace(key, json.dumps(val))
            else:
                c2_code = c2_code.replace(key, val)
    except Exception as p:
        build_msg += str(p)
        pass

class Apfell(PayloadType):
    name = "apfell"
    ...
    async def build ...

    async def on_new_callback(self, newCallback: PTOnNewCallbackAllData) -> PTOnNewCallbackResponse:
            new_task_resp = await SendMythicRPCTaskCreate(MythicRPCTaskCreateMessage(
                AgentCallbackUUID=newCallback.Callback.AgentCallbackID,
                CommandName="shell",
                Params="whoami",
            ))
            if new_task_resp.Success:
                return PTOnNewCallbackResponse(AgentCallbackUUID=newCallback.Callback.AgentCallbackID, Success=True)
            return PTOnNewCallbackResponse(AgentCallbackUUID=newCallback.Callback.AgentCallbackID, Success=False,
                                           Error=new_task_resp.Error)

class PTOtherServiceRPCMessage:
    """Request to call an RPC function of another C2 Profile or Payload Type

    Attributes:
        ServiceName (str): Name of the C2 Profile or Payload Type
        ServiceRPCFunction (str): Name of the function to call
        ServiceRPCFunctionArguments (dict): Arguments to that function

    Functions:
        to_json(self): return dictionary form of class
    """

    def __init__(self,
                 ServiceName: str = None,
                 service_name: str = None,
                 ServiceRPCFunction: str = None,
                 service_function: str = None,
                 ServiceRPCFunctionArguments: dict = None,
                 service_arguments: dict = None,
                 **kwargs):
        self.ServiceName = ServiceName
        if self.ServiceName is None:
            self.ServiceName = service_name
        self.ServiceRPCFunction = ServiceRPCFunction
        if self.ServiceRPCFunction is None:
            self.ServiceRPCFunction = service_function
        self.ServiceRPCFunctionArguments = ServiceRPCFunctionArguments
        if self.ServiceRPCFunctionArguments is None:
            self.ServiceRPCFunctionArguments = service_arguments
        for k, v in kwargs.items():
            logger.error(f"unknown kwarg {k} {v}")

    def to_json(self):
        return {
            "service_name": self.ServiceName,
            "service_function": self.ServiceRPCFunction,
            "service_arguments": self.ServiceRPCFunctionArguments
        }

    def __str__(self):
        return json.dumps(self.to_json(), sort_keys=True, indent=2)


class PTOtherServiceRPCMessageResponse:
    """Result of running an RPC call from another service

    Attributes:
        Success (bool): Did the RPC succeed or fail
        Error (str): Error message if the RPC check failed
        Result (dict): Result from the RPC

    Functions:
        to_json(self): return dictionary form of class
    """

    def __init__(self,
                 success: bool = None,
                 error: str = None,
                 result: dict = None,
                 Success: bool = None,
                 Error: str = None,
                 Result: dict = None,
                 **kwargs):
        self.Success = Success
        if self.Success is None:
            self.Success = success
        self.Error = Error
        if self.Error is None:
            self.Error = error
        self.Result = Result
        if self.Result is None:
            self.Result = result
        for k, v in kwargs.items():
            logger.error(f"unknown kwarg {k} {v}")

    def to_json(self):
        return {
            "success": self.Success,
            "error": self.Error,
            "result": self.Result
        }

    def __str__(self):
        return json.dumps(self.to_json(), sort_keys=True, indent=2)
custom_rpc_functions: dict[
        str, Callable[[PTOtherServiceRPCMessage], Awaitable[PTOtherServiceRPCMessageResponse]]] = {}