from django_auth_ldap3.conf import settings

from django.contrib.auth.models import Group
from django.contrib.auth import get_user_model
from ldap3.core.exceptions import LDAPSocketOpenError
from ldap3.core.exceptions import LDAPAttributeError
import hashlib
import ldap3
import logging
import ssl
import time
import random

User = get_user_model()

logger = logging.getLogger('django_auth_ldap3')

class ADConnection(object):
    def __init__(self, connection, legacy_bind):
        self.connection = connection
        self.legacy_bind = legacy_bind

class LDAPUser(object):
    """
    A class representing an LDAP user returned from the directory.
    """

    connection = None
    _attrib_keys = [settings.UID_ATTRIB, 'cn', 'givenName', 'sn', 'mail']

    def __init__(self, connection, attributes):
        self.connection = connection
        for k, v in attributes.items():
            # Flatten any lists into their first element
            if type(v) == list and len(v) >= 1:
                v = v[0]
            setattr(self, k, v)

        # Set any missing attributes
        for k in self._attrib_keys:
            if not hasattr(self, k):
                setattr(self, k, '')

    def __str__(self):
        return self.username

    @property
    def username(self):
        return getattr(self, settings.UID_ATTRIB)

class LDAPBackend(object):
    """
    An authentication backend for LDAP directories.
    """

    backend = None

    def __init__(self):
        tls_config = None

        if settings.TLS:
            tls_opts = {
                'validate': ssl.CERT_REQUIRED if settings.TLS_VALIDATE else ssl.CERT_NONE
            }

            if settings.TLS_CA_CERTS:
                tls_opts['ca_certs_file'] = settings.TLS_CA_CERTS

            if settings.TLS_PRIVATE_KEY:
                tls_opts['local_private_key_file'] = settings.TLS_PRIVATE_KEY

            if settings.TLS_LOCAL_CERT:
                tls_opts['local_certificate_file'] = settings.TLS_LOCAL_CERT

            tls_config = ldap3.Tls(**tls_opts)

        self.backend = ldap3.Server(settings.URI, use_ssl=settings.TLS, tls=tls_config)

    def __del__(self):
        # TODO: disconnect?
        pass

    def authenticate(self, request=None, username=None, password=None):
        """
        Required for Django auth. Authenticate the user against the LDAP
        backend and populate the local User model if this is the first
        time the user has authenticated.
        """

        # Authenticate against the LDAP backend and return an LDAPUser.
        ldap_user = self.retrieve_ldap_user(username, password)
        if ldap_user is None:
            logger.debug('Authentication failed for {}'.format(ldap_user))
            return None

        # If we get here, authentication is successful and we have an LDAPUser
        # instance populated with the user's attributes. We still need to check
        # group membership and populate a local User model.

        # Check LDAP group membership before creating a local user. The default
        # is '*' so any user can log in.
        if not self.check_group_membership(ldap_user, settings.LOGIN_GROUP):
            logger.debug('Failed group membership test: {} !memberOf {}'.format(ldap_user, settings.LOGIN_GROUP))
            return None

        # Check if this user is part of the admin group.
        admin = False
        if settings.ADMIN_GROUP:
            admin = self.check_group_membership(ldap_user, settings.ADMIN_GROUP)

        # Get or create the User object in Django's auth, populating it with
        # fields from the LDAPUser. Note we set the password to a random hash
        # as authentication should never occur directly off this user.
        django_user = User.objects.filter(username__iexact=username)
        if not django_user:
            # Create new user. We use `ldap_user.username` here as it is the
            # case-sensitive version
            django_user = User(username=ldap_user.username,
                    password=hashlib.sha1().hexdigest(),
                    first_name=ldap_user.givenName,
                    last_name=ldap_user.sn,
                    email=ldap_user.mail,
                    is_superuser=False,
                    is_staff=admin,
                    is_active=True
            )
            django_user.set_unusable_password() #bdm 11-23-20 - added to prevent password resets. -assumes current users have already set this property (since this is on creation only)
            django_user.save()
        else:
            # If the user wasn't created, update its fields from the directory.
            django_user = django_user[0]
            django_user.first_name = ldap_user.givenName
            django_user.last_name = ldap_user.sn
            django_user.email = ldap_user.mail
            django_user.is_staff = admin
            django_user.save()

        self.update_group_membership(ldap_user, django_user)

        return django_user

    def get_user(self, user_id):
        """
        Required for Django auth.
        """
        try:
            return User.objects.get(pk=user_id)
        except User.DoesNotExist:
            return None

    def check_group_membership(self, ldap_user, group_dn):
        """
        Check the LDAP user to see if it is a member of the given group.

        This is straightforward with OpenLDAP but tricky with AD as due to
        the weird way AD handles "primary" group membership, we must test for
        a separate attribute as well as the usual 'memberof' as the primary
        group is not returned with that filter.
        """

        # Don't bother search directory if '*' was given: this denotes any group
        # so pass the test immediately.
        if group_dn == '*':
            return True

        # Hack for AD: fetch the group's attributes and check for the
        # primaryGroupToken. This will return 0 results in OpenLDAP and hence
        # be ignored.
        pgt = None
        group_attribs = self.search_ldap(ldap_user.connection, '(distinguishedName={})' \
                .format(group_dn), attributes=['primaryGroupToken'])
        if group_attribs:
            pgt = group_attribs.get('primaryGroupToken', None)
            if type(pgt) == list:
                pgt = pgt[0]

        # Now perform our group membership test. If the primary group token is not-None,
        # then we wrap the filter in an OR and test for that too.
        search_filter = '(&(objectClass=user)({}={})(memberof={}))'.format(
                settings.UID_ATTRIB, str(ldap_user), group_dn)
        if pgt:
            search_filter = '(|{}(&(cn={})(primaryGroupID={})))'.format(search_filter, ldap_user.cn, pgt)

        # Return True if user is a member of group
        r = self.search_ldap(ldap_user.connection, search_filter)
        return r is not None

    def retrieve_ldap_user(self, username, password):
        """
        First attempt an AD bind, then attempt legacy LDAP bind
        """
        ldap_object = None
        try:
            ldap_object = self.bind_ad_user(username, password)
        except:
            logger.error("AD Bind function error reached wrapper for user %s" % username)
            for i in (0, 4):
                self.delay_return_none()
                return(self.delay_return_none())
        if not ldap_object:
            if settings.LEGACY_BIND == True:
                self.backend = ldap3.Server(settings.LEGACY_URI, use_ssl=settings.TLS, tls=None)
                ldap_object = self.bind_ldap_user(username, password)
        return(ldap_object)

    def search_ldap(self, connection, ldap_filter, **kwargs):
        """
        Searches the LDAP directory against the given LDAP filter in the form
        of '(attr=val)' e.g. '(uid=test)'.

        Any keyword arguments will be passed directly to the underlying search
        method.

        A dictionary of attributes will be returned.
        """
        for dn in settings.LEGACY_DN:
            connection.search(dn, ldap_filter, **kwargs)
            entry = None
            for d in connection.response:
                if d['type'] == 'searchResEntry':
                    entry = d['attributes']
                    entry['dn'] = d['dn']
                    break
            
            if entry:
                break
            
        return entry

    def delay_return_none(self):
        time.sleep(((random.randrange(100, 500))*.001))
        return(None)
        
    def get_ad_connection(self, username = None, password = None):
        if username == None:
            username = settings.BIND_USER
        if password == None:
            password = settings.BIND_PASS
        URI = settings.URI
        USERNAME_PREFIX = settings.USERNAME_PREFIX
        USERNAME_SUFFIX = settings.USERNAME_SUFFIX
        if USERNAME_SUFFIX:
            if (not (USERNAME_SUFFIX.startswith("@"))):
                USERNAME_SUFFIX = "@" + USERNAME_SUFFIX

        USERNAME = username
        PASSWORD = password

        ### AD Bind disallowed username check
        allowed_user = True
        disallowed_substrings = settings.DISALLOWED_USER_SUBSTRINGS
        disallowed_users = settings.DISALLOWED_USERS
        for ss in disallowed_substrings:
            if ss.lower().strip() in USERNAME.lower().strip():
                allowed_user = False
                break
        for du in disallowed_users:
            if du.lower().strip() == USERNAME.lower().strip():
                allowed_user = False
                break
                
        allowed_specials = settings.ALLOWED_SPECIALS
        for char in USERNAME:
            if not char.isalnum():
                if not char in allowed_specials:
                    allowed_user = False
                    break
                
        if not allowed_user:
            for i in range(0,4):
                self.delay_return_none()
            return(self.delay_return_none())
        ### end disallowed user check        

        BASE_DN = settings.BASE_DN
        server = ldap3.Server(URI, port=636, use_ssl=True)
        
        #try the new way of domain binding first, then fall back to legacy
        connection = None
        legacy_bind = False
        if USERNAME_SUFFIX:
            try:
                connection = ldap3.Connection(server, user='{}{}'.format(USERNAME, USERNAME_SUFFIX), password=PASSWORD, lazy=False, auto_bind=True, client_strategy=ldap3.SYNC, read_only=True)        
                connection.open()
                if not connection.bind():
                    connection = None
            except ldap3.core.exceptions.LDAPBindError:
                connection = None
        #legacy AD domain bind
        if not connection:
            try:
                connection = ldap3.Connection(server, user='{}\\{}'.format(USERNAME_PREFIX, USERNAME), password=PASSWORD, lazy=False, auto_bind=True, client_strategy=ldap3.SYNC, read_only=True)
                connection.open()
                if not connection.bind():
                    connection = None
                else:
                    legacy_bind = True
            except ldap3.core.exceptions.LDAPBindError:
                connection = None    
        return_connection = ADConnection(connection, legacy_bind)
        if return_connection.connection == None:
            return_connection = None
        return(return_connection)
        
        
    def get_ad_attribs(self, connection, username, attribs, legacy_attribs = None):
        connection = connection.connection
        #if attribs is string instead of list, do single mode
        single_mode = False
        if type(attribs) == type(str()):
            attribs = [attribs]
            single_mode = True
            
        #if we're doing legacy attributes, map them here in format
        #ldap_attrib : ad_attrib
        ad_attribs_to_get = []
        attrib_map = {
            'cn' : 'sAMAccountName',
            'dn' : 'distinguishedName',
            'loginTime' : 'lastLogon'
        }
        for attr in attribs:
            new_attr_out = None
            if legacy_attribs == True:
                for key in attrib_map.keys():
                    if key.lower() == attr.lower():
                        new_attr_out = attrib_map[key]
            if not new_attr_out:
                new_attr_out = attr
            ad_attribs_to_get.append(new_attr_out)
            
        #search for the user by samaccountname
        BASE_DN = settings.BASE_DN
        try:
            connection.search(BASE_DN, '(&(objectClass=person)(SamAccountName=%s))' % username, attributes=ad_attribs_to_get)
        except LDAPAttributeError:
            logger.error('Invalid attribute or attributes passed to AD: %s' % str(ad_attribs_to_get))
            return(self.delay_return_none())
        entry = None
        for res in connection.response:
            if res['type'] == 'searchResEntry':
                entry = res
                break

        if not entry:
            logger.error('LDAP search error: no results for %s' % username)
            return(self.delay_return_none())
            
        else:
            ad_attribs = entry['attributes']
            if not legacy_attribs:
                return(ad_attribs)
            else:
                try:
                    legacy_ldap_attribs = {}
                    for ad_key in ad_attribs.keys():
                        legacy_ldap_key = None
                        for legacy_attrib in attrib_map.keys():
                            if attrib_map[legacy_attrib] == ad_key:
                                legacy_ldap_key = legacy_attrib
                            if legacy_ldap_key:
                                break
                        if not legacy_ldap_key:
                            legacy_ldap_key = ad_key
                        legacy_ldap_attribs[legacy_ldap_key] = ad_attribs[ad_key]
                    if single_mode == False:
                        return(legacy_ldap_attribs)
                    else:
                        return(list(legacy_ldap_attribs.items())[0][1])
                except:
                    logger.error('Exception encountered translating ad_attribs to legacy_ldap_attribs')
                    return(self.delay_return_none())
        return(self.delay_return_none())        

        
    def bind_ad_user(self, username, password):
        """
        tests the user credentials and pulls attributes from AD
        then converts ad-provided attributes to those
        expected by the legacy django implementation
        ljones 2022-06-03
        """
        
        connection = self.get_ad_connection(username, password)

        if not connection:
            logger.error('Bind failed for user %s' % username)
            return(self.delay_return_none())
            
        ldap_attribs = self.get_ad_attribs(connection, username, ['mail', 'sn', 'givenName', 'sAMAccountName', 'distinguishedName'], legacy_attribs = True)
        
        if not ldap_attribs:
            logger.error('Failed to fetch attribs for user %s' % username)
            return(self.delay_return_none())
        
        ldap_user = LDAPUser(connection, ldap_attribs)
        
        if not ldap_user:
            logger.error('Failed to build legacy LDAP user from attribs for %s' % username)
            return(self.delay_return_none())
        return(ldap_user)

    def bind_ldap_user(self, username, password):
        """
        Attempts to bind the specified username and password and returns
        an LDAPUser object representing the user.

        Returns None if the bind was unsuccessful.

        This implements direct binding.
        """

        #search through BASE_DNs for user, store first BASE_DN where user found
        server = ldap3.Server(settings.LEGACY_URI)
        c = ldap3.Connection(server, auto_bind=True)
        ldap_bind_user = None
        sfilter = '(&(objectClass=inetOrgPerson)(cn=' + username + '))'
        
        for dn in settings.LEGACY_DN:
            c.search(search_base = dn, search_filter = sfilter, search_scope = ldap3.LEVEL, attributes = ['cn'], size_limit=1)
            if c.response:
                ldap_bind_user = c.response[0]['dn']
                break
        
        #check if we failed to find someone
        if not ldap_bind_user:
            return None
        
        try:
            c = ldap3.Connection(self.backend,
                    read_only=True,
                    lazy=False,
                    auto_bind=True,
                    client_strategy=ldap3.SYNC,
                    authentication=ldap3.SIMPLE,
                    user=ldap_bind_user,
                    password=password)
        except ldap3.core.exceptions.LDAPSocketOpenError as e:
            logger.error('LDAP connection error: ' + str(e))
            return None
        except ldap3.core.exceptions.LDAPBindError as e:
            if 'invalidCredentials' in str(e):
                # Invalid bind DN or password
                return None
            else:
                logger.error('LDAP bind error: ' + str(e))
                return None
        except Exception as e:
            logger.exception('Caught exception when trying to connect and bind to LDAP')
            raise

        # Search for the user using their full DN
        search_filter = '({}={})'.format(settings.UID_ATTRIB, username)
        attributes = self.search_ldap(c, search_filter, attributes=LDAPUser._attrib_keys, size_limit=1)
        if not attributes:
            logger.error('LDAP search error: no results for ' + search_filter)
            return None

        # Construct an LDAPUser instance for this user
        return LDAPUser(c, attributes)

    def update_group_membership(self, ldap_user, django_user):
        """Update the user's group memberships

        Checks settings.GROUP_MAP to determine group memberships
        that should be added.
        """

        if not settings.GROUP_MAP:
            return None

        groups = {'add': [], 'remove': []}
        for ldap_group, django_groups in settings.GROUP_MAP.items():
            if self.check_group_membership(ldap_user, ldap_group):
                groups['add'] += [group for group in django_groups if group not in groups['add']]
            else:
                groups['remove'] += [group for group in django_groups if group not in groups['remove']]

        for operation in ('remove', 'add'):
            grouplist = groups[operation]
            for group in grouplist:
                try:
                    g = Group.objects.get(name=group)
                except Group.DoesNotExist:
                    logger.error('Django group does not exist: {}'.format(group))
                    continue
                else:
                    getattr(django_user.groups, operation)(g)

        django_user.save()
