MiscUtil: CERN LDAP plugin improvement

* Improves the search function and introduces paged search to
  avoid exceeding the size limit of the CERN LDAP server.

* Adds function which returns the user data given a searchfilter
  and attribute list.

* Fixes kwalitee.

Signed-off-by: Jochen Klein <[email protected]>

Author: jochenklein

modules/miscutil/lib/ldap_cern.py

@@ -1,60 +1,155 @@
-## This file is part of Invenio.
-## Copyright (C) 2009, 2010, 2011, 2014 CERN.
-##
-## Invenio is free software; you can redistribute it and/or
-## modify it under the terms of the GNU General Public License as
-## published by the Free Software Foundation; either version 2 of the
-## License, or (at your option) any later version.
-##
-## Invenio is distributed in the hope that it will be useful, but
-## WITHOUT ANY WARRANTY; without even the implied warranty of
-## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
-## General Public License for more details.
-##
-## You should have received a copy of the GNU General Public License
-## along with Invenio; if not, write to the Free Software Foundation, Inc.,
-## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
-
-"""Invenio LDAP interface for CERN. """
+# This file is part of Invenio.
+# Copyright (C) 2009, 2010, 2011, 2014, 2015 CERN.
+#
+# Invenio is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License as
+# published by the Free Software Foundation; either version 2 of the
+# License, or (at your option) any later version.
+#
+# Invenio is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Invenio; if not, write to the Free Software Foundation, Inc.,
+# 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
+
+"""Invenio LDAP interface for CERN."""
 
-from time import sleep
 from thread import get_ident
+from time import sleep
 
 import ldap
+from ldap.controls import SimplePagedResultsControl
 import ldap.filter
 
 CFG_CERN_LDAP_URI = "ldap://xldap.cern.ch:389"
 CFG_CERN_LDAP_BASE = "OU=Users,OU=Organic Units,DC=cern,DC=ch"
+CFG_CERN_LDAP_PAGESIZE = 250
 
 _ldap_connection_pool = {}
 
 
+class LDAPError(Exception):
+
+    """Base class for exceptions in this module."""
+
+    pass
+
+
 def _cern_ldap_login():
-    """Get a connection from _ldap_connection_pool or create a new one"""
+    """Get a connection from _ldap_connection_pool or create a new one."""
     try:
         connection = _ldap_connection_pool[get_ident()]
     except KeyError:
         connection = _ldap_connection_pool[get_ident()] = ldap.initialize(CFG_CERN_LDAP_URI)
+
+    connection.set_option(ldap.OPT_PROTOCOL_VERSION, 3)
     return connection
 
 
 def _sanitize_input(query):
     """
-    Take the query, filter it through ldap.filter.escape_filter_chars and
+    Sanitize input query.
+
+    Take the query, filter it through ldap.filter.escape_filter_chars, and
     replace the dots with spaces.
     """
     query = ldap.filter.escape_filter_chars(query)
     query = query.replace(".", " ")
     return query
 
 
-def get_users_info_by_displayName(displayName):
+def _msgid(connection, req_ctrl, query_filter, attr_list=None):
+    """Run the search request using search_ext.
+
+    :param string query_filter: filter to apply in the LDAP search
+    :param list attr_list: retrieved LDAP attributes. If None, all attributes
+        are returned
+    :return: msgid
+    """
+    try:
+        return connection.search_ext(
+            CFG_CERN_LDAP_BASE,
+            ldap.SCOPE_SUBTREE,
+            query_filter,
+            attr_list,
+            attrsonly=0,
+            serverctrls=[req_ctrl])
+    except ldap.SERVER_DOWN as e:
+        raise LDAPError("Error: Connection to CERN LDAP failed. ({0})"
+                        .format(e))
+
+
+def _paged_search(connection, query_filter, attr_list=None):
+    """Search the CERN LDAP server using pagination.
+
+    :param string query_filter: filter to apply in the LDAP search
+    :param list attr_list: retrieved LDAP attributes. If None, all attributes
+        are returned
+    :return: list of tuples (result-type, result-data) or empty list,
+        where result-data contains the user dictionary
     """
+    req_ctrl = SimplePagedResultsControl(True, CFG_CERN_LDAP_PAGESIZE, "")
+    msgid = _msgid(connection, req_ctrl, query_filter, attr_list)
+    result_pages = 0
+    results = []
+
+    while True:
+        rtype, rdata, rmsgid, rctrls = connection.result3(msgid)
+        results.extend(rdata)
+        result_pages += 1
+
+        pctrls = [
+            c
+            for c in rctrls
+            if c.controlType == SimplePagedResultsControl.controlType
+        ]
+        if pctrls:
+            if pctrls[0].cookie:
+                req_ctrl.cookie = pctrls[0].cookie
+                msgid = _msgid(connection, req_ctrl,
+                               query_filter, attr_list)
+            else:
+                break
+
+    return results
+
+
+def get_users_records_data(query_filter, attr_list=None, decode_encoding=None):
+    """Get result-data of records.
+
+    :param string query_filter: filter to apply in the LDAP search
+    :param list attr_list: retrieved LDAP attributes. If None, all attributes
+        are returned
+    :param string decode_encoding: decode the values of the LDAP records
+    :return: list of LDAP records, but result-data only
+    """
+    connection = _cern_ldap_login()
+    records = _paged_search(connection, query_filter, attr_list)
+
+    records_data = []
+
+    if decode_encoding:
+        records_data = [
+            dict(
+                (k, [v[0].decode(decode_encoding)]) for (k, v) in x.iteritems()
+            )
+            for (dummy, x) in records]
+    else:
+        records_data = [x for (dummy, x) in records]
+
+    return records_data
+
+
+def get_users_info_by_displayName(displayName):
+    """Get user information, given the displayName.
+
     Query the CERN LDAP server for information about all users whose name
     contains the displayName.
     Return a list of user dictionaries (or empty list).
     """
-
     connection = _cern_ldap_login()
 
     # Split displayName and add each part of it to the search query
@@ -64,22 +159,20 @@ def get_users_info_by_displayName(displayName):
         query_filter = "& "
         for element in query_elements:
             query_filter += '(displayName=*%s*) ' % element
-        # Query will look like that: "(& (displayName=*john*) (displayName=*smith*)"
+        # Query will look like: "(& (displayName=*john*) (displayName=*smith*)"
         # Eliminate the secondary accounts (aliases, etc.)
         query_filter = "(& (%s) (| (employeetype=primary) (employeetype=external) (employeetype=ExCern) ) )" % query_filter
     else:
         return []
 
     try:
-        results = connection.search_st(CFG_CERN_LDAP_BASE, ldap.SCOPE_SUBTREE,
-                                       query_filter, timeout=5)
+        results = _paged_search(connection, query_filter)
     except ldap.LDAPError:
-        ## Mmh.. connection error? Let's reconnect at least once just in case
+        # Mmh.. connection error? Let's reconnect at least once just in case
         sleep(1)
         connection = _cern_ldap_login()
         try:
-            results = connection.search_st(CFG_CERN_LDAP_BASE, ldap.SCOPE_SUBTREE,
-                                           query_filter, timeout=5)
+            results = _paged_search(connection, query_filter)
         except ldap.LDAPError:
             # Another error (maybe the LDAP query size is too big, etc.)
             # TODO, if it's needed, here we can return various different
@@ -89,12 +182,12 @@ def get_users_info_by_displayName(displayName):
 
 
 def get_users_info_by_displayName_or_email(name):
-    """
-    Query the CERN LDAP server for information about all users whose displayName
-    or email contains the name.
+    """Get user information, given displayName or email address.
+
+    Query the CERN LDAP server for information about all users whose
+    displayName or email contains the name.
     Return a list of user dictionaries (or empty list).
     """
-
     connection = _cern_ldap_login()
 
     # Split name and add each part of it to the search query
@@ -114,15 +207,13 @@ def get_users_info_by_displayName_or_email(name):
         return []
 
     try:
-        results = connection.search_st(CFG_CERN_LDAP_BASE, ldap.SCOPE_SUBTREE,
-                                       query_filter, timeout=5)
+        results = _paged_search(connection, query_filter)
     except ldap.LDAPError:
-        ## Mmh.. connection error? Let's reconnect at least once just in case
+        # Mmh.. connection error? Let's reconnect at least once just in case
         sleep(1)
         connection = _cern_ldap_login()
         try:
-            results = connection.search_st(CFG_CERN_LDAP_BASE, ldap.SCOPE_SUBTREE,
-                                           query_filter, timeout=5)
+            results = _paged_search(connection, query_filter)
         except ldap.LDAPError:
             # Another error (maybe the LDAP query size is too big, etc.)
             # TODO, if it's needed, here we can return various different