Clean up Python APIs to C_FindObject*() and C_GetSlotList().

2 files changed, 20 insertions, 16 deletions

diff --git a/py11/__init__.py b/py11/__init__.py
index 80495d9..da79e34 100644
--- a/py11/__init__.py
+++ b/py11/__init__.py
@@ -95,10 +95,11 @@ class PKCS11 (object):
     self._C_Initialize_args = None
 
   def C_GetSlotList(self):
-    slots = (CK_SLOT_ID * 10)()
-    count = CK_ULONG(len(slots))
+    count = CK_ULONG()
+    self.so.C_GetSlotList(CK_TRUE, None, byref(count))
+    slots = (CK_SLOT_ID * count.value)()
     self.so.C_GetSlotList(CK_TRUE, slots, byref(count))
-    return [slots[i] for i in xrange(count.value)]
+    return tuple(slots[i] for i in xrange(count.value))
 
   def C_GetTokenInfo(self, slot_id):
     token_info = CK_TOKEN_INFO()
@@ -145,11 +146,11 @@ class PKCS11 (object):
       for i in xrange(count.value):
         yield objects[i]
 
-  def Find(self, session, template = None, **kwargs):
+  def FindObjects(self, session, template = None, **kwargs):
     self.C_FindObjectsInit(session, template, **kwargs)
-    for obj in self.C_FindObjects(session):
-      yield obj
+    result = tuple(self.C_FindObjects(session))
     self.C_FindObjectsFinal(session)
+    return result
 
   def _parse_GenerateKeyPair_template(self,
                                       # Attributes common to public and private templates
diff --git a/unit_tests.py b/unit_tests.py
index cd47886..75f2602 100644
--- a/unit_tests.py
+++ b/unit_tests.py
@@ -50,7 +50,7 @@ class TestDevice(unittest.TestCase):
     p11.C_CloseAllSessions(only_slot)
 
   def test_getSlots(self):
-    self.assertEqual(p11.C_GetSlotList(), [only_slot])
+    self.assertEqual(p11.C_GetSlotList(), (only_slot,))
 
   def test_getTokenInfo(self):
     token_info = p11.C_GetTokenInfo(only_slot)
@@ -62,8 +62,7 @@ class TestDevice(unittest.TestCase):
     ro_session = p11.C_OpenSession(only_slot, CKF_SERIAL_SESSION)
 
   def test_sessions_parallel(self):
-    # Cooked API doesn't allow the user to make this mistake, so we
-    # have to use the raw API to test it.
+    # Cooked API doesn't allow this mistake, must use raw API to test
     from ctypes import byref
     handle = CK_SESSION_HANDLE()
     notify = CK_NOTIFY()
@@ -96,6 +95,16 @@ class TestDevice(unittest.TestCase):
     self.assertIsInstance(random, str)
     self.assertEqual(len(random), n)
 
+  def test_findObjects(self):
+    session = p11.C_OpenSession(only_slot)
+    p11.C_FindObjectsInit(session, CKA_CLASS = CKO_PUBLIC_KEY)
+    with self.assertRaises(CKR_OPERATION_ACTIVE):
+      p11.C_FindObjectsInit(session, CKA_CLASS = CKO_PRIVATE_KEY)
+    for handle in p11.C_FindObjects(session):
+      self.assertIsInstance(handle, (int, long))
+    p11.C_FindObjectsFinal(session)
+
+
 class TestKeys(unittest.TestCase):
   """
   Tests involving keys.
@@ -118,7 +127,7 @@ class TestKeys(unittest.TestCase):
     p11.C_Login(self.session, CKU_USER, user_pin)
 
   def tearDown(self):
-    for handle in p11.Find(self.session):
+    for handle in p11.FindObjects(self.session):
       p11.C_DestroyObject(self.session, handle)
     p11.C_CloseAllSessions(only_slot)
     del self.session
@@ -223,9 +232,6 @@ class TestKeys(unittest.TestCase):
       CKA_LABEL         = "EC-P-256 test case from \"Suite B Implementer's Guide to FIPS 186-3\"",
       CKA_ID            = "EC-P-256",
       CKA_VERIFY        = True,
-      CKA_ENCRYPT       = False,
-      CKA_WRAP          = False,
-      CKA_TOKEN         = False,
       CKA_EC_POINT      = Q,
       CKA_EC_PARAMS     = self.oid_p256)
     p11.C_VerifyInit(self.session, CKM_ECDSA, handle)
@@ -244,9 +250,6 @@ class TestKeys(unittest.TestCase):
       CKA_LABEL         = "EC-P-384 test case from \"Suite B Implementer's Guide to FIPS 186-3\"",
       CKA_ID            = "EC-P-384",
       CKA_VERIFY        = True,
-      CKA_ENCRYPT       = False,
-      CKA_WRAP          = False,
-      CKA_TOKEN         = False,
       CKA_EC_POINT      = Q,
       CKA_EC_PARAMS     = self.oid_p384)
     p11.C_VerifyInit(self.session, CKM_ECDSA, handle)