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.

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

modules/miscutil/lib/ldap_cern.py

@@ -1,40 +1,48 @@
-## 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
 
 import ldap
 import ldap.filter
+from ldap.controls import SimplePagedResultsControl
 
 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):
+    pass
+
+
 def _cern_ldap_login():
     """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
 
 
@@ -48,6 +56,86 @@ def _sanitize_input(query):
     return query
 
 
+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: encoding format for decoding 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 = [
+            {(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):
     """
     Query the CERN LDAP server for information about all users whose name
@@ -71,15 +159,13 @@ def get_users_info_by_displayName(displayName):
         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
         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
@@ -114,15 +200,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
         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