More key unit tests.

1 files changed, 33 insertions, 12 deletions

diff --git a/py11/__init__.py b/py11/__init__.py
index 9d3837e..80495d9 100644
--- a/py11/__init__.py
+++ b/py11/__init__.py
@@ -119,14 +119,19 @@ class PKCS11 (object):
   def C_Login(self, session, user, pin):
     self.so.C_Login(session, user, pin, len(pin))
 
-  def C_GetAttributeValue(self, session_handle, object_handle, attributes):
+  def C_GetAttributeValue(self, session_handle, object_handle, *attributes):
+    if len(attributes) == 1 and isinstance(attributes[0], (list, tuple)):
+      attributes = attributes[0]
     template = self.adb.getvalue_create_template(attributes)
     self.so.C_GetAttributeValue(session_handle, object_handle, template, len(template))
     self.adb.getvalue_allocate_template(template)
     self.so.C_GetAttributeValue(session_handle, object_handle, template, len(template))
     return self.adb.from_ctypes(template)
 
-  def C_FindObjectsInit(self, session, template):
+  def C_FindObjectsInit(self, session, template = None, **kwargs):
+    if kwargs:
+      assert not template
+      template = kwargs
     if template:
       self.so.C_FindObjectsInit(session, self.adb.to_ctypes(template), len(template))
     else:
@@ -140,6 +145,12 @@ class PKCS11 (object):
       for i in xrange(count.value):
         yield objects[i]
 
+  def Find(self, session, template = None, **kwargs):
+    self.C_FindObjectsInit(session, template, **kwargs)
+    for obj in self.C_FindObjects(session):
+      yield obj
+    self.C_FindObjectsFinal(session)
+
   def _parse_GenerateKeyPair_template(self,
                                       # Attributes common to public and private templates
                                       CKA_ID,
@@ -159,19 +170,26 @@ class PKCS11 (object):
                                       **kwargs):
     if CKA_LABEL is None:
       CKA_LABEL = CKA_ID
-    public = dict(kwargs, CKA_LABEL = CKA_LABEL, CKA_ID = CKA_ID, CKA_TOKEN = public_CKA_TOKEN or CKA_TOKEN)
-    private = dict(CKA_LABEL = CKA_LABEL, CKA_ID = CKA_ID, CKA_TOKEN = private_CKA_TOKEN or CKA_TOKEN,
-                   CKA_SIGN = CKA_SIGN, CKA_DECRYPT = CKA_DECRYPT, CKA_UNWRAP = CKA_UNWRAP,
-                   CKA_SENSITIVE = CKA_SENSITIVE, CKA_PRIVATE = CKA_PRIVATE, CKA_EXTRACTABLE = CKA_EXTRACTABLE)
-    return self.adb.to_ctypes(public), self.adb.to_ctypes(private)
+    return (dict(kwargs,
+                 CKA_LABEL      = CKA_LABEL,
+                 CKA_ID         = CKA_ID,
+                 CKA_TOKEN      = public_CKA_TOKEN or CKA_TOKEN),
+            dict(CKA_LABEL      = CKA_LABEL,
+                 CKA_ID         = CKA_ID,
+                 CKA_TOKEN      = private_CKA_TOKEN or CKA_TOKEN,
+                 CKA_SIGN       = CKA_SIGN,
+                 CKA_DECRYPT    = CKA_DECRYPT,
+                 CKA_UNWRAP     = CKA_UNWRAP,
+                 CKA_SENSITIVE  = CKA_SENSITIVE,
+                 CKA_PRIVATE    = CKA_PRIVATE,
+                 CKA_EXTRACTABLE = CKA_EXTRACTABLE))
 
   def C_GenerateKeyPair(self, session, mechanism_type, public_template = None, private_template = None, **kwargs):
-    assert kwargs or (public_template and private_template)
     if kwargs:
+      assert not public_template and not private_template
       public_template, private_template = self._parse_GenerateKeyPair_template(**kwargs)
-    else:
-      public_template  = self.adb.to_ctypes(public_template)
-      private_template = self.adb.to_ctypes(private_template)
+    public_template  = self.adb.to_ctypes(public_template)
+    private_template = self.adb.to_ctypes(private_template)
     mechanism = CK_MECHANISM(mechanism_type, None, 0)
     public_handle  = CK_OBJECT_HANDLE()
     private_handle = CK_OBJECT_HANDLE()
@@ -215,7 +233,10 @@ class PKCS11 (object):
   def C_VerifyFinal(self, session, signature):
     self.so.C_VerifyFinal(session, signature, len(signature))
 
-  def C_CreateObject(self, session, template):
+  def C_CreateObject(self, session, template = None, **kwargs):
+    if kwargs:
+      assert not template
+      template = kwargs
     template = self.adb.to_ctypes(template)
     handle = CK_OBJECT_HANDLE()
     self.so.C_CreateObject(session, template, len(template), byref(handle))