#!/usr/bin/env python

import unittest
import os

from py11 import *
from py11.mutex import MutexDB

p11 = None

so_pin    = "fnord"
user_pin  = "fnord"
only_slot = 0

class TestInit(unittest.TestCase):

  def test_no_lock(self):
    p11.C_Initialize()

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

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

  def test_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):

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

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

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

  def test_slots(self):
    self.assertEqual(p11.C_GetSlotList(), [only_slot])

  def test_serial_sessions(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_parallel_sessions(self):
    # Cooked API doesn't allow the user to make this mistake, so we
    # have to use the raw API to test it.
    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, user_pin)
    p11.C_CloseSession(ro_session)
    p11.C_Login(rw_session, CKU_SO, user_pin)
    self.assertRaises(CKR_SESSION_READ_WRITE_SO_EXISTS, p11.C_OpenSession, only_slot, CKF_SERIAL_SESSION)
    p11.C_Logout(rw_session)

def setUpModule():
  from os import environ
  from os.path import abspath
  global p11
  p11 = PKCS11("./libpkcs11.so")
  environ["PKCS11_DATABASE"] = abspath("unit_tests.db")
  delete_db()
  set_pin()

def tearDownModule():
  delete_db()

def delete_db():
  from os import environ, unlink
  print "Deleting database", environ["PKCS11_DATABASE"]
  try:
    unlink(environ["PKCS11_DATABASE"])
  except OSError, e:
    if e.errno != 2:
      raise

def set_pin():
  from subprocess import Popen, PIPE
  print "Creating database and setting PINs"
  Popen(("./p11util", "-sup"), stdin = PIPE).communicate("%s\n%s\n" % (so_pin, user_pin))

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