From 4135d00dd3a6e0c0c38156b32cd85bd7778537b5 Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Tue, 15 Dec 2020 18:24:50 +0100
Subject: [PATCH] LdapServer: Add paged_search()

---
 LdapServer.py | 70 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/LdapServer.py b/LdapServer.py
index 17f2b7c..361be1a 100644
--- a/LdapServer.py
+++ b/LdapServer.py
@@ -5,6 +5,7 @@ import datetime
 import dateutil.parser
 import dateutil.tz
 import ldap
+from ldap.controls import SimplePagedResultsControl
 import ldap.modlist as modlist
 import logging
 import pytz
@@ -87,6 +88,75 @@ class LdapServer(object):
 		result = self.search(dn, filterstr=filterstr, scope='base', attrs=attrs)
 		return result[dn] if dn in result else None
 
+	def paged_search(self, basedn, filterstr, attrs, scope='sub', pagesize=500):
+		# Initialize SimplePagedResultsControl object
+		page_control = SimplePagedResultsControl(
+			True,
+			size=pagesize,
+			cookie=''  # Start without cookie
+		)
+		ret = {}
+		pages_count = 0
+		self.logger.debug(
+			"LdapServer - Paged search with base DN '%s', filter '%s', scope '%s', pagesize=%d and attrs=%s",
+			basedn,
+			filterstr,
+			scope,
+			pagesize,
+			attrs
+		)
+		while True:
+			pages_count += 1
+			self.logger.debug(
+				"LdapServer - Paged search: request page %d with a maximum of %d objects (current total count: %d)",
+				pages_count,
+				pagesize,
+				len(ret)
+			)
+			try:
+				res_id = self.con.search_ext(
+					basedn,
+					self.get_scope(scope),
+					filterstr,
+					attrs,
+					serverctrls=[page_control]
+				)
+			except ldap.LDAPError, e:
+				self._error('LdapServer - Error running paged search on LDAP server: %s' % e, logging.CRITICAL)
+				return False
+			try:
+				rtype, rdata, rmsgid, rctrls = self.con.result3(res_id)
+			except ldap.LDAPError, e:
+				self._error('LdapServer - Error pulling paged search result from LDAP server: %s' % e, logging.CRITICAL)
+				return False
+
+			# Detect and catch PagedResultsControl answer from rctrls
+			result_page_control = None
+			if rctrls:
+				for rctrl in rctrls:
+					if rctrl.controlType == SimplePagedResultsControl.controlType:
+						result_page_control = rctrl
+						break
+
+			# If PagedResultsControl answer not detected, paged serach
+			if not result_page_control:
+				self._error('LdapServer - Server ignores RFC2696 control, paged search can not works', logging.CRITICAL)
+				return False
+
+			# Store results of this page
+			for obj_dn, obj_attrs in rdata:
+				ret[obj_dn] = obj_attrs
+
+			# If no cookie returned, we are done
+			if not result_page_control.cookie:
+				break
+
+			# Otherwise, set cookie for the next search
+			page_control.cookie = result_page_control.cookie
+
+		self.logger.debug("LdapServer - Paged search end: %d object(s) retreived in %d page(s) of %d object(s)", len(ret), pages_count, pagesize)
+		return ret
+
 	def add_object(self,dn,attrs):
 		ldif = modlist.addModlist(attrs)
 		try: