From d4908be50cce0a1010822bd53fc157039818097e Mon Sep 17 00:00:00 2001
From: Jakob Schlyter <jakob@kirei.se>
Date: Tue, 12 Apr 2022 21:47:58 +0200
Subject: [PATCH] better tests set salt on password only

---
 src/cryptojwt/jwe/fernet.py | 11 +++++------
 tests/test_07_jwe.py        |  7 +++++++
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/cryptojwt/jwe/fernet.py b/src/cryptojwt/jwe/fernet.py
index 556ca42..90b02d3 100644
--- a/src/cryptojwt/jwe/fernet.py
+++ b/src/cryptojwt/jwe/fernet.py
@@ -17,17 +17,13 @@ class FernetEncrypter(Encrypter):
     def __init__(
         self,
         password: Optional[str] = None,
-        key: Optional[bytes] = None,
         salt: Optional[bytes] = "",
+        key: Optional[bytes] = None,
         hash_alg: Optional[str] = "SHA256",
         digest_size: Optional[int] = 0,
         iterations: Optional[int] = DEFAULT_ITERATIONS,
     ):
         Encrypter.__init__(self)
-        if not salt:
-            salt = os.urandom(16)
-        else:
-            salt = as_bytes(salt)
 
         if password is not None:
             _alg = getattr(hashes, hash_alg)
@@ -36,12 +32,15 @@ def __init__(
                 _algorithm = _alg(digest_size)
             else:
                 _algorithm = _alg()
+            salt = as_bytes(salt) if salt else os.urandom(16)
             kdf = PBKDF2HMAC(algorithm=_algorithm, length=32, salt=salt, iterations=iterations)
             self.key = base64.urlsafe_b64encode(kdf.derive(as_bytes(password)))
         elif key is not None:
+            if not isinstance(key, bytes):
+                raise TypeError("Raw key must be bytes")
             if len(key) != 32:
                 raise ValueError("Raw key must be 32 bytes")
-            self.key = base64.urlsafe_b64encode(as_bytes(key))
+            self.key = base64.urlsafe_b64encode(key)
         else:
             self.key = Fernet.generate_key()
 
diff --git a/tests/test_07_jwe.py b/tests/test_07_jwe.py
index 6cb94a7..82a3160 100644
--- a/tests/test_07_jwe.py
+++ b/tests/test_07_jwe.py
@@ -668,6 +668,13 @@ def test_fernet_symkey():
     assert resp == plain
 
 
+def test_fernet_bad():
+    with pytest.raises(TypeError):
+        encrypter = FernetEncrypter(key="xyzzy")
+    with pytest.raises(ValueError):
+        encrypter = FernetEncrypter(key=os.urandom(16))
+
+
 def test_fernet_bytes():
     key = os.urandom(32)