Merge branch 'kaminski_k/ldapuserdir-master'

ldapuserdir/ldapuserdir.py

@@ -11,6 +11,7 @@ with an LDAP based user directory service
 """
 
 import ldap
+from ldap.controls import SimplePagedResultsControl
 #import ldap.ldapobject
 import sys
 import re
@@ -40,6 +41,10 @@ class LdapUserDir(object):
        base path for groups
     user_ou  : str, optional
        base path for users
+    page_size : int
+       page size for paged retrieval of results in search_s_reconn.
+       The default value is '500'.  A value of '0' disables paged
+       results.
     logger  : logger instance, optional
 
     Attributes
@@ -49,6 +54,7 @@ class LdapUserDir(object):
     user_ou : str
     user_dn : str
     user_pw : str
+    page_size : int
     logger : logger instance
 
     Raises
@@ -62,12 +68,14 @@ class LdapUserDir(object):
                  user_pw,
                  group_ou = 'ou=example.com',
                  user_ou = 'ou=example.com',
+                 page_size = 0,
                  logger = None):
         self.serverurl = serverurl
         self.group_ou = group_ou
         self.user_ou = user_ou
         self.user_dn = user_dn
         self.user_pw = user_pw
+        self.page_size = page_size
 
         if logger == None:
             self.logger = logging.getLogger('LdapUserDir')
@@ -85,6 +93,11 @@ class LdapUserDir(object):
                                     trace_file=sys.stderr)
         self.logger.debug('binding to: %s\n' % serverurl)
         self.logger.debug('binding as user: %s\n' % user_dn)
+
+        # Without this, paged results don't work (see the python-ldap FAQ for a
+        # hint as to why)
+        self._ldap.set_option(ldap.OPT_REFERRALS,0)
+
         try:
             self._ldap.bind_s(self.user_dn, self.user_pw)
         except ldap.INVALID_CREDENTIALS, e:
@@ -129,6 +142,44 @@ class LdapUserDir(object):
             raise RuntimeError("failed to convert DN to CN (%s)" % dn)
 
     
+    def _search_s(self, base, scope, filterstr='(objectClass=*)',
+                 attrlist=None, attrsonly=0):
+        """Helper for search_s_reconn. Wraps ldap.search_ext to use paged results if
+desired (see self.page_size)."""
+        if self.page_size == 0:
+            # Do not use paged results
+            return self._ldap.search_s(base, scope, filterstr, attrlist,
+                                       attrsonly)
+        else:
+            # Use paged results
+            page_ctrl = SimplePagedResultsControl(criticality=True,
+                                                  size=self.page_size,
+                                                  cookie='')
+            msgid = self._ldap.search_ext(base, scope, filterstr, attrlist,
+                                          attrsonly,
+                                          serverctrls=[page_ctrl])
+
+            results = []
+            while True:
+                _, rdata, _, resp_ctrls = self._ldap.result3(msgid)
+                results.extend(rdata)
+
+                # Extract the SimplePagedResultsControl to get the cookie.
+                page_ctrls = [c for c in resp_ctrls if c.controlType == SimplePagedResultsControl.controlType]
+                if page_ctrls == [] or page_ctrls[0].cookie == '':
+                    # We're done.
+                    break
+                else:
+                    # Update the cookie to retrieve the next page.
+                    page_ctrl.cookie = page_ctrls[0].cookie
+
+                msgid = self._ldap.search_ext(base, scope, filterstr, attrlist, attrsonly,
+                                              serverctrls=[page_ctrl])
+
+            # result4 return triples instead of tuples, despite what the
+            # python-ldap documentation says. Drop the third element.
+            return [r[:2] for r in results]
+
     def search_s_reconn(self, base, scope, filterstr='(objectClass=*)',
                         attrlist=None, attrsonly=0, recon_attempts = 2):
         """wrapper of standard ldap.search_s synchronous search that
@@ -161,18 +212,19 @@ class LdapUserDir(object):
         list of tuples
            list of tuples of the form (dn, attributes)
         """
+        
         attempts = 0
         ok = False
         while ok == False:
             try:
                 ok = True
                 attempts += 1
-                repl = self._ldap.search_s(base, scope, filterstr, attrlist,
-                                       attrsonly)
+                repl = self._search_s(base, scope, filterstr, attrlist,
+                                      attrsonly)
             except Exception, err:
                 ok = False
 
-                self.logger.warning("Got ldap error: %s" % str(err))
+                self.logger.warning("Got ldap error: %s" % (err,))
                 if attempts >= recon_attempts:
                     raise
 
@@ -181,7 +233,7 @@ class LdapUserDir(object):
                     del self._ldap
                 except Exception, err:
                     self.logger.warning("failed to delete LDAP object: %s"
-                                        % str(err))
+                                        % (err,))
 
                 self.logger.warning("Trying reconnecting to ldap (attempt %s)"
                                     % attempts)
@@ -194,12 +246,12 @@ class LdapUserDir(object):
                     self.logger.warning("ldap initialization error" +
                                         ", server down (server: %s)" %
                                         self.serverurl
-                                        + ": %s" % str(err))
+                                        + ": %s" % (err,))
                 except Exception, err:
                     self.logger.warning("ldap initialization error" +
                                         " (server: %s)" %
                                         self.serverurl
-                                        + ": %s" % str(err))
+                                        + ": %s" % (err,))
 
                 try:
                     self._ldap.bind_s(self.user_dn, self.user_pw)
@@ -210,7 +262,7 @@ class LdapUserDir(object):
                     self.logger.warning("ldap binding error" +
                                         " (server: %s)" %
                                         self.serverurl
-                                        + ": %s" % str(err))
+                                        + ": %s" % (err,))