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

logging.basicConfig(level=logging.INFO)

# Wrapper function for Inline-EA
# BoF Credits: https://github.com/EricEsquivel/Inline-EA

# Completion callback for BOF execution
# This captures output from the subtask and displays it in the main command
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
        # Add a small delay to ensure the response is available
        import asyncio
        await asyncio.sleep(1.0)  # Wait 1 second for response to be processed
        
        responses = await SendMythicRPCResponseSearch(MythicRPCResponseSearchMessage(TaskID=subtask_id))
        
        logging.info(f"[coff_completion_callback] Searching responses for subtask {subtask_id}, Success={responses.Success}")
        logging.info(f"[coff_completion_callback] Response object type: {type(responses)}")
        logging.info(f"[coff_completion_callback] Response object attributes: {dir(responses)}")
        
        if not responses.Success:
            error_detail = getattr(responses, 'Error', 'Unknown error') if hasattr(responses, 'Error') else 'Unknown error'
            logging.warning(f"[coff_completion_callback] Failed to search responses: {error_detail}")
            # Don't return error, just log it and continue
        
        # Process responses if they exist
        responses_list = None
        if responses.Success:
            if hasattr(responses, 'Responses'):
                responses_list = responses.Responses
                logging.info(f"[coff_completion_callback] Found responses.Responses: {len(responses_list) if responses_list else 0} items")
            elif hasattr(responses, 'responses'):
                responses_list = responses.responses
                logging.info(f"[coff_completion_callback] Found responses.responses: {len(responses_list) if responses_list else 0} items")
            elif hasattr(responses, 'Response'):
                responses_list = [responses.Response] if responses.Response else []
                logging.info(f"[coff_completion_callback] Found responses.Response: {len(responses_list) if responses_list else 0} items")
            else:
                logging.warning(f"[coff_completion_callback] No known response attribute found in response object")
        
        if responses_list:
            logging.info(f"[coff_completion_callback] Processing {len(responses_list)} response items")
            for idx, output in enumerate(responses_list):
                logging.info(f"[coff_completion_callback] Processing response item {idx}, type: {type(output)}")
                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 hasattr(output, 'output'):
                    out += str(output.output)
                elif isinstance(output, str):
                    out += output
                elif isinstance(output, bytes):
                    out += output.decode('utf-8', errors='ignore')
                else:
                    logging.info(f"[coff_completion_callback] Converting output to string, type: {type(output)}")
                    out += str(output)
        
        # If no output was found, indicate successful execution
        if not out:
            logging.warning(f"[coff_completion_callback] No output found for subtask {subtask_id}")
            out = "[Assembly 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
        return PTTaskCompletionFunctionMessageResponse(Success=False, TaskStatus="error", Completed=True, Error=str(e))

class InlineExecuteAssemblyArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="assembly_name",
                cli_name="Assembly",
                display_name="Assembly",
                type=ParameterType.ChooseOne,
                dynamic_query_function=self.get_files,
                description="Already existing .NET assembly to execute (e.g. SharpUp.exe)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=True,
                        group_name="Default",
                        ui_position=1
                    )
                ]),
            CommandParameter(
                name="assembly_file",
                display_name="New Assembly",
                type=ParameterType.File,
                description="A new .NET assembly to execute. After uploading once, you can just supply the -Assembly parameter",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=True, 
                        group_name="New Assembly", 
                        ui_position=1,
                    )
                ]
            ),
            CommandParameter(
                name="assembly_arguments",
                cli_name="Arguments",
                display_name="Arguments",
                type=ParameterType.String,
                description="Arguments to pass to the assembly.",
                default_value="",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False, group_name="Default", ui_position=2,
                    ),
                    ParameterGroupInfo(
                        required=False, group_name="New Assembly", ui_position=2
                    ),
                ],
            ),
            CommandParameter(
                name="patch_exit",
                cli_name="-patchexit",
                display_name="patchexit",
                type=ParameterType.Boolean,
                description="Patches System.Environment.Exit to prevent Beacon process from exiting",
                default_value=False,
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False, group_name="Default", ui_position=3,
                    ),
                    ParameterGroupInfo(
                        required=False, group_name="New Assembly", ui_position=3
                    ),
                ],
            ),
            CommandParameter(
                name="amsi",
                cli_name="-amsi",
                display_name="amsi",
                type=ParameterType.Boolean,
                description="Bypass AMSI by patching clr.dll instead of amsi.dll to avoid common detections",
                default_value=False,
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False, group_name="Default", ui_position=4,
                    ),
                    ParameterGroupInfo(
                        required=False, group_name="New Assembly", ui_position=4
                    ),
                ],
            ),
            CommandParameter(
                name="etw",
                cli_name="-etw",
                display_name="etw",
                type=ParameterType.Boolean,
                description="Bypass ETW by EAT Hooking advapi32.dll!EventWrite to point to a function that returns right away",
                default_value=False,
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False, group_name="Default", ui_position=5,
                    ),
                    ParameterGroupInfo(
                        required=False, group_name="New Assembly", ui_position=5
                    ),
                ],
            ),
        ]
    
    async def get_files(self, callback: PTRPCDynamicQueryFunctionMessage) -> PTRPCDynamicQueryFunctionMessageResponse:
        response = PTRPCDynamicQueryFunctionMessageResponse()
        file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
            CallbackID=callback.Callback,
            LimitByCallback=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(".exe"):
                    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

    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):
        # This is called when using the modal UI or when called as a subtask
        # Just load the arguments directly
        self.load_args_from_dictionary(dictionary_arguments)

class InlineExecuteAssemblyCommand(CommandBase):
    cmd = "inline_execute_assembly"
    needs_admin = False
    help_cmd = "inline_execute_assembly -Assembly [file] [-Arguments [assembly args] [--patchexit] [--amsi] [--etw]]"
    description = "Execute a .NET Assembly in the current process using @EricEsquivel's BOF \"Inline-EA\" (e.g., inline_execute_assembly -Assembly SharpUp.exe -Arguments \"audit\" --patchexit --amsi --etw)"
    version = 1
    author = "@c0rnbread (adapted for Angerona)"
    script_only = True
    attackmapping = []
    argument_class = InlineExecuteAssemblyArguments
    completion_functions = {"coff_completion_callback": coff_completion_callback}
    attributes = CommandAttributes(
        dependencies=["inline_execute"],
        alias=True,
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about .NET assembly execution risks, especially with bypass flags.
        """
        # Check if bypass flags are enabled
        amsi_enabled = taskData.args.get_arg("amsi") or False
        etw_enabled = taskData.args.get_arg("etw") or False
        patch_exit_enabled = taskData.args.get_arg("patch_exit") or False
        
        blocked = False
        message = "⚠️ OPSEC WARNING - .NET Assembly In-Process Execution Detected\n\n"
        
        message += ".NET assembly execution is monitored by:\n"
        message += "  • EDR/XDR solutions (CLR loading detection)\n"
        message += "  • AMSI (Anti-Malware Scan Interface)\n"
        message += "  • ETW (Event Tracing for Windows)\n"
        message += "  • Memory scanning (executable page allocation)\n"
        message += "  • Behavioral analysis engines\n\n"
        
        if amsi_enabled or etw_enabled:
            message += "🚨 HIGH RISK DETECTED - Bypass Flags Enabled\n\n"
            message += "⚠️ WARNING: You have enabled bypass flags that significantly increase detection risk:\n\n"
            
            if amsi_enabled:
                message += "  • --amsi: AMSI bypass by patching clr.dll\n"
                message += "    - Patching clr.dll may trigger behavioral detections\n"
                message += "    - More evasive than patching amsi.dll, but still detectable\n"
                message += "    - Advanced EDR solutions monitor CLR modifications\n\n"
            
            if etw_enabled:
                message += "  • --etw: ETW bypass via EAT hooking\n"
                message += "    - EAT hooking advapi32.dll!EventWrite is monitored\n"
                message += "    - Advanced EDR solutions detect export table modifications\n"
                message += "    - May trigger immediate alerts\n\n"
            
            message += "🔍 BYPASS DETECTION MECHANISMS:\n"
            message += "  • CLR DLL patching detection\n"
            message += "  • Export Address Table (EAT) hooking detection\n"
            message += "  • Behavioral analysis of bypass techniques\n"
            message += "  • Memory integrity checks\n\n"
            
            message += "⚠️ DETECTION RISK: HIGH (with bypass flags)\n"
            message += "Bypass techniques are heavily monitored and may trigger immediate alerts.\n\n"
            
            # Block if both bypasses are enabled
            if amsi_enabled and etw_enabled:
                blocked = True
                message += "🚨 BLOCKED: Both AMSI and ETW bypasses enabled.\n"
                message += "This combination has very high detection risk.\n"
                message += "Consider using only one bypass or none if possible.\n\n"
        else:
            message += "🔍 DETECTION MECHANISMS:\n"
            message += "  • CLR loading into memory\n"
            message += "  • .NET assembly reflection/loading\n"
            message += "  • Memory allocation (executable pages)\n"
            message += "  • Assembly execution patterns\n\n"
            
            message += "💡 CONSIDERATIONS:\n"
            message += "  • In-process execution reduces detection vectors (no process injection)\n"
            message += "  • Better OPSEC than process spawning (no new process creation)\n"
            message += "  • CLR loading may be detected by EDR\n"
            message += "  • Memory allocation may be scanned\n\n"
            
            message += "⚠️ DETECTION RISK: MEDIUM (without bypass flags)\n"
            message += "Usually not immediately detected, but may be flagged by behavioral analysis.\n\n"
        
        if patch_exit_enabled:
            message += "  • --patchexit: Prevents assembly from exiting (recommended)\n\n"
        
        message += "💡 ALTERNATIVES:\n"
        message += "  • Use `inline_execute` for BOFs (lower detection risk)\n"
        message += "  • Consider executing assemblies without bypass flags first\n\n"
        
        if not blocked:
            message += "✅ This is a WARNING only - command will proceed.\n"
        else:
            message += "❌ This command is BLOCKED due to high detection risk.\n"
        
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=blocked,
            OpsecPreBypassRole="other_operator" if blocked else "operator",
            OpsecPreMessage=message
        )

    async def opsec_post(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPostTaskMessageResponse:
        """
        OPSEC check after creating the task.
        Warns about artifacts that will be created.
        """
        amsi_enabled = taskData.args.get_arg("amsi") or False
        etw_enabled = taskData.args.get_arg("etw") or False
        
        message = "⚠️ OPSEC POST-CHECK - .NET Assembly Execution Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Create: .NET assembly execution in-process\n"
        message += "  • CLR Loading: .NET Common Language Runtime loaded into memory\n"
        message += "  • Memory Allocation: Executable memory pages for assembly\n"
        
        if amsi_enabled:
            message += "  • AMSI Bypass: clr.dll patching (high detection risk)\n"
        if etw_enabled:
            message += "  • ETW Bypass: EAT hooking advapi32.dll!EventWrite (high detection risk)\n"
        
        message += "\n💡 ARTIFACT DETAILS:\n"
        message += "  • Process Create artifact is logged in Mythic\n"
        message += "  • CLR loading may be detected by EDR\n"
        message += "  • Memory allocation may be scanned by EDR memory protection\n"
        
        if amsi_enabled or etw_enabled:
            message += "  • Bypass techniques may trigger immediate behavioral alerts\n"
        
        message += "\n✅ Review artifacts in the Artifacts tab after execution."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostMessage=message
        )

    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()
            
            if groupName == "New Assembly":
                file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
                    TaskID=taskData.Task.ID,
                    AgentFileID=taskData.args.get_arg("assembly_file")
                ))
                if file_resp.Success:
                    if len(file_resp.Files) > 0:
                        pass
                    else:
                        raise Exception("Failed to find that file")
                else:
                    raise Exception("Error from Mythic trying to get file: " + str(file_resp.Error))
                
                # Set display parameters
                response.DisplayParams = "-Assembly {} -Arguments {} --patchexit {} --amsi {} --etw {}".format(
                    file_resp.Files[0].Filename,
                    taskData.args.get_arg("assembly_arguments"),
                    taskData.args.get_arg("patch_exit"),
                    taskData.args.get_arg("amsi"),
                    taskData.args.get_arg("etw")
                )
                
                taskData.args.add_arg("assembly_name", file_resp.Files[0].Filename)
                taskData.args.remove_arg("assembly_file")
            
            elif groupName == "Default":
                # 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("assembly_name"),
                    LimitByCallback=False,                                
                    MaxResults=1
                ))
                if file_resp.Success:
                    if len(file_resp.Files) > 0:
                        logging.info(f"Found existing Assembly with File ID : {file_resp.Files[0].AgentFileId}")

                        taskData.args.remove_arg("assembly_name")    # Don't need this anymore
                        
                        # Set display parameters
                        response.DisplayParams = "-Assembly {} -Arguments {} --patchexit {} --amsi {} --etw {}".format(
                            file_resp.Files[0].Filename,
                            taskData.args.get_arg("assembly_arguments"),
                            taskData.args.get_arg("patch_exit"),
                            taskData.args.get_arg("amsi"),
                            taskData.args.get_arg("etw")
                        )

                    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))

            # Get the file contents of the .NET assembly ( base64 ( assembly bytes ) )
            assembly_contents = await SendMythicRPCFileGetContent(
                MythicRPCFileGetContentMessage(AgentFileId=file_resp.Files[0].AgentFileId)
            )
            
            # Arguments depend on the BOF
            # Note: Using bof_name instead of bof_file to match Xenon's implementation
            file_name = "inline-ea.x64.o"
            arguments = [
                [
                    "bytes", 
                    assembly_contents.Content.hex()                             # Raw bytes of Assembly
                ],
                [
                    "int32",               
                    len(assembly_contents.Content)                             # Assembly length
                ],
                [
                    "wchar", 
                    taskData.args.get_arg("assembly_arguments") or ""                 # Assembly argument string
                ],
                [
                    "int32",
                    1 if taskData.args.get_arg("patch_exit") else 0  # BOOL: convert to 1 or 0
                ],
                [
                    "int32",
                    1 if taskData.args.get_arg("amsi") else 0  # BOOL: convert to 1 or 0
                ],
                [
                    "int32",
                    1 if taskData.args.get_arg("etw") else 0  # BOOL: convert to 1 or 0
                ]
                
            ]
            
            # Find the BOF file first
            bof_file_resp = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
                TaskID=taskData.Task.ID,
                Filename=file_name,
                LimitByCallback=False,
                MaxResults=1
            ))
            
            if not bof_file_resp.Success or len(bof_file_resp.Files) == 0:
                raise Exception(f"Failed to find BOF file '{file_name}': {bof_file_resp.Error if bof_file_resp.Success else 'File not found'}")
            
            # Create subtask to execute inline_execute with the BOF
            # Use coff_completion_callback to capture output and display it in the main command
            subtask = await SendMythicRPCTaskCreateSubtask(
                MythicRPCTaskCreateSubtaskMessage(
                    taskData.Task.ID,
                    CommandName="inline_execute",
                    SubtaskCallbackFunction="coff_completion_callback",
                    Params=json.dumps({
                        "bof_file": bof_file_resp.Files[0].AgentFileId,
                        "bof_arguments": arguments
                    }),
                    Token=taskData.Task.TokenID,
                )
            )
            
            return response

        except Exception as e:
            raise Exception("Error from Mythic: " + str(e))
        
    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        return resp

