HSMAUTH: change type annotation in _password_to_key

Yubico · Aug 9, 2023 · 3ef0b0c · 3ef0b0c
1 parent 58ce91e
commit 3ef0b0c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 14 deletions.
diff --git a/tests/test_hsmauth.py b/tests/test_hsmauth.py
@@ -51,18 +51,15 @@ def test_password_to_key(self):
             a2b_hex("592fd483f759e29909a04c4505d2ce0a"),
         ) == _password_to_key("password")
 
-        assert (
-            a2b_hex("090b47dbed595654901dee1cc655e420"),
-            a2b_hex("592fd483f759e29909a04c4505d2ce0a"),
-        ) == _password_to_key(b"password")
-
     def test__password_to_key_utf8(self):
         assert (
             a2b_hex("f320972c667ba5cd4d35119a6b0271a1"),
             a2b_hex("f10050ca688e5a6ce62b1ffb0f6f6869"),
         ) == _password_to_key("κόσμε")
 
-        assert (
-            a2b_hex("f320972c667ba5cd4d35119a6b0271a1"),
-            a2b_hex("f10050ca688e5a6ce62b1ffb0f6f6869"),
-        ) == _password_to_key(a2b_hex("cebae1bdb9cf83cebcceb5"))
+    def test_password_to_key_bytes_fails(self):
+        with pytest.raises(AttributeError):
+            _password_to_key(b"password")
+
+        with pytest.raises(AttributeError):
+            _password_to_key(a2b_hex("cebae1bdb9cf83cebcceb5"))
diff --git a/yubikit/hsmauth.py b/yubikit/hsmauth.py
@@ -142,15 +142,13 @@ def _parse_select(response):
     return Version.from_bytes(data)
 
 
-def _password_to_key(password: Union[str, bytes]) -> Tuple[bytes, bytes]:
+def _password_to_key(password: str) -> Tuple[bytes, bytes]:
     """Derive encryption and MAC key from a password.
 
     :return: A tuple containing the encryption key, and MAC key.
     """
-    if isinstance(password, str):
-        pw_bytes = password.encode()
-    else:
-        pw_bytes = password
+    pw_bytes = password.encode()
+
     key = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
         length=32,