import base64
from impacket.smbconnection import SMBConnection
from impacket.spnego import SPNEGO_NegTokenInit, TypesMech
from binascii import unhexlify
from ldap3 import ANONYMOUS
from lib.logger import logger
import ldap3
import sys
import ssl
import os



def get_dn(domain):
    components = domain.split('.')
    base = ''
    for comp in components:
        base += f',DC={comp}'
    
    return base[1:]

def get_machine_name(domain_controller, domain):
    if domain_controller is not None:
        s = SMBConnection(domain_controller, domain_controller)
    else:
        s = SMBConnection(domain, domain)
    try:
        s.login('', '')
    except Exception:
        if s.getServerName() == '':
            raise Exception('Error while anonymous logging into %s' % domain)
    else:
        s.logoff()
    return s.getServerName()

def init_ldap_connection(target, tls_version, domain, username, password, lmhash, nthash, domain_controller, kerberos, hashes, aesKey, use_channel_binding, signing=False):
    if signing:
        if use_channel_binding:
            logger.info("LDAP signing and channel binding are mutually exclusive.")
            sys.exit(1)
        if tls_version is not None:
            logger.info("LDAP signing is used over plain LDAP; LDAPS already provides integrity.")
            sys.exit(1)
        from lib.ldap_signed import build_signed_session
        logger.debug(f'[LDAP] Attempting signed bind | target={target} kerberos={bool(kerberos)} hashes={bool(hashes)}')
        ldap_session = build_signed_session(
            target=target,
            domain=domain,
            username=username,
            password=password,
            lmhash=lmhash,
            nthash=nthash,
            domain_controller=domain_controller,
            kerberos=kerberos,
            hashes=hashes,
            aes_key=aesKey,
        )
        return None, ldap_session

    user = '%s\\%s' % (domain, username)
    if tls_version is not None:
        use_ssl = True
        port = 636
        tls = ldap3.Tls(validate=ssl.CERT_NONE, version=tls_version)
    else:
        use_ssl = False
        port = 389
        tls = None
    ldap_server = ldap3.Server(target, get_info=ldap3.ALL, port=port, use_ssl=use_ssl, tls=tls)

    channel_binding = dict()
    if use_channel_binding:
        channel_binding = dict(channel_binding=ldap3.TLS_CHANNEL_BINDING)
    if kerberos:
        if channel_binding is True:
            logger.info("Kerberos auth + channel binding isn't supported yet.")
            sys.exit(1)
        logger.debug(f'[LDAP] Attempting Kerberos bind | target={target} port={port} ssl={use_ssl}')
        ldap_session = ldap3.Connection(ldap_server)
        ldap_session.bind()
        ldap3_kerberos_login(ldap_session, target, username, password, domain, lmhash, nthash, aesKey, kdcHost=domain_controller)
    elif hashes is not None:
        if lmhash == "":
            lmhash = "aad3b435b51404eeaad3b435b51404ee"
        logger.debug(f'[LDAP] Attempting NTLM bind (hash) | target={target} port={port} ssl={use_ssl} user={user}')
        try:
            ldap_session = ldap3.Connection(ldap_server, user=user, password=lmhash + ":" + nthash, authentication=ldap3.NTLM, auto_bind=True, **channel_binding)
        except Exception as e:
            logger.debug(f'[LDAP] Bind failed | {e}')
            raise
    elif username == '' and password == '':
        logger.debug(f'[LDAP] Attempting anonymous bind | target={target} port={port} ssl={use_ssl}')
        try:
            ldap_session = ldap3.Connection(ldap_server, authentication=ANONYMOUS, auto_bind=True, **channel_binding)
        except Exception as e:
            logger.debug(f'[LDAP] Bind failed | {e}')
            raise
    else:
        logger.debug(f'[LDAP] Attempting NTLM bind (password) | target={target} port={port} ssl={use_ssl} user={user}')
        try:
            ldap_session = ldap3.Connection(ldap_server, user=user, password=password, authentication=ldap3.NTLM, auto_bind=True, **channel_binding)
        except Exception as e:
            logger.debug(f'[LDAP] Bind failed | {e}')
            raise

    return ldap_server, ldap_session

def init_ldap_session(domain, username, password, lmhash, nthash, kerberos, domain_controller, ldaps, hashes, aesKey, channel_binding, signing=False):
    if signing and channel_binding:
        logger.info("LDAP signing and channel binding are mutually exclusive.")
        sys.exit(1)
    if signing and ldaps:
        logger.info("LDAP signing is used over plain LDAP; LDAPS already provides integrity.")
        sys.exit(1)

    if kerberos:
        #target = domain_controller
        netbiosname = get_machine_name(domain_controller, domain)
        #dontlookhereok
        target = netbiosname + "." + domain
    else:
        if domain_controller is not None:
            target = domain_controller
        else:
            target = domain

    if ldaps:
        try:
            return init_ldap_connection(target, ssl.PROTOCOL_TLSv1_2, domain, username, password, lmhash, nthash, domain_controller, kerberos, hashes, aesKey, channel_binding, signing=signing)
        except ldap3.core.exceptions.LDAPSocketOpenError:
            return init_ldap_connection(target, ssl.PROTOCOL_TLSv1, domain, username, password, lmhash, nthash, domain_controller, kerberos, hashes, aesKey, channel_binding, signing=signing)
    else:
        return init_ldap_connection(target, None, domain, username, password, lmhash, nthash, domain_controller, kerberos, hashes, aesKey, channel_binding, signing=signing)

def ldap3_kerberos_login(connection, target, user, password, domain='', lmhash='', nthash='', aesKey='', kdcHost=None, TGT=None, TGS=None, useCache=True, admin_service=False):
    from pyasn1.codec.ber import encoder, decoder
    from pyasn1.type.univ import noValue
    """
    logins into the target system explicitly using Kerberos. Hashes are used if RC4_HMAC is supported.
    :param string user: username
    :param string password: password for the user
    :param string domain: domain where the account is valid for (required)
    :param string lmhash: LMHASH used to authenticate using hashes (password is not used)
    :param string nthash: NTHASH used to authenticate using hashes (password is not used)
    :param string aesKey: aes256-cts-hmac-sha1-96 or aes128-cts-hmac-sha1-96 used for Kerberos authentication
    :param string kdcHost: hostname or IP Address for the KDC. If None, the domain will be used (it needs to resolve tho)
    :param struct TGT: If there's a TGT available, send the structure here and it will be used
    :param struct TGS: same for TGS. See smb3.py for the format
    :param bool useCache: whether or not we should use the ccache for credentials lookup. If TGT or TGS are specified this is False
    :return: True, raises an Exception if error.
    """

    # Normalize optional inputs that may arrive as None from CLI options.
    user = '' if user is None else user
    password = '' if password is None else password
    domain = '' if domain is None else domain
    lmhash = '' if lmhash is None else lmhash
    nthash = '' if nthash is None else nthash
    aesKey = '' if aesKey is None else aesKey

    if lmhash != '' or nthash != '':
        if len(lmhash) % 2:
            lmhash = '0' + lmhash
        if len(nthash) % 2:
            nthash = '0' + nthash
        try:  # just in case they were converted already
            lmhash = unhexlify(lmhash)
            nthash = unhexlify(nthash)
        except TypeError:
            pass

    # Importing down here so pyasn1 is not required if kerberos is not used.
    from impacket.krb5.ccache import CCache
    from impacket.krb5.asn1 import AP_REQ, Authenticator, TGS_REP, seq_set
    from impacket.krb5.kerberosv5 import getKerberosTGT, getKerberosTGS
    from impacket.krb5 import constants
    from impacket.krb5.types import Principal, KerberosTime, Ticket
    import datetime

    if TGT is not None or TGS is not None:
        useCache = False

    if useCache:
        try:
            ccache = CCache.loadFile(os.getenv('KRB5CCNAME'))
            if ccache is None:
                raise Exception('CCache file is not found. Skipping...')
        except Exception as e:
            # No cache present
            logger.warning(e)
            pass
        else:
            # retrieve domain information from CCache file if needed
            if domain == '':
                domain = ccache.principal.realm['data'].decode('utf-8')


            logger.debug('Using Kerberos Cache: %s' % os.getenv('KRB5CCNAME'))

            # Build the service principal name
            if admin_service:
                service_principal = 'HTTP/%s@%s' % (target.upper(), domain.upper())
            else:
                service_principal = 'ldap/%s@%s' % (target.upper(), domain.upper())

            # First, try to get the TGS (service ticket) from cache
            creds = ccache.getCredential(service_principal)
            if creds is None:
                # No TGS in cache, try to get TGT instead
                tgt_principal = 'krbtgt/%s@%s' % (domain.upper(), domain.upper())
                creds = ccache.getCredential(tgt_principal)
                if creds is not None:
                    TGT = creds.toTGT()
                    logger.debug('Using TGT from cache')
                else:
                    logger.debug('No valid credentials found in cache')
            else:
                # Found TGS in cache - we can reuse it!
                TGS = creds.toTGS(service_principal)
                logger.debug('Using TGS from cache for principal: %s' % service_principal)

            # retrieve user information from CCache file if needed
            if user == '' and creds is not None:
                user = creds['client'].prettyPrint().split(b'@')[0].decode('utf-8')
                logger.debug('Username retrieved from CCache: %s' % user)
            elif user == '' and len(ccache.principal.components) > 0:
                user = ccache.principal.components[0]['data'].decode('utf-8')
                logger.debug('Username retrieved from CCache: %s' % user)

    # First of all, we need to get a TGT for the user
    userName = Principal(user, type=constants.PrincipalNameType.NT_PRINCIPAL.value)
    kdc_str = kdcHost if kdcHost else domain
    logger.debug(f'[KRB5] principal={user}@{domain.upper()} target={target} kdc={kdc_str}')
    if TGT is None:
        if TGS is None:
            if password == '' and lmhash == '' and nthash == '' and aesKey == '':
                raise ValueError(
                    'Kerberos authentication requires one of: a valid Kerberos ccache, '
                    'password, NTLM hash, or AES key.'
                )
            auth_method = 'aes' if aesKey else ('nthash' if nthash else 'password')
            logger.debug(f'[KRB5] Requesting TGT (AS-REQ) | kdc={kdc_str} auth={auth_method}')
            try:
                tgt, cipher, oldSessionKey, sessionKey = getKerberosTGT(userName, password, domain, lmhash, nthash, aesKey, kdcHost)
            except Exception as e:
                logger.debug(f'[KRB5] TGT request failed | {e}')
                raise
            logger.debug(f'[KRB5] TGT obtained | enctype={cipher.enctype}')
    else:
        tgt = TGT['KDC_REP']
        cipher = TGT['cipher']
        sessionKey = TGT['sessionKey']
        logger.debug(f'[KRB5] Using provided TGT | enctype={cipher.enctype}')

    if TGS is None:
        # No TGS in cache, need to request one from KDC
        if admin_service:
            serverName = Principal('HTTP/%s' % target, type=constants.PrincipalNameType.NT_SRV_INST.value)
        else:
            serverName = Principal('ldap/%s' % target, type=constants.PrincipalNameType.NT_SRV_INST.value)
        logger.debug(f'[KRB5] Requesting TGS (TGS-REQ) | service={serverName} kdc={kdc_str}')
        try:
            tgs, cipher, oldSessionKey, sessionKey = getKerberosTGS(serverName, domain, kdcHost, tgt, cipher, sessionKey)
        except Exception as e:
            logger.debug(f'[KRB5] TGS request failed | {e}')
            raise
        logger.debug(f'[KRB5] TGS obtained | enctype={cipher.enctype}')
    else:
        # Reuse cached TGS - no KDC call needed!
        logger.debug('Reusing cached TGS - skipping KDC request')
        tgs = TGS['KDC_REP']
        cipher = TGS['cipher']
        sessionKey = TGS['sessionKey']
        logger.debug(f'[KRB5] Using cached TGS | enctype={cipher.enctype}')

        # Let's build a NegTokenInit with a Kerberos REQ_AP

    blob = SPNEGO_NegTokenInit()

    # Kerberos
    blob['MechTypes'] = [TypesMech['MS KRB5 - Microsoft Kerberos 5']]

    # Let's extract the ticket from the TGS
    tgs = decoder.decode(tgs, asn1Spec=TGS_REP())[0]
    ticket = Ticket()
    ticket.from_asn1(tgs['ticket'])

    # Now let's build the AP_REQ
    apReq = AP_REQ()
    apReq['pvno'] = 5
    apReq['msg-type'] = int(constants.ApplicationTagNumbers.AP_REQ.value)

    opts = []
    apReq['ap-options'] = constants.encodeFlags(opts)
    seq_set(apReq, 'ticket', ticket.to_asn1)

    authenticator = Authenticator()
    authenticator['authenticator-vno'] = 5
    authenticator['crealm'] = domain
    seq_set(authenticator, 'cname', userName.components_to_asn1)
    now = datetime.datetime.utcnow()

    authenticator['cusec'] = now.microsecond
    authenticator['ctime'] = KerberosTime.to_asn1(now)

    encodedAuthenticator = encoder.encode(authenticator)

    # Key Usage 11
    # AP-REQ Authenticator (includes application authenticator
    # subkey), encrypted with the application session key
    # (Section 5.5.1)
    encryptedEncodedAuthenticator = cipher.encrypt(sessionKey, 11, encodedAuthenticator, None)

    apReq['authenticator'] = noValue
    apReq['authenticator']['etype'] = cipher.enctype
    apReq['authenticator']['cipher'] = encryptedEncodedAuthenticator

    blob['MechToken'] = encoder.encode(apReq)


    if admin_service:
        token = base64.b64encode(blob.getData()).decode('ascii')
        return f'Negotiate {token}'
    else:
        request = ldap3.operation.bind.bind_operation(connection.version, ldap3.SASL, user, None, 'GSS-SPNEGO',
                                                  blob.getData())

        # Done with the Kerberos saga, now let's get into LDAP
        logger.debug(f'[KRB5] Sending SASL GSS-SPNEGO bindRequest | user={user} server={connection.server}')
        if connection.closed:  # try to open connection if closed
            connection.open(read_server_info=False)

        connection.sasl_in_progress = True
        response = connection.post_send_single_response(connection.send('bindRequest', request, None))
        connection.sasl_in_progress = False
        if response[0]['result'] != 0:
            result_code = response[0].get('result', 'unknown')
            description = response[0].get('description', '')
            message = response[0].get('message', '')
            logger.debug(f'[LDAP] Kerberos bind failed | result={result_code} description={description!r} message={message!r}')
            raise Exception(response)

        connection.bound = True

    return True
