Initial LDAP group member iteration support

Add interface for group member iteration on internal auth providers and implement support in the LDAP interface.

data/users/externalldap.py

@@ -2,6 +2,8 @@ import ldap
 import logging
 import os
 
+from ldap.controls import SimplePagedResultsControl
+
 from collections import namedtuple
 from data.users.federated import FederatedUsers, UserInformation
 from util.itertoolrecipes import take
@@ -10,6 +12,7 @@ logger = logging.getLogger(__name__)
 
 _DEFAULT_NETWORK_TIMEOUT = 10.0 # seconds
 _DEFAULT_TIMEOUT = 10.0 # seconds
+_DEFAULT_PAGE_SIZE = 500
 
 
 class LDAPConnectionBuilder(object):
@@ -84,6 +87,7 @@ class LDAPUsers(FederatedUsers):
 
     # 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])
@@ -174,7 +178,7 @@ class LDAPUsers(FederatedUsers):
     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 _credential_for_user(self, response):
+  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)
 
@@ -194,7 +198,7 @@ class LDAPUsers(FederatedUsers):
 
     logger.debug('Found user for LDAP username or email %s', username_or_email)
     _, found_response = found_user
-    return self._credential_for_user(found_response)
+    return self._build_user_information(found_response)
 
   def query_users(self, query, limit=20):
     """ Queries LDAP for matching users. """
@@ -208,7 +212,7 @@ class LDAPUsers(FederatedUsers):
 
     final_results = []
     for result in results[0:limit]:
-      credentials, err_msg = self._credential_for_user(result.attrs)
+      credentials, err_msg = self._build_user_information(result.attrs)
       if err_msg is not None:
         continue
 
@@ -253,4 +257,65 @@ class LDAPUsers(FederatedUsers):
       logger.debug('Invalid LDAP credentials')
       return (None, 'Invalid password')
 
-    return self._credential_for_user(found_response)
+    return self._build_user_information(found_response)
+
+  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):
+    with self._ldap.get_connection() as conn:
+      lc = ldap.controls.libldap.SimplePagedResultsControl(criticality=True, size=page_size,
+                                                           cookie='')
+
+      search_flt = '(memberOf=%s,%s)' % (group_dn, self._base_dn)
+
+      for user_search_dn in self._user_dns:
+        # Conduct the initial search for users that are a member of the group.
+        if disable_pagination:
+          msgid = conn.search(user_search_dn, ldap.SCOPE_SUBTREE, search_flt)
+        else:
+          msgid = conn.search_ext(user_search_dn, ldap.SCOPE_SUBTREE, search_flt, serverctrls=[lc])
+
+        while True:
+          if disable_pagination:
+            _, rdata = conn.result(msgid)
+          else:
+            _, rdata, _, serverctrls = conn.result3(msgid)
+
+          # Yield any users found.
+          for userdata in rdata:
+            yield self._build_user_information(userdata[1])
+
+          # If pagination is disabled, nothing more to do.
+          if disable_pagination:
+            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:
+                msgid = conn.search_ext(user_search_dn, ldap.SCOPE_SUBTREE, search_flt,
+                                        serverctrls=[lc])
+                continue
+            else:
+              # No additional results.
+              break
+          else:
+            # Pagintation is not supported.
+            logger.debug('Pagination is not supported for this LDAP server')
+            break