Add attribute database based on attributes.yaml.

Simplify prototype definitions and move them to separate module.

1 files changed, 13 insertions, 121 deletions

diff --git a/py11/__init__.py b/py11/__init__.py
index 7acb7aa..b67ce8a 100644
--- a/py11/__init__.py
+++ b/py11/__init__.py
@@ -6,120 +6,23 @@ from ctypes      import *
 from .exceptions import *
 from .types      import *
 from .constants  import *
+from .attributes import *
+from .prototypes import *
 
-# Prototypes for the PKCS #11 public functions.
-#
-# We don't bother implementing C_GetFunctionList(), because it would
-# be extremely tedious and it's not all that useful to us in Python.
-# Instead, we emulate its behavior in the PKCS11 class.
-
-PKCS11_Prototypes = (
-  ("C_Initialize",           [CK_VOID_PTR]),
-  ("C_Finalize",             [CK_VOID_PTR]),
-  ("C_GetInfo",              [CK_INFO_PTR]),
-# ("C_GetFunctionList",      [CK_FUNCTION_LIST_PTR_PTR]),
-  ("C_GetSlotList",          [CK_BBOOL, CK_SLOT_ID_PTR, CK_ULONG_PTR]),
-  ("C_GetSlotInfo",          [CK_SLOT_ID, CK_SLOT_INFO_PTR]),
-  ("C_GetTokenInfo",         [CK_SLOT_ID, CK_TOKEN_INFO_PTR]),
-  ("C_GetMechanismList",     [CK_SLOT_ID, CK_MECHANISM_TYPE_PTR, CK_ULONG_PTR]),
-  ("C_GetMechanismInfo",     [CK_SLOT_ID, CK_MECHANISM_TYPE, CK_MECHANISM_INFO_PTR]),
-  ("C_InitToken",            [CK_SLOT_ID, CK_UTF8CHAR_PTR, CK_ULONG, CK_UTF8CHAR_PTR]),
-  ("C_InitPIN",              [CK_SESSION_HANDLE, CK_UTF8CHAR_PTR, CK_ULONG]),
-  ("C_SetPIN",               [CK_SESSION_HANDLE, CK_UTF8CHAR_PTR, CK_ULONG, CK_UTF8CHAR_PTR, CK_ULONG]),
-  ("C_OpenSession",          [CK_SLOT_ID, CK_FLAGS, CK_VOID_PTR, CK_NOTIFY, CK_SESSION_HANDLE_PTR]),
-  ("C_CloseSession",         [CK_SESSION_HANDLE]),
-  ("C_CloseAllSessions",     [CK_SLOT_ID]),
-  ("C_GetSessionInfo",       [CK_SESSION_HANDLE, CK_SESSION_INFO_PTR]),
-  ("C_GetOperationState",    [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_SetOperationState",    [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_OBJECT_HANDLE, CK_OBJECT_HANDLE]),
-  ("C_Login",                [CK_SESSION_HANDLE, CK_USER_TYPE, CK_UTF8CHAR_PTR, CK_ULONG]),
-  ("C_Logout",               [CK_SESSION_HANDLE]),
-  ("C_CreateObject",         [CK_SESSION_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG, CK_OBJECT_HANDLE_PTR]),
-  ("C_CopyObject",           [CK_SESSION_HANDLE, CK_OBJECT_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG, CK_OBJECT_HANDLE_PTR]),
-  ("C_DestroyObject",        [CK_SESSION_HANDLE, CK_OBJECT_HANDLE]),
-  ("C_GetObjectSize",        [CK_SESSION_HANDLE, CK_OBJECT_HANDLE, CK_ULONG_PTR]),
-  ("C_GetAttributeValue",    [CK_SESSION_HANDLE, CK_OBJECT_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG]),
-  ("C_SetAttributeValue",    [CK_SESSION_HANDLE, CK_OBJECT_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG]),
-  ("C_FindObjectsInit",      [CK_SESSION_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG]),
-  ("C_FindObjects",          [CK_SESSION_HANDLE, CK_OBJECT_HANDLE_PTR, CK_ULONG, CK_ULONG_PTR]),
-  ("C_FindObjectsFinal",     [CK_SESSION_HANDLE]),
-  ("C_EncryptInit",          [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_Encrypt",              [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_EncryptUpdate",        [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_EncryptFinal",         [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DecryptInit",          [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_Decrypt",              [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DecryptUpdate",        [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DecryptFinal",         [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DigestInit",           [CK_SESSION_HANDLE, CK_MECHANISM_PTR]),
-  ("C_Digest",               [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DigestUpdate",         [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_DigestKey",            [CK_SESSION_HANDLE, CK_OBJECT_HANDLE]),
-  ("C_DigestFinal",          [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_SignInit",             [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_Sign",                 [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_SignUpdate",           [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_SignFinal",            [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_SignRecoverInit",      [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_SignRecover",          [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_VerifyInit",           [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_Verify",               [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG]),
-  ("C_VerifyUpdate",         [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_VerifyFinal",          [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_VerifyRecoverInit",    [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE]),
-  ("C_VerifyRecover",        [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DigestEncryptUpdate",  [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DecryptDigestUpdate",  [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_SignEncryptUpdate",    [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_DecryptVerifyUpdate",  [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG, CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_GenerateKey",          [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_ATTRIBUTE_PTR, CK_ULONG, CK_OBJECT_HANDLE_PTR]),
-  ("C_GenerateKeyPair",      [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_ATTRIBUTE_PTR, CK_ULONG,
-                              CK_ATTRIBUTE_PTR, CK_ULONG, CK_OBJECT_HANDLE_PTR, CK_OBJECT_HANDLE_PTR]),
-  ("C_WrapKey",              [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE, CK_OBJECT_HANDLE,
-                              CK_BYTE_PTR, CK_ULONG_PTR]),
-  ("C_UnwrapKey",            [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE, CK_BYTE_PTR, CK_ULONG,
-                              CK_ATTRIBUTE_PTR, CK_ULONG, CK_OBJECT_HANDLE_PTR]),
-  ("C_DeriveKey",            [CK_SESSION_HANDLE, CK_MECHANISM_PTR, CK_OBJECT_HANDLE, CK_ATTRIBUTE_PTR, CK_ULONG,
-                              CK_OBJECT_HANDLE_PTR]),
-  ("C_SeedRandom",           [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_GenerateRandom",       [CK_SESSION_HANDLE, CK_BYTE_PTR, CK_ULONG]),
-  ("C_GetFunctionStatus",    [CK_SESSION_HANDLE]),
-  ("C_CancelFunction",       [CK_SESSION_HANDLE]),
-  ("C_WaitForSlotEvent",     [CK_FLAGS, CK_SLOT_ID_PTR, CK_VOID_PTR]))
-
-
-# Need to convert attributes between a sane Pythonic representation
-# and something ctypes can handle.  Pythoninc representation would be
-# an iteraable of two element sequences, or perhaps a dict.
-
-def attributes_to_ctypes(dict_or_iterable):
-  if isinstance(dict_or_iterable, dict):
-    dict_or_iterable = dict_or_iterable.iteritems()
-  items = tuple(dict_or_iterable)
-  a = (CK_ATTRIBUTE * len(items))()
-  for i, kv in enumerate(items):
-    k, v = kv
-    if isinstance(v, bool):
-      v = pack("B", int(v))
-    elif isinstance(v, (int, long)):
-      v = pack("L", v)
-    a[i].type = k
-    a[i].pValue = create_string_buffer(v)
-    a[i].ulValueLen = len(v)
-  return a
 
 class PKCS11 (object):
 
-  def __init__(self, so_name = "libpkcs11.so"):
+  def __init__(self, so_name = "libpkcs11.so", attributes_initializer = "attributes.yaml"):
     self.so_name = so_name
     self.so = CDLL(so_name)
     def raise_on_failure(rv):
       if rv != CKR_OK:
         raise CKR_Exception.ckr_map[rv]
-    for name, args in PKCS11_Prototypes:
+    for name, args in Prototypes.iteritems():
       func = getattr(self.so, name)
       func.restype = raise_on_failure
       func.argtypes = args
+    self.adb = AttributeDB(attributes_initializer)
 
   def __getattr__(self, name):
     return getattr(self.so, name)
@@ -164,41 +67,30 @@ class PKCS11 (object):
     self.so.C_Login(session, user, pin, len(pin))
 
   def C_GetAttributeValue(self, session_handle, object_handle, attributes):
-    if not isinstance(attributes, (list, tuple)):
-      attributes = tuple(attributes)
-    if not all(isinstance(a, (int, long)) for a in attributes):
-      raise TypeError
-    template = (CK_ATTRIBUTE * len(attributes))()
-    for i in xrange(len(attributes)):
-      template[i].type = attributes[i]
-      template[i].pValue = None
-      template[i].ulValueLen = 0
+    template = self.adb.getvalue_create_template(attributes)
     self.so.C_GetAttributeValue(session_handle, object_handle, template, len(template))
-    for t in template:
-      t.pValue = create_string_buffer(t.ulValueLen)
+    self.adb.getvalue_allocate_template(template)
     self.so.C_GetAttributeValue(session_handle, object_handle, template, len(template))
-    return dict((t.type, t.pValue[:t.ulValueLen]) for t in template)
+    return self.adb.from_ctypes(template)
 
   def C_FindObjectsInit(self, session, template):
     if template:
-      self.so.C_FindObjectsInit(session, attributes_to_ctypes(template), len(template))
+      self.so.C_FindObjectsInit(session, self.adb.to_ctypes(template), len(template))
     else:
       self.so.C_FindObjectsInit(session, None, 0)
 
   def C_FindObjects(self, session, chunk_size = 10):
     objects = (CK_OBJECT_HANDLE * chunk_size)()
-    count   = CK_ULONG()
-    while True:
+    count   = CK_ULONG(1)
+    while count.value > 0:
       self.so.C_FindObjects(session, objects, len(objects), byref(count))
       for i in xrange(count.value):
         yield objects[i]
-      if count.value == 0:
-        break
 
   def C_GenerateKeyPair(self, session, mechanism_type, public_template, private_template):
     mechanism = CK_MECHANISM(mechanism_type, None, 0)
-    public_template  = attributes_to_ctypes(public_template)
-    private_template = attributes_to_ctypes(private_template)
+    public_template  = self.adb.to_ctypes(public_template)
+    private_template = self.adb.to_ctypes(private_template)
     public_handle  = CK_OBJECT_HANDLE()
     private_handle = CK_OBJECT_HANDLE()
     self.so.C_GenerateKeyPair(session, byref(mechanism),