Sort password policies properly with --pkey-only

Password policy plugin sorts password policies by its COS priority.
However, when the pwpolicy-find command is run with --pkey-only,
the resulting entries do not contain COS priority and the sort
function crashes.

This patch makes sure that cospriority is present in the time
of the result sorting process and removes the cospriority again
when the sorting is done. This way, the entries are sorted properly
both with and without --pkey-only flag.

Previous entries_sortfn member attribute of LDAPSearch class
containing custom user sorting function was replaced just with
a flag indicating if a sorting in LDAPSearch shall be done at all.
This change makes it possible to sort entries in a custom
post_callback which is much more powerful (and essential for
sorting like in pwpolicy plugin) approach than a plain sorting
function.

https://fedorahosted.org/freeipa/ticket/2676

2 files changed, 31 insertions, 12 deletions

diff --git a/ipalib/plugins/baseldap.py b/ipalib/plugins/baseldap.py
index 7ee5fa1e1..d37a20d1f 100644
--- a/ipalib/plugins/baseldap.py
+++ b/ipalib/plugins/baseldap.py
@@ -1675,9 +1675,10 @@ class LDAPSearch(BaseLDAPCommand, crud.Search):
     member_param_incl_doc = _('Search for %(searched_object)s with these %(relationship)s %(ldap_object)s.')
     member_param_excl_doc = _('Search for %(searched_object)s without these %(relationship)s %(ldap_object)s.')
 
-    # pointer to function for entries sorting
-    # if no function is assigned the entries are sorted by their primary key value
-    entries_sortfn = None
+    # LDAPSearch sorts all matched records in the end using their primary key
+    # as a key attribute
+    # Set the following attribute to False to turn sorting off
+    sort_result_entries = True
 
     takes_options = (
         Int('timelimit?',
@@ -1846,12 +1847,10 @@ class LDAPSearch(BaseLDAPCommand, crud.Search):
             else:
                 callback(self, ldap, entries, truncated, *args, **options)
 
-        if not self.entries_sortfn:
+        if self.sort_result_entries:
             if self.obj.primary_key:
                 sortfn=lambda x,y: cmp(x[1][self.obj.primary_key.name][0].lower(), y[1][self.obj.primary_key.name][0].lower())
                 entries.sort(sortfn)
-        else:
-            entries.sort(self.entries_sortfn)
 
         if not options.get('raw', False):
             for e in entries:
diff --git a/ipalib/plugins/pwpolicy.py b/ipalib/plugins/pwpolicy.py
index 90a2ea110..330f9f7e5 100644
--- a/ipalib/plugins/pwpolicy.py
+++ b/ipalib/plugins/pwpolicy.py
@@ -456,6 +456,9 @@ api.register(pwpolicy_show)
 class pwpolicy_find(LDAPSearch):
     __doc__ = _('Search for group password policies.')
 
+    # this command does custom sorting in post_callback
+    sort_result_entries = False
+
     def sort_priority(self,x,y):
         # global policy will be always last in the output
         if x[1]['cn'][0] == global_policy_name:
@@ -464,17 +467,34 @@ class pwpolicy_find(LDAPSearch):
             return -1
         else:
             # policies with higher priority will be at the beginning of the list
-            return cmp(int(x[1]['cospriority'][0]), int(y[1]['cospriority'][0]))
-
-    entries_sortfn = sort_priority
+            try:
+                x =  cmp(int(x[1]['cospriority'][0]), int(y[1]['cospriority'][0]))
+            except KeyError:
+                # if cospriority is not present in the entry, rather return 0
+                # than crash
+                x = 0
+            return x
 
     def post_callback(self, ldap, entries, truncated, *args, **options):
-        if options.get('pkey_only', False):
-            return False
         for e in entries:
             # attribute rights are not allowed for pwpolicy_find
             self.obj.add_cospriority(e[1], e[1]['cn'][0], rights=False)
-            self.obj.convert_time_for_output(e[1], **options)
+            if options.get('pkey_only', False):
+                # when pkey_only flag is on, entries should contain only a cn
+                # and a cospriority attribute that will be used for sorting
+                # When the entries are sorted, cosentry is removed
+                self.obj.convert_time_for_output(e[1], **options)
+
+        # do custom entry sorting by its cospriority
+        entries.sort(self.sort_priority)
+
+        if options.get('pkey_only', False):
+            # remove cospriority that was used for sorting
+            for e in entries:
+                try:
+                    del e[1]['cospriority']
+                except KeyError:
+                    pass
 
 api.register(pwpolicy_find)