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

# Completion callback for BOF execution
async def coff_completion_callback(completionMsg: PTTaskCompletionFunctionMessage) -> PTTaskCompletionFunctionMessageResponse:
    try:
        out = ""
        response = PTTaskCompletionFunctionMessageResponse(Success=True, TaskStatus="success", Completed=True)
        
        # Get parent task ID
        parent_task_id = completionMsg.TaskData.Task.ID if completionMsg.TaskData and completionMsg.TaskData.Task else None
        if not parent_task_id:
            return PTTaskCompletionFunctionMessageResponse(Success=False, TaskStatus="error", Completed=True, Error="No parent task ID found")
        
        # Get subtask ID
        subtask_id = None
        if hasattr(completionMsg, 'SubtaskData') and completionMsg.SubtaskData:
            if hasattr(completionMsg.SubtaskData, 'Task') and completionMsg.SubtaskData.Task:
                subtask_id = completionMsg.SubtaskData.Task.ID
        
        if not subtask_id:
            # Try alternative attribute names
            if hasattr(completionMsg, 'Subtask') and completionMsg.Subtask:
                if hasattr(completionMsg.Subtask, 'ID'):
                    subtask_id = completionMsg.Subtask.ID
                elif hasattr(completionMsg.Subtask, 'TaskID'):
                    subtask_id = completionMsg.Subtask.TaskID
        
        if not subtask_id:
            await SendMythicRPCResponseCreate(MythicRPCResponseCreateMessage(
                TaskID=parent_task_id,
                Response="[ERROR] No subtask ID found in completion message"
            ))
            return PTTaskCompletionFunctionMessageResponse(Success=False, TaskStatus="error", Completed=True, Error="No subtask ID found")
        
        # Search for responses from the subtask
        responses = await SendMythicRPCResponseSearch(MythicRPCResponseSearchMessage(TaskID=subtask_id))
        
        if not responses.Success:
            error_detail = getattr(responses, 'Error', 'Unknown error') if hasattr(responses, 'Error') else 'Unknown error'
            await SendMythicRPCResponseCreate(MythicRPCResponseCreateMessage(
                TaskID=parent_task_id,
                Response=f"[ERROR] Failed to search responses for subtask {subtask_id}: {error_detail}"
            ))
            return PTTaskCompletionFunctionMessageResponse(Success=False, TaskStatus="error", Completed=True, Error=error_detail)
        
        # Process responses if they exist
        responses_list = None
        if hasattr(responses, 'Responses'):
            responses_list = responses.Responses
        elif hasattr(responses, 'responses'):
            responses_list = responses.responses
        elif hasattr(responses, 'Response'):
            responses_list = [responses.Response] if responses.Response else []
        
        if responses_list:
            for output in responses_list:
                if hasattr(output, 'Response'):
                    out += str(output.Response)
                elif hasattr(output, 'response'):
                    out += str(output.response)
                elif hasattr(output, 'Message'):
                    out += str(output.Message)
                elif hasattr(output, 'message'):
                    out += str(output.message)
                elif isinstance(output, str):
                    out += output
                elif isinstance(output, bytes):
                    out += output.decode('utf-8', errors='ignore')
                else:
                    out += str(output)
        
        # If no output was found, indicate successful execution
        if not out:
            out = "[BOF executed successfully but produced no output]"
            
        # Send the aggregated output back to the parent task
        await SendMythicRPCResponseCreate(MythicRPCResponseCreateMessage(
            TaskID=parent_task_id,
            Response=out
        ))
        return response
        
    except Exception as e:
        import traceback
        error_msg = f"[ERROR] Exception in coff_completion_callback: {str(e)}\n{traceback.format_exc()}"
        try:
            parent_task_id = completionMsg.TaskData.Task.ID if completionMsg.TaskData and completionMsg.TaskData.Task else None
            if parent_task_id:
                await SendMythicRPCResponseCreate(MythicRPCResponseCreateMessage(
                    TaskID=parent_task_id,
                    Response=error_msg
                ))
        except:
            pass  # If we can't send the error, at least return failure
        return PTTaskCompletionFunctionMessageResponse(Success=False, TaskStatus="error", Completed=True, Error=str(e))

class InlineExecuteArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="bof_name",
                cli_name="BOF",
                display_name="BOF",
                type=ParameterType.ChooseOne,
                dynamic_query_function=self.get_files,
                description="Already existing BOF to execute (e.g. whoami.x64.o)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=True,
                        group_name="Default",
                        ui_position=1
                    )
                ]
            ),
            CommandParameter(
                name="bof_file",
                display_name="New BOF",
                type=ParameterType.File,
                description="A new BOF to execute. After uploading once, you can just supply the bof_name parameter",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=True, 
                        group_name="New", 
                        ui_position=1,
                    )
                ]
            ),
            CommandParameter(
                name="bof_arguments",
                cli_name="Arguments",
                display_name="Arguments",
                type=ParameterType.TypedArray,
                default_value=[],
                choices=["int16", "int32", "string", "wchar", "base64", "bytes"],
                description="""Arguments to pass to the BOF via the following way:
                -s:123 or int16:123
                -i:123 or int32:123
                -z:hello or string:hello
                -Z:hello or wchar:hello
                -b:abc== or base64:abc==
                bytes:hexstring (hexadecimal string of bytes)""",
                typedarray_parse_function=self.get_arguments,
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=4
                    ),
                    ParameterGroupInfo(
                        required=False,
                        group_name="New",
                        ui_position=4
                    ),
                ]
            ),
        ]
        
    async def parse_arguments(self):
        # This is only called when user types command directly
        # Subtasks use parse_dictionary instead
        if len(self.command_line) > 0:
            raise ValueError("Must supply named arguments or use the modal")
        # If command_line is empty, do nothing (subtask will use parse_dictionary)

    async def parse_dictionary(self, dictionary_arguments):
         # Allowed argument keys
        expected_args = {"bof_arguments", "bof_file", "bof_name"}

        # Check if dictionary contains only allowed keys
        invalid_keys = set(dictionary_arguments.keys()) - expected_args
        if invalid_keys:
            raise ValueError(f"Invalid arguments provided: {', '.join(invalid_keys)}")
        
        # Load only allowed arguments
        filtered_arguments = {k: v for k, v in dictionary_arguments.items() if k in expected_args}
        self.load_args_from_dictionary(filtered_arguments)
   
    async def get_arguments(self, arguments: PTRPCTypedArrayParseFunctionMessage) -> PTRPCTypedArrayParseFunctionMessageResponse:
        argumentSplitArray = []
        for argValue in arguments.InputArray:
            argSplitResult = argValue.split(" ")
            for spaceSplitArg in argSplitResult:
                argumentSplitArray.append(spaceSplitArg)
        bof_arguments = []
        for argument in argumentSplitArray:
            if ":" not in argument:
                continue
            argType,value = argument.split(":",1)
            value = value.strip("\'").strip("\"")
            if argType == "":
                pass
            elif argType == "int16" or argType == "-s" or argType == "s":
                bof_arguments.append(["int16", int(value)])
            elif argType == "int32" or argType == "-i" or argType == "i":
                bof_arguments.append(["int32", int(value)])
            elif argType == "string" or argType == "-z" or argType == "z":
                bof_arguments.append(["string",value])
            elif argType == "wchar" or argType == "-Z" or argType == "Z":
                bof_arguments.append(["wchar",value])
            elif argType == "base64" or argType == "-b" or argType == "b":
                bof_arguments.append(["base64",value])
            elif argType == "bytes":
                # bytes type expects a hexadecimal string
                bof_arguments.append(["bytes",value])
            else:
                return PTRPCTypedArrayParseFunctionMessageResponse(Success=False,
                                                                   Error=f"Failed to parse argument: {argument}: Unknown value type.")

        argumentResponse = PTRPCTypedArrayParseFunctionMessageResponse(Success=True, TypedArray=bof_arguments)
        return argumentResponse

    async def get_files(self, callback: PTRPCDynamicQueryFunctionMessage) -> PTRPCDynamicQueryFunctionMessageResponse:
        response = PTRPCDynamicQueryFunctionMessageResponse()
        file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
            CallbackID=callback.Callback,
            LimitByCallback=False,
            IsDownloadFromAgent=False,
            IsScreenshot=False,
            IsPayload=False,
            Filename="",
        ))
        if file_resp.Success:
            file_names = []
            for f in file_resp.Files:
                if f.Filename not in file_names and f.Filename.endswith(".o"):
                    file_names.append(f.Filename)
            response.Success = True
            response.Choices = file_names
            return response
        else:
            await SendMythicRPCOperationEventLogCreate(MythicRPCOperationEventLogCreateMessage(
                CallbackId=callback.Callback,
                Message=f"Failed to get files: {file_resp.Error}",
                MessageLevel="warning"
            ))
            response.Error = f"Failed to get files: {file_resp.Error}"
            return response


class InlineExecuteCommand(CommandBase):
    cmd = "inline_execute"
    needs_admin = False
    help_cmd = "inline_execute -BOF [COFF.o] [-Arguments [optional arguments]]"
    description = "Execute a Beacon Object File in the current process thread. (e.g., inline_execute -BOF listmods.x64.o -Arguments int32:1234) \n [!!WARNING!! Incorrect argument types can crash the Agent process.]"
    version = 1
    author = "@c0rnbread (adapted for Angerona)"
    attackmapping = []
    argument_class = InlineExecuteArguments
    completion_functions = {"coff_completion_callback": coff_completion_callback}
    attributes = CommandAttributes(
        builtin=False,
        supported_os=[ SupportedOS.Windows ],
        suggested_command=False
    )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        
        try:
            groupName = taskData.args.get_parameter_group_name()
            
            # Check if bof_file is provided (from subtask or "New" group)
            if taskData.args.get_arg("bof_file"):
                file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
                    TaskID=taskData.Task.ID,
                    AgentFileID=taskData.args.get_arg("bof_file")
                ))
                if file_resp.Success:
                    if len(file_resp.Files) > 0:
                        # Set display parameters
                        if groupName != "New":  # Don't override if already set
                            response.DisplayParams = "-bof_file {} -bof_arguments {}".format(
                                file_resp.Files[0].Filename,
                                taskData.args.get_arg("bof_arguments")
                            )
                    else:
                        raise Exception("Failed to find that file")
                else:
                    raise Exception("Error from Mythic trying to get file: " + str(file_resp.Error))
            elif groupName == "Default" or taskData.args.get_arg("bof_name"):
                # We're trying to find an already existing file and use that
                file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
                    TaskID=taskData.Task.ID,
                    Filename=taskData.args.get_arg("bof_name"),
                    LimitByCallback=False,
                    MaxResults=1
                ))
                if file_resp.Success:
                    if len(file_resp.Files) > 0:
                        taskData.args.add_arg("bof_file", file_resp.Files[0].AgentFileId)
                        taskData.args.remove_arg("bof_name")    # Don't need this anymore
                        
                        # Set display parameters
                        response.DisplayParams = "-bof_file {} -bof_arguments {}".format(
                            file_resp.Files[0].Filename,
                            taskData.args.get_arg("bof_arguments")
                        )

                    elif len(file_resp.Files) == 0:
                        raise Exception("Failed to find the named file. Have you uploaded it before? Did it get deleted?")
                else:
                    raise Exception("Error from Mythic trying to search files:\n" + str(file_resp.Error))
            else:
                raise Exception("Either bof_name or bof_file must be provided")
        except Exception as e:
            raise Exception("Error from Mythic: " + str(e))
        
        return response

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about BOF execution risks.
        """
        message = "⚠️ OPSEC WARNING - BOF Execution Detected\n\n"
        message += "BOF (Beacon Object File) execution is monitored by:\n"
        message += "  • EDR/XDR solutions (memory allocation monitoring)\n"
        message += "  • API call monitoring (LoadLibraryA, GetProcAddress)\n"
        message += "  • Memory scanning (executable page allocation)\n"
        message += "  • Symbol resolution tracking\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • PAGE_EXECUTE_READWRITE memory allocation\n"
        message += "  • External symbol resolution (LoadLibraryA/GetProcAddress)\n"
        message += "  • COFF loader activity patterns\n"
        message += "  • Beacon API function calls\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • In-process execution reduces some detection vectors (no new process)\n"
        message += "  • Better OPSEC than `shell` (no cmd.exe spawn)\n"
        message += "  • Memory allocation may be scanned by EDR\n"
        message += "  • API calls may be logged by security tools\n\n"
        
        message += "⚠️ DETECTION RISK: MEDIUM\n"
        message += "Usually not immediately detected, but may be flagged by behavioral analysis.\n\n"
        message += "✅ This is a WARNING only - command will proceed.\n"
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=False,  # Medium risk, don't block
            OpsecPreBypassRole="operator",
            OpsecPreMessage=message
        )

    async def opsec_post(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPostTaskMessageResponse:
        """
        OPSEC check after creating the task.
        Warns about artifacts that will be created.
        """
        message = "⚠️ OPSEC POST-CHECK - BOF Execution Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Create: BOF execution in current process\n"
        message += "  • Memory Allocation: Executable memory pages (PAGE_EXECUTE_READWRITE)\n"
        message += "  • API Calls: LoadLibraryA, GetProcAddress, symbol resolution\n\n"
        
        message += "💡 ARTIFACT DETAILS:\n"
        message += "  • Process Create artifact is logged in Mythic\n"
        message += "  • Memory allocation may be detected by EDR memory scanning\n"
        message += "  • API calls may be logged by security tools\n\n"
        
        message += "✅ Review artifacts in the Artifacts tab after execution."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostMessage=message
        )

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        return resp

