import ldap
import logging
import os

from collections import namedtuple
from datetime import datetime
from data.users.federated import FederatedUsers, VerifiedCredentials

logger = logging.getLogger(__name__)


class LDAPConnectionBuilder(object):
  def __init__(self, ldap_uri, user_dn, user_pw):
    self._ldap_uri = ldap_uri
    self._user_dn = user_dn
    self._user_pw = user_pw

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


class LDAPConnection(object):
  def __init__(self, ldap_uri, user_dn, user_pw):
    self._ldap_uri = ldap_uri
    self._user_dn = user_dn
    self._user_pw = user_pw
    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.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):
    super(LDAPUsers, self).__init__('ldap')
    self._ldap = LDAPConnectionBuilder(ldap_uri, admin_dn, admin_passwd)
    self._ldap_uri = ldap_uri
    self._base_dn = base_dn
    self._user_rdn = user_rdn
    self._uid_attr = uid_attr
    self._email_attr = email_attr

  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(self, username_or_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__())
      user_search_dn = ','.join(self._user_rdn + self._base_dn)
      query = u'(|({0}={2})({1}={2}))'.format(self._uid_attr, self._email_attr,
                                              username_or_email)

      logger.debug('Conducting user search: %s under %s', query, user_search_dn)
      try:
        pairs = conn.search_s(user_search_dn, ldap.SCOPE_SUBTREE, query.encode('utf-8'))
      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)
          pairs = conn.search_s(referral_dn, ldap.SCOPE_BASE, subquery)
        except ldap.LDAPError:
          logger.exception('LDAP referral search exception')
          return (None, 'Username not found')

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

      logger.debug('Found matching pairs: %s', pairs)

      results = [LDAPUsers._LDAPResult(*pair) for pair in pairs]

      # Filter out pairs without DNs. Some LDAP impls will return such
      # pairs.
      with_dns = [result for result in results if result.dn]
      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 results if result.attrs.get(self._email_attr)]
      return (with_mail[0] if with_mail else with_dns[0], 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_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')):
        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')):
          pass
      except ldap.INVALID_CREDENTIALS:
        logger.exception('Invalid LDAP credentials')
        return (None, 'Invalid password')

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

    # Now check if we have a federated login for this user
    if not found_response.get(self._uid_attr):
      return (None, 'Missing uid field "%s" in user record' % self._uid_attr)

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

    username = found_response[self._uid_attr][0].decode('utf-8')
    email = found_response[self._email_attr][0]
    return (VerifiedCredentials(username=username, email=email), None)