✨ Implement paging for searching users and roles in update_user_roles job

cli/__init__.py

@@ -34,7 +34,7 @@ def add_roles_to_users(
     root_dn,
     user_role_list,
 ):
-    cli.ldap.user.process_user_roles_list(
+    cli.ldap_cmds.user.process_user_roles_list(
         user_role_list,
         user_ou,
         root_dn,
@@ -73,7 +73,7 @@ def update_user_home_areas(
     user_ou,
     root_dn,
 ):
-    cli.ldap.user.change_home_areas(
+    cli.ldap_cmds.user.change_home_areas(
         old_home_area,
         new_home_area,
         user_ou,
@@ -118,32 +118,10 @@ def update_user_home_areas(
     help="Remove role from users",
     is_flag=True,
 )
-@click.option(
-    "-rf",
-    "--role-filter",
-    help='Comma separated string to generate roles filter from eg "role1,role2,role3"',
-    required=False,
-    default="*",
-)
-@click.option(
-    "-uf",
-    "--user-filter",
-    help="Filter to find users",
-    required=False,
-    default="(userSector=*)",
-)
-def update_user_roles(
-    roles,
-    user_ou,
-    root_dn,
-    add,
-    remove,
-    update_notes,
-    user_note,
-    user_filter,
-    role_filter,
-):
-    cli.ldap.user.update_roles(
+@click.option("-uf", "--user-filter", help="Filter to find users", required=False, default="(objectclass=*)")
+@click.option("--roles-to-filter", help="Roles to filter", required=False, default="*")
+def update_user_roles(roles, user_ou, root_dn, add, remove, update_notes, user_note, user_filter, roles_to_filter):
+    cli.ldap_cmds.user.update_roles(
         roles,
         user_ou,
         root_dn,
@@ -152,7 +130,7 @@ def update_user_roles(
         update_notes,
         user_note=user_note,
         user_filter=user_filter,
-        role_filter=role_filter,
+        roles_to_filter=roles_to_filter,
     )
 
 

cli/ldap_cmds/user.py

@@ -9,13 +9,17 @@
     env,
 )
 
+import ldap
+from ldap.controls import SimplePagedResultsControl
+import ldap.modlist as modlist
+
 from cli.ldap_cmds import (
     ldap_connect,
 )
 from ldap3 import (
     MODIFY_REPLACE,
     MODIFY_DELETE,
-    DEREF_NEVER,
+    DEREF_ALWAYS,
 )
 
 import cli.database
@@ -150,89 +154,126 @@ def process_user_roles_list(user_role_list, user_ou="ou=Users", root_dn="dc=moj,
 #########################################
 
 
-def update_roles(
-    roles, user_ou, root_dn, add, remove, update_notes, user_note, user_filter="(userSector=*)", role_filter="*"
-):
+def update_roles(roles, user_ou, root_dn, add, remove, update_notes, user_note, user_filter, roles_to_filter):
     if update_notes and (user_note is None or len(user_note) < 1):
         log.error("User note must be provided when updating notes")
         raise Exception("User note must be provided when updating notes")
 
     try:
-        ldap_connection_user_filter = ldap_connect(
-            env.vars.get("LDAP_HOST"),
-            env.vars.get("LDAP_USER"),
-            env.secrets.get("LDAP_BIND_PASSWORD"),
-        )
+        ldap_connection_user_filter = ldap.initialize("ldap://" + env.vars.get("LDAP_HOST"))
+        ldap_connection_user_filter.simple_bind_s(env.vars.get("LDAP_USER"), env.secrets.get("LDAP_BIND_PASSWORD"))
     except Exception as e:
         log.exception("Failed to connect to LDAP")
         raise e
 
     # # Search for users matching the user_filter
+
+    user_filter = f"(&(objectclass=NDUser){user_filter})"
+    log.debug(f"User filter: {user_filter}")
     try:
-        ldap_connection_user_filter.search(
+        user_filter_results = ldap_connection_user_filter.search_s(
             ",".join([user_ou, root_dn]),
+            ldap.SCOPE_SUBTREE,
             user_filter,
-            attributes=["cn"],
+            ["cn"],
         )
     except Exception as e:
         log.exception("Failed to search for users")
         raise e
 
-    users_found = sorted([entry.cn.value for entry in ldap_connection_user_filter.entries if entry.cn.value])
+    users_found = sorted(set([entry[1]["cn"][0].decode("utf-8") for entry in user_filter_results]))
     log.debug("users found from user filter")
     log.debug(users_found)
+    log.info(f"Found {len(users_found)} users matching the user filter")
     ldap_connection_user_filter.unbind()
 
-    roles_filter_list = role_filter.split(",")
     roles = roles.split(",")
 
-    # create role filter
-    if len(roles_filter_list) > 0:
-        full_role_filter = (
-            f"(&(objectclass=NDRoleAssociation)(|{''.join(['(cn=' + role + ')' for role in roles_filter_list])}))"
-        )
+    # 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:
         full_role_filter = "(&(objectclass=NDRoleAssociation)(cn=*))"
 
-    # Search for roles matching the role_filter
+    log.debug(full_role_filter)
 
     try:
-        ldap_connection_role_filter = ldap_connect(
-            env.vars.get("LDAP_HOST"),
-            env.vars.get("LDAP_USER"),
-            env.secrets.get("LDAP_BIND_PASSWORD"),
-        )
-    except Exception as e:
+        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"))
+        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:
-        ldap_connection_role_filter.search(
-            ",".join([user_ou, root_dn]),
-            full_role_filter,
-            attributes=["cn"],
-            dereference_aliases=DEREF_NEVER,
+        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({entry.entry_dn.split(",")[1].split("=")[1] for entry in ldap_connection_role_filter.entries})
-    )
-    log.debug("users found from roles filter: ")
-    log.debug(roles_found)
+    finally:
+        ldap_connection_role_filter.unbind_s()
+
+    roles_found = sorted(set({dn.split(",")[1].split("=")[1] for dn, entry in roles_search_result}))
 
-    ldap_connection_role_filter.unbind()
+    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")
 
     # generate a list of matches in roles and users
-    matched_users = set(users_found) & set(roles_found)
+    users_found_set = set(users_found)
+    roles_found_set = set(roles_found)
+
+    log.debug(users_found_set)
+    log.debug(roles_found_set)
+
+    matched_users = sorted(users_found_set.intersection(roles_found_set))
     log.debug("matched users: ")
     log.debug(matched_users)
 
     # cartesian_product = [(user, role) for user in matched_users for role in roles]
-
     cartesian_product = list(product(matched_users, roles))
+    log.info(f"Created {len(cartesian_product)} combinations of users and roles")
     log.debug("cartesian product: ")
     log.debug(cartesian_product)
 
@@ -246,6 +287,9 @@ def update_roles(
         log.exception("Failed to connect to LDAP")
         raise e
 
+    actioned = 0
+    not_actioned = 0
+    failed = 0
     for item in cartesian_product:
         if add:
             try:
@@ -262,23 +306,44 @@ def update_roles(
                 raise e
             if ldap_connection_action.result["result"] == 0:
                 log.info(f"Successfully added role '{item[1]}' to user '{item[0]}'")
+                actioned = actioned + 1
             elif ldap_connection_action.result["result"] == 68:
                 log.info(f"Role '{item[1]}' already present for user '{item[0]}'")
+                not_actioned = not_actioned + 1
             else:
                 log.e(f"Failed to add role '{item[1]}' to user '{item[0]}'")
                 log.debug(ldap_connection_action.result)
         elif remove:
+            removed = 0
+            not_removed = 0
+            failed = 0
             ldap_connection_action.delete(f"cn={item[1]},cn={item[0]},{user_ou},{root_dn}")
             if ldap_connection_action.result["result"] == 0:
                 log.info(f"Successfully removed role '{item[1]}' from user '{item[0]}'")
+                actioned = actioned + 1
             elif ldap_connection_action.result["result"] == 32:
                 log.info(f"Role '{item[1]}' already absent for user '{item[0]}'")
+                not_actioned = not_actioned + 1
             else:
                 log.error(f"Failed to remove role '{item[1]}' from user '{item[0]}'")
                 log.debug(ldap_connection_action.result)
+                failed = failed + 1
         else:
             log.error("No action specified")
 
+    log.info("\n==========================\n\tSUMMARY\n==========================")
+    log.info("User/role searches:")
+    log.info(f"    - Found {len(roles_found)} users with roles matching the role filter")
+    log.info(f"    - Found {len(users_found)} users matching the user filter")
+
+    log.info("This produces the following matches:")
+    log.info(f"    - Found {len(matched_users)} users with roles matching the role filter and user filter")
+
+    log.info("Actions:")
+    log.info(f"    - Successfully actioned {actioned} roles")
+    log.info(f"    - Roles already in desired state for {not_actioned} users")
+    log.info(f"    - Failed to remove {failed} roles due to errors")
+
     if update_notes:
         connection = cli.database.connection()
         log.debug("Created database cursor successfully")

requirements.txt

@@ -8,4 +8,4 @@ python-dotenv==1.0.1
 Jinja2==3.1.4
 python-ldap
 requests==2.32.3
-setuptools==71.1.0
+setuptools==72.0.0

setup.py

@@ -14,7 +14,7 @@
 
 setup(
     name="ldap-automation",
-    version="0.1",
+    version="0.2",
     packages=find_packages(),
     install_requires=all_reqs,
     entry_points="""