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

import sys
import os

from libs.pentestlymodule import PentestlyModule
from libs.wmiexec import WMIEXEC

class Module(PentestlyModule):
    meta = {
        'name': 'Execute Command',
        'author': 'Cory Duplantis (@ctfhacker)',
        'description': 'Attempts to execute commands on systems using valid credentials.',
        'query': 'SELECT username,password,domain,host FROM pentestly_creds WHERE success="True" AND execute!="False"',
        'options': (
            ('username', '', False, 'Single user to execute commands. Leave blank for all valid credentials'),
            ('host', '', False, 'Single host to execute commands. Leave blank for all valid hosts'),
            ('domain', '', False, 'Single domain to execute commands. Leave blank for all valid domains'),
            ('share', 'C$', True, 'Share to execute commands against'),
            ('command', '', True, 'Command to execute'),
            ('get_results', '', True, 'Retrieve output from command'),
        ),
    }

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

    def module_run(self, creds):
        username = self.options['username']
        host = self.options['host']
        domain = self.options['domain']
        share = self.options['share']
        command = self.options['command']
        get_results = self.options['get_results']

        seen = set()
        # Reduce creds if options are given
        if username:
            creds = [cred for cred in creds if username == cred[0]]
        if host:
            creds = [cred for cred in creds if host == cred[3]]
        if domain:
            creds = [cred for cred in creds if domain == cred[2]]

        """
        # Not sure if we need this moving forward
        # Ensure we are only executing one command per box, and not per user
        pruned = []
        for username, password, domain, host in creds:
            if host not in seen:
                seen.add(host)
                pruned.append((username, password, domain, host))
        creds = pruned
        """

        for username, password, domain, host in creds:
            self.output("Executing command {}\{}:{}@{} - {}".format(domain, username, password, host, command))

            try:
                self.smb.hosts[host] = {'user':username, 'passwd':password, 'domain':domain, 'port': 445}

                result = self.wmi_execute(username, password, domain, host, command, get_results=get_results)
                self.output(result)
                self.pentestly_update(new_vals={'execute':'True'}, conditions={'username':username,
                                                                         'password':password,
                                                                         'host':host})
            except Exception as e:
                self.output("Failed to execute: {}\{}:{}@{}".format(domain, username, password, host))
                self.pentestly_update(new_vals={'execute':'False'}, conditions={'username':username,
                                                                          'password':password,
                                                                          'host':host})
                continue

    def module_post(self):
        self.stop_webserver()
