add AES in counter mode support

PyKCS11/__init__.py

@@ -812,6 +812,29 @@ def to_native(self):
         return self._mech
 
 
+class AES_CTR_Mechanism(object):
+    """CKM_AES_CTR encryption mechanism"""
+
+    def __init__(self, counterBits, counterBlock):
+        """
+        :param counterBits: the number of incremented bits in the counter block
+        :param counterBlock: a 16-byte initial value of the counter block
+        """
+        self._param = PyKCS11.LowLevel.CK_AES_CTR_PARAMS()
+
+        self._source_cb = ckbytelist(counterBlock)
+        self._param.ulCounterBits = counterBits
+        self._param.cb = self._source_cb
+
+        self._mech = PyKCS11.LowLevel.CK_MECHANISM()
+        self._mech.mechanism = CKM_AES_CTR
+        self._mech.pParameter = self._param
+        self._mech.ulParameterLen = PyKCS11.LowLevel.CK_AES_CTR_PARAMS_LENGTH
+
+    def to_native(self):
+        return self._mech
+
+
 class RSAOAEPMechanism(object):
     """RSA OAEP Wrapping mechanism"""
 

src/opensc/pkcs11.h

@@ -756,6 +756,11 @@ struct ck_gcm_params {
   unsigned long ulTagBits;
 } ;
 
+struct ck_aes_ctr_params {
+  unsigned long ulCounterBits;
+  unsigned char cb[16];
+} ;
+
 struct ck_ecdh1_derive_params {
   unsigned long kdf;
   unsigned long ulSharedDataLen;
@@ -1335,6 +1340,9 @@ typedef struct ck_rsa_pkcs_pss_params *CK_RSA_PKCS_PSS_PARAMS_PTR;
 
 typedef struct ck_gcm_params CK_GCM_PARAMS;
 
+typedef struct ck_aes_ctr_params CK_AES_CTR_PARAMS;
+typedef struct ck_aes_ctr_params *CK_AES_CTR_PARAMS_PTR;
+
 typedef struct ck_ecdh1_derive_params CK_ECDH1_DERIVE_PARAMS;
 typedef struct ck_ecdh1_derive_params *CK_ECDH1_DERIVE_PARAMS_PTR;
 

src/pykcs11.i

@@ -245,23 +245,61 @@ typedef struct CK_DATE{
     }
     else
     {
-      // If the value being set is of CK_RSA_PKCS_OAEP_PARAMS type:
-      int res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_RSA_PKCS_OAEP_PARAMS*), 0);
-      if (!SWIG_IsOK(res2)) {
-          res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_RSA_PKCS_PSS_PARAMS*), 0);
-          if (!SWIG_IsOK(res2)) {
+        // If the value isn't a ckbytelist, then it must be a pointer to a mechanism parameter
+        int res2 = -1;
+        do { // Add mechanism parameters here
+            res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_RSA_PKCS_OAEP_PARAMS*), 0);
+            if( SWIG_IsOK( res2 ) )
+                break;
+
+            res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_RSA_PKCS_PSS_PARAMS*), 0);
+            if( SWIG_IsOK( res2 ) )
+                break;
+
             res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_GCM_PARAMS*), 0);
-            if (!SWIG_IsOK(res2)) {
-              res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_ECDH1_DERIVE_PARAMS*), 0);
-              if (!SWIG_IsOK(res2)) {
-                SWIG_exception_fail(SWIG_ArgError(res2), "unsupported CK_MECHANISM Parameter type.");
-              }
-            }
-          }
-      }
+            if( SWIG_IsOK( res2 ) )
+                break;
+
+            res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_ECDH1_DERIVE_PARAMS*), 0);
+            if( SWIG_IsOK( res2 ) )
+                break;
+
+            res2 = SWIG_ConvertPtr($input, &arg2, $descriptor(CK_AES_CTR_PARAMS*), 0);
+            if( SWIG_IsOK( res2 ) )
+                break;
+        } while(0);
+
+        if (!SWIG_IsOK(res2)) {
+            SWIG_exception_fail(SWIG_ArgError(res2), "unsupported CK_MECHANISM Parameter type.");
+        }
     }
 }
 
+// typemap for CK_BYTE static arrays
+%typemap(in) unsigned char[ANY](unsigned char out[$1_dim0]) {
+    vector<unsigned char> *vect;
+    // Expect a value of ckbytelist type:
+    int res = SWIG_ConvertPtr($input, (void **)&vect, $descriptor(vector<unsigned char> *), 0);
+    if (SWIG_IsOK(res))
+    {
+        if (vect->size() != $1_dim0)
+        {
+            SWIG_exception_fail(SWIG_ValueError, "Expected a ckbytelist with $1_dim0 elements");
+        }
+
+        for (size_t i = 0; i < $1_dim0; i++)
+        {
+            out[i] = (*vect)[i];
+        }
+    }
+    else
+    {
+        // If a mechanism parameter has a CK_BYTE array as a member, it must be represented as a ckbytelist
+        SWIG_exception_fail(SWIG_ArgError(res), "CK_BYTE arrays of CK_* mechanism params must be represented as ckbytelist type");
+    }
+    $1 = &out[0];
+}
+
 typedef struct CK_MECHANISM {
   unsigned long mechanism;
   void*       pParameter;
@@ -318,6 +356,24 @@ typedef struct CK_GCM_PARAMS {
 
 %constant int CK_GCM_PARAMS_LENGTH = sizeof(CK_GCM_PARAMS);
 
+typedef struct CK_AES_CTR_PARAMS {
+    unsigned long ulCounterBits;
+    unsigned char cb[16];
+} CK_AES_CTR_PARAMS;
+
+%extend CK_AES_CTR_PARAMS
+{
+    CK_AES_CTR_PARAMS()
+    {
+        CK_AES_CTR_PARAMS *p = new CK_AES_CTR_PARAMS();
+        p->ulCounterBits = 128;
+        memset(p->cb, 0, sizeof(p->cb));
+        return p;
+    }
+};
+
+%constant int CK_AES_CTR_PARAMS_LENGTH = sizeof(CK_AES_CTR_PARAMS);
+
 typedef struct CK_RSA_PKCS_OAEP_PARAMS {
   unsigned long hashAlg;
   unsigned long mgf;

test/test_symetric.py

@@ -93,6 +93,17 @@ def test_symetric(self):
         # 2nd block
         self.assertSequenceEqual(DataOut[16:], DataECBOut2)
 
+        # AES CTR with IV
+        mechanism = PyKCS11.AES_CTR_Mechanism(128, "1234567812345678")
+
+        DataOut = self.session.encrypt(symKey, DataIn, mechanism)
+        # print("DataOut", DataOut)
+
+        DataCheck = self.session.decrypt(symKey, DataOut, mechanism)
+        # print("DataCheck:", DataCheck)
+
+        self.assertSequenceEqual(DataIn, DataCheck)
+
         #
         # test CK_GCM_PARAMS
         #