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

logging.basicConfig(level=logging.INFO)


class ExecuteAssemblyArguments(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=True, group_name="Default", ui_position=2,
                    ),
                    ParameterGroupInfo(
                        required=True, group_name="New Assembly", ui_position=2
                    ),
                ],
            ),
        ]
    
    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):
        if len(self.command_line) == 0:
            raise ValueError("Must supply arguments")
        raise ValueError("Must supply named arguments or use the modal")

class ExecuteAssemblyCommand(CommandBase):
    cmd = "execute_assembly"
    needs_admin = False
    help_cmd = "execute_assembly -File [Assembly Filename] [-Arguments [optional arguments]]"
    description = "Execute a .NET Assembly. Use an already uploaded assembly file or upload one with the command. (e.g., execute_assembly -File SharpUp.exe -Arguments \"audit\")"
    version = 1
    author = "@c0rnbread (adapted for Cazalla)"
    script_only = True
    attackmapping = []
    argument_class = ExecuteAssemblyArguments
    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()
            
            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 {}".format(
                    file_resp.Files[0].Filename,
                    taskData.args.get_arg("assembly_arguments")
                )
                
                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 {}".format(
                            file_resp.Files[0].Filename,
                            taskData.args.get_arg("assembly_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))

            # Get the file contents of the .NET assembly
            assembly_contents = await SendMythicRPCFileGetContent(
                MythicRPCFileGetContentMessage(AgentFileId=file_resp.Files[0].AgentFileId)
            )

            # Try to import donut - required for execute_assembly
            try:
                import donut
            except ImportError:
                raise Exception("The 'donut' Python module is not installed. Please install it with: pip install donut-shellcode")
            
            # Need a physical path for donut.create()
            fd, temppath = tempfile.mkstemp(suffix='.exe')
            logging.info(f"Writing Assembly Contents to temporary file \"{temppath}\"")
            with os.fdopen(fd, 'wb') as tmp:
                tmp.write(assembly_contents.Content)

            # Bypass=None, ExitOption=exit process
            assembly_shellcode = donut.create(file=temppath, params=taskData.args.get_arg("assembly_arguments"), bypass=1, exit_opt=2)
            # Clean up temp file
            os.remove(temppath)
            
            logging.info(f"Converted .NET into Shellcode {len(assembly_shellcode)} bytes")
            
            # .NET shellcode stub in Mythic
            shellcode_file_resp = await SendMythicRPCFileCreate(
                MythicRPCFileCreateMessage(TaskID=taskData.Task.ID, FileContents=assembly_shellcode, DeleteAfterFetch=True)
            )
            

            if shellcode_file_resp.Success:
                shellcode_file_uuid = shellcode_file_resp.AgentFileId
            else:
                raise Exception("Failed to register execute_assembly binary: " + shellcode_file_resp.Error)
            
            # Send subtask to inject shellcode
            # Note: This requires an inject_shellcode command to exist
            # For now, we'll just register the file and note that inject_shellcode needs to be implemented
            subtask = await SendMythicRPCTaskCreateSubtask(
                MythicRPCTaskCreateSubtaskMessage(
                    taskData.Task.ID,
                    CommandName="inject_shellcode",
                    SubtaskCallbackFunction="coff_completion_callback",
                    Params=json.dumps({
                        "shellcode_file": shellcode_file_uuid,
                        "method": "default"
                    }),
                    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

