From 287dd870832b5bd6caf8d2c609d8964f451b4e96 Mon Sep 17 00:00:00 2001
From: Max Wolfe <max@securitywolfe.com>
Date: Sat, 6 Feb 2021 19:07:34 -0800
Subject: [PATCH] Add Implementations and tests for C_InitToken and C_InitPin

Added as Token.init_token and Session.init_pin
---
 pkcs11/_pkcs11.pyx             | 46 ++++++++++++++++++++++++++++++++++
 tests/test_sessions.py         | 13 ++++++++++
 tests/test_slots_and_tokens.py | 20 +++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/pkcs11/_pkcs11.pyx b/pkcs11/_pkcs11.pyx
index 8e0260a..597586e 100644
--- a/pkcs11/_pkcs11.pyx
+++ b/pkcs11/_pkcs11.pyx
@@ -253,6 +253,31 @@ class Slot(types.Slot):
 class Token(types.Token):
     """Extend Token with implementation."""
 
+    def init_token(self, token_label, so_pin):
+        cdef CK_SLOT_ID slot_id = self.slot.slot_id
+        cdef CK_UTF8CHAR *pin_data
+        cdef CK_ULONG pin_length
+        cdef CK_UTF8CHAR *label
+
+        if token_label is None or so_pin is None:
+            raise ArgumentsBad("Set both `token_label` and `so_pin`")
+
+        pin = so_pin.encode('utf-8')
+        tlabel = token_label.encode('utf-8')
+
+        if pin and tlabel:
+            pin_data = pin
+            pin_length = len(pin)
+            label = tlabel
+
+            with nogil:
+                assertRV(_funclist.C_InitToken(slot_id, pin_data, pin_length,
+                                               label))
+
+            return True
+
+        return False
+
     def open(self, rw=False, user_pin=None, so_pin=None):
         cdef CK_SLOT_ID slot_id = self.slot.slot_id
         cdef CK_SESSION_HANDLE handle
@@ -373,6 +398,27 @@ def merge_templates(default_template, *user_templates):
 class Session(types.Session):
     """Extend Session with implementation."""
 
+    def init_pin(self, user_pin):
+        cdef CK_OBJECT_HANDLE handle = self._handle
+        cdef CK_UTF8CHAR *pin_data
+        cdef CK_ULONG pin_length
+
+        if user_pin is None:
+            raise ArgumentsBad("Set `user_pin`")
+
+        pin = user_pin.encode('utf-8')
+
+        if pin:
+            pin_data = pin
+            pin_length = len(pin)
+
+            with nogil:
+                assertRV(_funclist.C_InitPIN(handle, pin_data, pin_length))
+
+            return True
+
+        return False
+
     def close(self):
         cdef CK_OBJECT_HANDLE handle = self._handle
 
diff --git a/tests/test_sessions.py b/tests/test_sessions.py
index f9a24e7..a70dd7b 100644
--- a/tests/test_sessions.py
+++ b/tests/test_sessions.py
@@ -26,6 +26,19 @@ def test_open_session_and_login_so(self):
         with self.token.open(rw=True, so_pin=TOKEN_SO_PIN) as session:
             self.assertIsInstance(session, pkcs11.Session)
 
+    @Only.softhsm2  # We don't have credentials to do this for other platforms
+    def test_init_pin(self):
+        temp_token_pin = "bearsbeetsbattlestargalactica"
+
+        with self.token.open(rw=True, so_pin=TOKEN_SO_PIN) as session:
+            self.assertTrue(session.init_pin(temp_token_pin))
+
+        with self.token.open(user_pin=temp_token_pin) as session:
+            self.assertIsInstance(session, pkcs11.Session)
+
+        with self.token.open(rw=True, so_pin=TOKEN_SO_PIN) as session:
+            self.assertTrue(session.init_pin(TOKEN_PIN))
+
     @requires(pkcs11.Mechanism.AES_KEY_GEN)
     def test_generate_key(self):
         with self.token.open(user_pin=TOKEN_PIN) as session:
diff --git a/tests/test_slots_and_tokens.py b/tests/test_slots_and_tokens.py
index 140e126..f073ee7 100644
--- a/tests/test_slots_and_tokens.py
+++ b/tests/test_slots_and_tokens.py
@@ -72,3 +72,23 @@ def test_get_token(self):
         self.assertEqual(token.label, TOKEN)
         self.assertIn(pkcs11.TokenFlag.TOKEN_INITIALIZED, token.flags)
         self.assertIn(pkcs11.TokenFlag.LOGIN_REQUIRED, token.flags)
+
+    @Only.softhsm2
+    def test_init_token(self):
+        lib = pkcs11.lib(LIB)
+        tokens = lib.get_tokens()
+        temp_token_pin = "bearsbeetsbattlestargalactica"
+        temp_token_label = "schrute"
+
+        for token in tokens:
+            if pkcs11.TokenFlag.TOKEN_INITIALIZED not in token.flags:
+                self.assertTrue(token.init_token(temp_token_label,
+                                                 temp_token_pin))
+                break
+        else:
+            raise AssertionError("No Uninitialized token found")
+
+        token, *_ = lib.get_tokens(token_label=temp_token_label)
+
+        self.assertIn(pkcs11.TokenFlag.TOKEN_INITIALIZED, token.flags)
+        self.assertNotIn(pkcs11.TokenFlag.USER_PIN_INITIALIZED, token.flags)