path: root/unit_tests.py

#!/usr/bin/env python

import unittest

from py11 import *
from py11.mutex import MutexDB

p11       = None
so_pin    = "fnord"
user_pin  = "fnord"
only_slot = 0
verbose   = True

class TestInit(unittest.TestCase):
  """
  Test all the flavors of C_Initialize().
  """

  def test_mutex_none(self):
    p11.C_Initialize()

  def test_mutex_os(self):
    p11.C_Initialize(CKF_OS_LOCKING_OK)

  def test_mutex_user(self):
    mdb = MutexDB()
    p11.C_Initialize(0, mdb.create, mdb.destroy, mdb.lock, mdb.unlock)

  def test_mutex_both(self):
    mdb = MutexDB()
    p11.C_Initialize(CKF_OS_LOCKING_OK, mdb.create, mdb.destroy, mdb.lock, mdb.unlock)

  def tearDown(self):
    p11.C_Finalize()

class TestDevice(unittest.TestCase):
  """
  Test basic device stuff like C_GetSlotList(), C_OpenSession(), and C_Login().
  """

  @classmethod
  def setUpClass(cls):
    p11.C_Initialize()

  @classmethod
  def tearDownClass(cls):
    p11.C_Finalize()

  def tearDown(self):
    p11.C_CloseAllSessions(only_slot)

  def test_getSlots(self):
    self.assertEqual(p11.C_GetSlotList(), (only_slot,))

  def test_getTokenInfo(self):
    token_info = p11.C_GetTokenInfo(only_slot)
    self.assertIsInstance(token_info, CK_TOKEN_INFO)
    self.assertEqual(token_info.label.rstrip(), "Cryptech Token")

  def test_sessions_serial(self):
    rw_session = p11.C_OpenSession(only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
    ro_session = p11.C_OpenSession(only_slot, CKF_SERIAL_SESSION)

  def test_sessions_parallel(self):
    # Cooked API doesn't allow this mistake, must use raw API to test
    from ctypes import byref
    handle = CK_SESSION_HANDLE()
    notify = CK_NOTIFY()
    with self.assertRaises(CKR_SESSION_PARALLEL_NOT_SUPPORTED):
      p11.so.C_OpenSession(only_slot, CKF_RW_SESSION, None, notify, byref(handle))
    with self.assertRaises(CKR_SESSION_PARALLEL_NOT_SUPPORTED):
      p11.so.C_OpenSession(only_slot, 0,              None, notify, byref(handle))

  def test_login_user(self):
    rw_session = p11.C_OpenSession(only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
    ro_session = p11.C_OpenSession(only_slot, CKF_SERIAL_SESSION)
    p11.C_Login(ro_session, CKU_USER, user_pin)
    p11.C_Logout(ro_session)

  def test_login_so(self):
    rw_session = p11.C_OpenSession(only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
    ro_session = p11.C_OpenSession(only_slot, CKF_SERIAL_SESSION)
    self.assertRaises(CKR_SESSION_READ_ONLY_EXISTS, p11.C_Login, ro_session, CKU_SO, so_pin)
    p11.C_CloseSession(ro_session)
    p11.C_Login(rw_session, CKU_SO, so_pin)
    self.assertRaises(CKR_SESSION_READ_WRITE_SO_EXISTS, p11.C_OpenSession, only_slot, CKF_SERIAL_SESSION)
    p11.C_Logout(rw_session)

  def test_random(self):
    # Testing that what this produces really is random seems beyond
    # the scope of a unit test.
    session = p11.C_OpenSession(only_slot)
    n = 17
    random = p11.C_GenerateRandom(session, n)
    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.
  """

  oid_p256 = "".join(chr(i) for i in (0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07))
  oid_p384 = "".join(chr(i) for i in (0x06, 0x05, 0x2b, 0x81, 0x04, 0x00, 0x22))
  oid_p521 = "".join(chr(i) for i in (0x06, 0x05, 0x2b, 0x81, 0x04, 0x00, 0x23))

  @classmethod
  def setUpClass(cls):
    p11.C_Initialize()

  @classmethod
  def tearDownClass(cls):
    p11.C_Finalize()

  def setUp(self):
    self.session = p11.C_OpenSession(only_slot)
    p11.C_Login(self.session, CKU_USER, user_pin)

  def tearDown(self):
    for handle in p11.FindObjects(self.session):
      p11.C_DestroyObject(self.session, handle)
    p11.C_CloseAllSessions(only_slot)
    del self.session

  def assertIsKeypair(self, public_handle, private_handle = None):
    if isinstance(public_handle, tuple) and private_handle is None:
      public_handle, private_handle = public_handle
    self.assertEqual(p11.C_GetAttributeValue(self.session, public_handle,  CKA_CLASS), {CKA_CLASS: CKO_PUBLIC_KEY})
    self.assertEqual(p11.C_GetAttributeValue(self.session, private_handle, CKA_CLASS), {CKA_CLASS: CKO_PRIVATE_KEY})

  def test_keygen_token_vs_session(self):
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN, CKA_TOKEN = False,
                            CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                            CKA_SIGN = True, CKA_VERIFY = True))
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN, CKA_TOKEN = True,
                            CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                            CKA_SIGN = True, CKA_VERIFY = True))
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                            public_CKA_TOKEN = False, private_CKA_TOKEN = True,
                            CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                            CKA_SIGN = True, CKA_VERIFY = True))
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                            public_CKA_TOKEN = True, private_CKA_TOKEN = False,
                            CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                            CKA_SIGN = True, CKA_VERIFY = True))

  def test_gen_sign_verify_ecdsa_p256_sha256(self):
    public_key, private_key = p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                                    CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                                                    CKA_SIGN = True, CKA_VERIFY = True)
    self.assertIsKeypair(public_key, private_key)
    hamster = "Your mother was a hamster"
    p11.C_SignInit(self.session, CKM_ECDSA_SHA256, private_key)
    sig = p11.C_Sign(self.session, hamster)
    self.assertIsInstance(sig, str)
    p11.C_VerifyInit(self.session, CKM_ECDSA_SHA256, public_key)
    p11.C_Verify(self.session, hamster, sig)

  @unittest.skip("SHA-384 not available in current build")
  def test_gen_sign_verify_ecdsa_p384_sha384(self):
    public_key, private_key = p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                                    CKA_ID = "EC-P384", CKA_EC_PARAMS = self.oid_p384,
                                                    CKA_SIGN = True, CKA_VERIFY = True)
    self.assertIsKeypair(public_key, private_key)
    hamster = "Your mother was a hamster"
    p11.C_SignInit(self.session, CKM_ECDSA_SHA384, private_key)
    sig = p11.C_Sign(self.session, hamster)
    self.assertIsInstance(sig, str)
    p11.C_VerifyInit(self.session, CKM_ECDSA_SHA384, public_key)
    p11.C_Verify(self.session, hamster, sig)

  @unittest.skip("SHA-512 not available in current build")
  def test_gen_sign_verify_ecdsa_p521_sha512(self):
    public_key, private_key = p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                                    CKA_ID = "EC-P521", CKA_EC_PARAMS = self.oid_p521,
                                                    CKA_SIGN = True, CKA_VERIFY = True)
    self.assertIsKeypair(public_key, private_key)
    hamster = "Your mother was a hamster"
    p11.C_SignInit(self.session, CKM_ECDSA_SHA512, private_key)
    sig = p11.C_Sign(self.session, hamster)
    self.assertIsInstance(sig, str)
    p11.C_VerifyInit(self.session, CKM_ECDSA_SHA512, public_key)
    p11.C_Verify(self.session, hamster, sig)

  def test_gen_rsa_1024(self):
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_RSA_PKCS_KEY_PAIR_GEN, CKA_MODULUS_BITS = 1024,
                            CKA_ID = "RSA-1024", CKA_SIGN = True, CKA_VERIFY = True))

  @unittest.skip("RSA key generation is still painfully slow")
  def test_gen_rsa_2048(self):
    self.assertIsKeypair(
      p11.C_GenerateKeyPair(self.session, CKM_RSA_PKCS_KEY_PAIR_GEN, CKA_MODULUS_BITS = 2048,
                            CKA_ID = "RSA-1024", CKA_SIGN = True, CKA_VERIFY = True))

  @staticmethod
  def _build_ecpoint(x, y):
    bytes_per_coordinate = (max(x.bit_length(), y.bit_length()) + 15) / 16
    value = chr(0x04) + ("%0*x%0*x" % (bytes_per_coordinate, x, bytes_per_coordinate, y)).decode("hex")
    if len(value) < 128:
      length = chr(len(value))
    else:
      n = len(value).bit_length()
      length = chr((n + 7) / 8) + ("%0*x" % ((n + 15) / 16, len(value))).decode("hex")
    tag = chr(0x04)
    return tag + length + value

  def test_canned_ecdsa_p256_verify(self):
    Q = self._build_ecpoint(0x8101ece47464a6ead70cf69a6e2bd3d88691a3262d22cba4f7635eaff26680a8,
                            0xd8a12ba61d599235f67d9cb4d58f1783d3ca43e78f0a5abaa624079936c0c3a9)
    H = "7c3e883ddc8bd688f96eac5e9324222c8f30f9d6bb59e9c5f020bd39ba2b8377".decode("hex")
    r = "7214bc9647160bbd39ff2f80533f5dc6ddd70ddf86bb815661e805d5d4e6f27c".decode("hex")
    s = "7d1ff961980f961bdaa3233b6209f4013317d3e3f9e1493592dbeaa1af2bc367".decode("hex")
    handle = p11.C_CreateObject(
      session           = self.session,
      CKA_CLASS         = CKO_PUBLIC_KEY,
      CKA_KEY_TYPE      = CKK_EC,
      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_EC_POINT      = Q,
      CKA_EC_PARAMS     = self.oid_p256)
    p11.C_VerifyInit(self.session, CKM_ECDSA, handle)
    p11.C_Verify(self.session, H, r + s)

  def test_canned_ecdsa_p384_verify(self):
    Q = self._build_ecpoint(0x1fbac8eebd0cbf35640b39efe0808dd774debff20a2a329e91713baf7d7f3c3e81546d883730bee7e48678f857b02ca0,
                            0xeb213103bd68ce343365a8a4c3d4555fa385f5330203bdd76ffad1f3affb95751c132007e1b240353cb0a4cf1693bdf9)
    H = "b9210c9d7e20897ab86597266a9d5077e8db1b06f7220ed6ee75bd8b45db37891f8ba5550304004159f4453dc5b3f5a1".decode("hex")
    r = "a0c27ec893092dea1e1bd2ccfed3cf945c8134ed0c9f81311a0f4a05942db8dbed8dd59f267471d5462aa14fe72de856".decode("hex")
    s = "20ab3f45b74f10b6e11f96a2c8eb694d206b9dda86d3c7e331c26b22c987b7537726577667adadf168ebbe803794a402".decode("hex")
    handle = p11.C_CreateObject(
      session           = self.session,
      CKA_CLASS         = CKO_PUBLIC_KEY,
      CKA_KEY_TYPE      = CKK_EC,
      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_EC_POINT      = Q,
      CKA_EC_PARAMS     = self.oid_p384)
    p11.C_VerifyInit(self.session, CKM_ECDSA, handle)
    p11.C_Verify(self.session, H, r + s)




def setUpModule():
  global p11
  p11 = PKCS11("./libpkcs11.so")
  import os, subprocess
  if verbose:
    print "Initializing database"
  db = os.path.abspath("unit_tests.db")
  if os.path.exists(db):
    os.unlink(db)
  os.environ["PKCS11_DATABASE"] = db
  subprocess.Popen(("./p11util", "-sup"), stdin = subprocess.PIPE).communicate("%s\n%s\n" % (so_pin, user_pin))
  if verbose:
    print "Setup complete"

def tearDownModule():
  import os
  os.unlink(os.environ["PKCS11_DATABASE"])

if __name__ == "__main__":
  unittest.main(verbosity = 2 if verbose else 1)