import ldap
import logging
import os

from ldap.controls import SimplePagedResultsControl
from ldap.filter import filter_format, escape_filter_chars

from collections import namedtuple
from data.users.federated import FederatedUsers, UserInformation
from util.itertoolrecipes import take

logger = logging.getLogger(__name__)

_DEFAULT_NETWORK_TIMEOUT = 10.0 # seconds
_DEFAULT_TIMEOUT = 10.0 # seconds
_DEFAULT_PAGE_SIZE = 1000


class LDAPConnectionBuilder(object):
  def __init__(self, ldap_uri, user_dn, user_pw, allow_tls_fallback=False,
               timeout=None, network_timeout=None):
    self._ldap_uri = ldap_uri
    self._user_dn = user_dn
    self._user_pw = user_pw
    self._allow_tls_fallback = allow_tls_fallback
    self._timeout = timeout
    self._network_timeout = network_timeout

  def get_connection(self):
    return LDAPConnection(self._ldap_uri, self._user_dn, self._user_pw, self._allow_tls_fallback,
                          self._timeout, self._network_timeout)


class LDAPConnection(object):
  def __init__(self, ldap_uri, user_dn, user_pw, allow_tls_fallback=False,
               timeout=None, network_timeout=None):
    self._ldap_uri = ldap_uri
    self._user_dn = user_dn
    self._user_pw = user_pw
    self._allow_tls_fallback = allow_tls_fallback
    self._timeout = timeout
    self._network_timeout = network_timeout
    self._conn = None

  def __enter__(self):
    trace_level = 2 if os.environ.get('USERS_DEBUG') == '1' else 0

    self._conn = ldap.initialize(self._ldap_uri, trace_level=trace_level)
    self._conn.set_option(ldap.OPT_REFERRALS, 1)
    self._conn.set_option(ldap.OPT_NETWORK_TIMEOUT,
                          self._network_timeout or _DEFAULT_NETWORK_TIMEOUT)
    self._conn.set_option(ldap.OPT_TIMEOUT, self._timeout or _DEFAULT_TIMEOUT)

    if self._allow_tls_fallback:
      logger.debug('TLS Fallback enabled in LDAP')
      self._conn.set_option(ldap.OPT_X_TLS_TRY, 1)

    self._conn.simple_bind_s(self._user_dn, self._user_pw)
    return self._conn

  def __exit__(self, exc_type, value, tb):
    self._conn.unbind_s()


class LDAPUsers(FederatedUsers):
  _LDAPResult = namedtuple('LDAPResult', ['dn', 'attrs'])

  def __init__(self, ldap_uri, base_dn, admin_dn, admin_passwd, user_rdn, uid_attr, email_attr,
               allow_tls_fallback=False, secondary_user_rdns=None, requires_email=True,
               timeout=None, network_timeout=None, force_no_pagination=False):
    super(LDAPUsers, self).__init__('ldap', requires_email)

    self._ldap = LDAPConnectionBuilder(ldap_uri, admin_dn, admin_passwd, allow_tls_fallback,
                                       timeout, network_timeout)
    self._ldap_uri = ldap_uri
    self._uid_attr = uid_attr
    self._email_attr = email_attr
    self._allow_tls_fallback = allow_tls_fallback
    self._requires_email = requires_email
    self._force_no_pagination = force_no_pagination

    # Note: user_rdn is a list of RDN pieces (for historical reasons), and secondary_user_rds
    # is a list of RDN strings.
    relative_user_dns = [','.join(user_rdn)] + (secondary_user_rdns or [])

    def get_full_rdn(relative_dn):
      prefix = relative_dn.split(',') if relative_dn else []
      return ','.join(prefix + base_dn)

    # Create the set of full DN paths.
    self._user_dns = [get_full_rdn(relative_dn) for relative_dn in relative_user_dns]
    self._base_dn = ','.join(base_dn)

  def _get_ldap_referral_dn(self, referral_exception):
    logger.debug('Got referral: %s', referral_exception.args[0])
    if not referral_exception.args[0] or not referral_exception.args[0].get('info'):
      logger.debug('LDAP referral missing info block')
      return None

    referral_info = referral_exception.args[0]['info']
    if not referral_info.startswith('Referral:\n'):
      logger.debug('LDAP referral missing Referral header')
      return None

    referral_uri = referral_info[len('Referral:\n'):]
    if not referral_uri.startswith('ldap:///'):
      logger.debug('LDAP referral URI does not start with ldap:///')
      return None

    referral_dn = referral_uri[len('ldap:///'):]
    return referral_dn

  def _ldap_user_search_with_rdn(self, conn, username_or_email, user_search_dn, suffix=''):
    query = u'(|({0}={2}{3})({1}={2}{3}))'.format(self._uid_attr, self._email_attr,
                                                  escape_filter_chars(username_or_email),
                                                  suffix)
    logger.debug('Conducting user search: %s under %s', query, user_search_dn)
    try:
      return (conn.search_s(user_search_dn, ldap.SCOPE_SUBTREE, query.encode('utf-8')), None)
    except ldap.REFERRAL as re:
      referral_dn = self._get_ldap_referral_dn(re)
      if not referral_dn:
        return (None, 'Failed to follow referral when looking up username')

      try:
        subquery = u'(%s=%s)' % (self._uid_attr, username_or_email)
        return (conn.search_s(referral_dn, ldap.SCOPE_BASE, subquery), None)
      except ldap.LDAPError:
        logger.debug('LDAP referral search exception')
        return (None, 'Username not found')

    except ldap.LDAPError:
      logger.debug('LDAP search exception')
      return (None, 'Username not found')

  def _ldap_user_search(self, username_or_email, limit=20, suffix=''):
    if not username_or_email:
      return (None, 'Empty username/email')

    # Verify the admin connection works first. We do this here to avoid wrapping
    # the entire block in the INVALID CREDENTIALS check.
    try:
      with self._ldap.get_connection():
        pass
    except ldap.INVALID_CREDENTIALS:
      return (None, 'LDAP Admin dn or password is invalid')

    with self._ldap.get_connection() as conn:
      logger.debug('Incoming username or email param: %s', username_or_email.__repr__())

      for user_search_dn in self._user_dns:
        (pairs, err_msg) = self._ldap_user_search_with_rdn(conn, username_or_email, user_search_dn,
                                                           suffix=suffix)
        if pairs is not None and len(pairs) > 0:
          break

      if err_msg is not None:
        return (None, err_msg)

      logger.debug('Found matching pairs: %s', pairs)
      results = [LDAPUsers._LDAPResult(*pair) for pair in take(limit, pairs)]

      # Filter out pairs without DNs. Some LDAP impls will return such pairs.
      with_dns = [result for result in results if result.dn]
      return (with_dns, None)

  def _ldap_single_user_search(self, username_or_email):
    with_dns, err_msg = self._ldap_user_search(username_or_email)
    if err_msg is not None:
      return (None, err_msg)

    # Make sure we have at least one result.
    if len(with_dns) < 1:
      return (None, 'Username not found')

    # If we have found a single pair, then return it.
    if len(with_dns) == 1:
      return (with_dns[0], None)

    # Otherwise, there are multiple pairs with DNs, so find the one with the mail
    # attribute (if any).
    with_mail = [result for result in with_dns if result.attrs.get(self._email_attr)]
    return (with_mail[0] if with_mail else with_dns[0], None)

  def _build_user_information(self, response):
    if not response.get(self._uid_attr):
      return (None, 'Missing uid field "%s" in user record' % self._uid_attr)

    if self._requires_email and not response.get(self._email_attr):
      return (None, 'Missing mail field "%s" in user record' % self._email_attr)

    username = response[self._uid_attr][0].decode('utf-8')
    email = response.get(self._email_attr, [None])[0]
    return (UserInformation(username=username, email=email, id=username), None)

  def ping(self):
    try:
      with self._ldap.get_connection():
        pass
    except ldap.INVALID_CREDENTIALS:
      return (False, 'LDAP Admin dn or password is invalid')
    except ldap.LDAPError as lde:
      logger.exception('Exception when trying to health check LDAP')
      return (False, lde.message)

    return (True, None)

  def at_least_one_user_exists(self):
    logger.debug('Checking if any users exist in LDAP')
    try:
      with self._ldap.get_connection():
        pass
    except ldap.INVALID_CREDENTIALS:
      return (None, 'LDAP Admin dn or password is invalid')

    has_pagination = not self._force_no_pagination
    with self._ldap.get_connection() as conn:
      for user_search_dn in self._user_dns:
        lc = ldap.controls.libldap.SimplePagedResultsControl(criticality=True, size=1, cookie='')
        try:
          if has_pagination:
            msgid = conn.search_ext(user_search_dn, ldap.SCOPE_SUBTREE, serverctrls=[lc])
            _, rdata, _, serverctrls = conn.result3(msgid)
          else:
            msgid = conn.search(user_search_dn, ldap.SCOPE_SUBTREE)
            _, rdata = conn.result(msgid)

          for entry in rdata: # Handles both lists and iterators.
            return (True, None)

        except ldap.LDAPError as lde:
          return (False, str(lde) or 'Could not find DN %s' % user_search_dn)

    return (False, None)

  def get_user(self, username_or_email):
    """ Looks up a username or email in LDAP. """
    logger.debug('Looking up LDAP username or email %s', username_or_email)
    (found_user, err_msg) = self._ldap_single_user_search(username_or_email)
    if err_msg is not None:
      return (None, err_msg)

    logger.debug('Found user for LDAP username or email %s', username_or_email)
    _, found_response = found_user
    return self._build_user_information(found_response)

  def query_users(self, query, limit=20):
    """ Queries LDAP for matching users. """
    if not query:
      return (None, self.federated_service, 'Empty query')

    logger.debug('Got query %s with limit %s', query, limit)
    (results, err_msg) = self._ldap_user_search(query, limit=limit, suffix='*')
    if err_msg is not None:
      return (None, self.federated_service, err_msg)

    final_results = []
    for result in results[0:limit]:
      credentials, err_msg = self._build_user_information(result.attrs)
      if err_msg is not None:
        continue

      final_results.append(credentials)

    logger.debug('For query %s found results %s', query, final_results)
    return (final_results, self.federated_service, None)

  def verify_credentials(self, username_or_email, password):
    """ Verify the credentials with LDAP. """
    # Make sure that even if the server supports anonymous binds, we don't allow it
    if not password:
      return (None, 'Anonymous binding not allowed')

    (found_user, err_msg) = self._ldap_single_user_search(username_or_email)
    if found_user is None:
      return (None, err_msg)

    found_dn, found_response = found_user
    logger.debug('Found user for LDAP username %s; validating password', username_or_email)
    logger.debug('DN %s found: %s', found_dn, found_response)

    # First validate the password by binding as the user
    try:
      with LDAPConnection(self._ldap_uri, found_dn, password.encode('utf-8'),
                          self._allow_tls_fallback):
        pass
    except ldap.REFERRAL as re:
      referral_dn = self._get_ldap_referral_dn(re)
      if not referral_dn:
        return (None, 'Invalid username')

      try:
        with LDAPConnection(self._ldap_uri, referral_dn, password.encode('utf-8'),
                            self._allow_tls_fallback):
          pass
      except ldap.INVALID_CREDENTIALS:
        logger.debug('Invalid LDAP credentials')
        return (None, 'Invalid password')

    except ldap.INVALID_CREDENTIALS:
      logger.debug('Invalid LDAP credentials')
      return (None, 'Invalid password')

    return self._build_user_information(found_response)

  def service_metadata(self):
    return {
      'base_dn': self._base_dn,
    }

  def check_group_lookup_args(self, group_lookup_args, disable_pagination=False):
    if not group_lookup_args.get('group_dn'):
      return (False, 'Missing group_dn')

    (it, err) = self.iterate_group_members(group_lookup_args, page_size=1,
                                           disable_pagination=disable_pagination)
    if err is not None:
      return (False, err)

    if not next(it, False):
      return (False, 'Group does not exist or is empty')

    return (True, None)

  def iterate_group_members(self, group_lookup_args, page_size=None, disable_pagination=False):
    try:
      with self._ldap.get_connection():
        pass
    except ldap.INVALID_CREDENTIALS:
      return (None, 'LDAP Admin dn or password is invalid')

    group_dn = group_lookup_args['group_dn']
    page_size = page_size or _DEFAULT_PAGE_SIZE
    return (self._iterate_members(group_dn, page_size, disable_pagination), None)

  def _iterate_members(self, group_dn, page_size, disable_pagination):
    has_pagination = not(self._force_no_pagination or disable_pagination)
    with self._ldap.get_connection() as conn:
      search_flt = filter_format('(memberOf=%s,%s)', (group_dn, self._base_dn))
      attributes = [self._uid_attr, self._email_attr]

      for user_search_dn in self._user_dns:
        lc = ldap.controls.libldap.SimplePagedResultsControl(criticality=True, size=page_size,
                                                             cookie='')

        # Conduct the initial search for users that are a member of the group.
        logger.debug('Conducting LDAP search of DN: %s and filter %s', user_search_dn, search_flt)
        try:
          if has_pagination:
            msgid = conn.search_ext(user_search_dn, ldap.SCOPE_SUBTREE, search_flt,
                                    serverctrls=[lc], attrlist=attributes)
          else:
            msgid = conn.search(user_search_dn, ldap.SCOPE_SUBTREE, search_flt, attrlist=attributes)
        except ldap.LDAPError as lde:
          logger.exception('Got error when trying to search %s with filter %s: %s',
                           user_search_dn, search_flt, lde.message)
          break

        while True:
          try:
            if has_pagination:
              _, rdata, _, serverctrls = conn.result3(msgid)
            else:
              _, rdata = conn.result(msgid)

            # Yield any users found.
            found_results = 0
            for userdata in rdata:
              found_results = found_results + 1
              yield self._build_user_information(userdata[1])

            logger.debug('Found %s users in group %s; %s', found_results, user_search_dn,
                         search_flt)
          except ldap.NO_SUCH_OBJECT as nsoe:
            logger.debug('NSO when trying to lookup results of search %s with filter %s: %s',
                         user_search_dn, search_flt, nsoe.message)
          except ldap.LDAPError as lde:
            logger.exception('Error when trying to lookup results of search %s with filter %s: %s',
                             user_search_dn, search_flt, lde.message)
            break

          # If no additional results, nothing more to do.
          if not found_results:
            break

          # If pagination is disabled, nothing more to do.
          if not has_pagination:
            logger.debug('Pagination is disabled, no further queries')
            break

          # Filter down the controls with which the server responded, looking for the paging
          # control type. If not found, then the server does not support pagination and we already
          # got all of the results.
          pctrls = [control for control in serverctrls
                    if control.controlType == ldap.controls.SimplePagedResultsControl.controlType]

          if pctrls:
            # Server supports pagination. Update the cookie so the next search finds the next page,
            # then conduct the next search.
            cookie = lc.cookie = pctrls[0].cookie
            if cookie:
              logger.debug('Pagination is supported for this LDAP server; trying next page')
              msgid = conn.search_ext(user_search_dn, ldap.SCOPE_SUBTREE, search_flt,
                                      serverctrls=[lc], attrlist=attributes)
              continue
            else:
              # No additional results.
              logger.debug('Pagination is supported for this LDAP server but on last page')
              break
          else:
            # Pagination is not supported.
            logger.debug('Pagination is not supported for this LDAP server')
            break