From 3d76fb36d311f86a285158cc8e00482648548881 Mon Sep 17 00:00:00 2001 From: George Taylor Date: Thu, 25 Jul 2024 18:25:23 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Use=20paged=20search=20to=20support?= =?UTF-8?q?=20larger=20data=20sets=20(#58)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * :ambulance: correct logic for matched user sets + role filtering * Formatted code with black --line-length 120 * Update role filter logic in user.py * Formatted code with black --line-length 120 * aliasing * Update user.py * Update user.py * Formatted code with black --line-length 120 * implement paging because we love ldap :heart: * remove test case + add var for page size * Formatted code with black --line-length 120 --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com> --- cli/ldap_cmds/user.py | 63 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 52 insertions(+), 11 deletions(-) diff --git a/cli/ldap_cmds/user.py b/cli/ldap_cmds/user.py index 3f67140..2ad4a02 100644 --- a/cli/ldap_cmds/user.py +++ b/cli/ldap_cmds/user.py @@ -10,6 +10,8 @@ ) import ldap +from ldap.controls import SimplePagedResultsControl +import ldap.modlist as modlist from cli.ldap_cmds import ( ldap_connect, @@ -187,7 +189,7 @@ def update_roles(roles, user_ou, root_dn, add, remove, update_notes, user_note, roles = roles.split(",") - # create role filter + # Create role filter if len(roles_to_filter) > 0: full_role_filter = f"(&(objectclass=NDRoleAssociation)(|{''.join(['(cn=' + role + ')' for role in roles_to_filter.split(',')])}))" else: @@ -195,30 +197,69 @@ def update_roles(roles, user_ou, root_dn, add, remove, update_notes, user_note, log.debug(full_role_filter) - # Search for roles matching the role_filter - try: ldap_connection_role_filter = ldap.initialize("ldap://" + env.vars.get("LDAP_HOST")) ldap_connection_role_filter.simple_bind_s(env.vars.get("LDAP_USER"), env.secrets.get("LDAP_BIND_PASSWORD")) - except Exception as e: + ldap_connection_role_filter.set_option(ldap.OPT_REFERRALS, 0) + except ldap.LDAPError as e: log.exception("Failed to connect to LDAP") raise e + roles_search_result = [] + pages = 0 + if env.vars.get("LDAP_PAGE_SIZE") is None: + ldap_page_size = 100 + else: + try: + ldap_page_size = int(env.vars.get("LDAP_PAGE_SIZE")) + except ValueError: + log.error("LDAP_PAGE_SIZE must be an integer") + raise ValueError("LDAP_PAGE_SIZE must be an integer") + + page_control = SimplePagedResultsControl(True, size=ldap_page_size, cookie="") + try: - role_filter_results = ldap_connection_role_filter.search_ext_s( - ",".join([user_ou, root_dn]), ldap.SCOPE_SUBTREE, full_role_filter, ["cn"], sizelimit=0 + response = ldap_connection_role_filter.search_ext( + ",".join([user_ou, root_dn]), ldap.SCOPE_SUBTREE, full_role_filter, ["cn"], serverctrls=[page_control] ) - except Exception as e: + + while True: + pages += 1 + log.debug(f"Processing page {pages}") + try: + rtype, rdata, rmsgid, serverctrls = ldap_connection_role_filter.result3(response) + roles_search_result.extend(rdata) + cookie = serverctrls[0].cookie + print(cookie) + if cookie: + page_control.cookie = cookie + response = ldap_connection_role_filter.search_ext( + ",".join([user_ou, root_dn]), + ldap.SCOPE_SUBTREE, + full_role_filter, + ["cn"], + serverctrls=[page_control], + ) + else: + break + except ldap.LDAPError as e: + log.exception("Error retrieving LDAP results") + raise e + + except ldap.LDAPError as e: log.exception("Failed to search for roles") raise e - roles_found = sorted(set({dn.split(",")[1].split("=")[1] for dn, entry in role_filter_results})) - log.debug("users found from roles filter: ") + finally: + ldap_connection_role_filter.unbind_s() + + roles_found = sorted(set({dn.split(",")[1].split("=")[1] for dn, entry in roles_search_result})) + + roles_found = sorted(roles_found) + log.debug("Users found from roles filter: ") log.debug(roles_found) log.info(f"Found {len(roles_found)} users with roles matching the role filter") - ldap_connection_role_filter.unbind() - # generate a list of matches in roles and users users_found_set = set(users_found) roles_found_set = set(roles_found)