SQUASHME: use a fixture for test rsa keys

Styria-Digital · Feb 18, 2020 · d99c71a · d99c71a
1 parent 1ee09fc
commit d99c71a
Showing 1 changed file with 41 additions and 38 deletions.
diff --git a/tests/views/test_authentication.py b/tests/views/test_authentication.py
@@ -7,7 +7,7 @@
 
 from jwt import get_unverified_header
 
-from pytest import skip
+from pytest import skip, fixture
 
 from rest_framework import status
 
@@ -17,28 +17,31 @@
 
 from sys import version_info
 
-try:
-    from cryptography.hazmat.backends import default_backend
-    from cryptography.hazmat.primitives.asymmetric import rsa
+@fixture
+def rsa_keys(scope="session"):
+    try:
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives.asymmetric import rsa
 
-    # DO NOT copy key_size=512 for any real code!
-    srsa1 = rsa.generate_private_key(
-        public_exponent=65537,
-        key_size=512,
-        backend=default_backend()
-    )
+        nkeys = 2
+        secret = {}
+        public = {}
 
-    srsa2 = rsa.generate_private_key(
-        public_exponent=65537,
-        key_size=512,
-        backend=default_backend()
-    )
-    prsa1 = srsa1.public_key()
-    prsa2 = srsa2.public_key()
+        for i in range(1, nkeys):
+            name = "rsa%d" % i
+            secret[name] = rsa.generate_private_key(
+                public_exponent=65537,
+                # key_size=512 is probably unsafe for any real world code!
+                key_size=512,
+                backend=default_backend()
+            )
+            public[name] = secret[name].public_key()
 
-except:
-    pass
+        rsa_keys = {"secret": secret, "public": public}
+        return rsa_keys
 
+    except ImportError:
+        return None
 
 def test_empty_credentials_returns_validation_error(call_auth_endpoint):
     expected_output = {
@@ -208,19 +211,19 @@ def test_multi_keys_hash_hash(
 
 
 def test_multi_keys_rsa_rsa(
-    monkeypatch, user, call_auth_endpoint
+    monkeypatch, user, call_auth_endpoint, rsa_keys
 ):
 
-    try:
-        srsa1
-    except:
+    if not rsa_keys:
         skip("no rsa available")
         return
 
     monkeypatch.setattr(api_settings, "JWT_ALGORITHM", "RS256")
-    monkeypatch.setattr(api_settings, "JWT_PUBLIC_KEY", [prsa2, prsa1])
+    monkeypatch.setattr(
+        api_settings, "JWT_PUBLIC_KEY", list(rsa_keys["public"].values())
+    )
 
-    for skey in [srsa1, srsa2]:
+    for skey in rsa_keys["secret"].values():
         monkeypatch.setattr(api_settings, "JWT_PRIVATE_KEY", skey)
 
         response = call_auth_endpoint("username", "password")
@@ -234,17 +237,19 @@ def test_multi_keys_rsa_rsa(
 
 
 def test_multi_algo(
-    monkeypatch, user, call_auth_endpoint
+        monkeypatch, user, call_auth_endpoint, rsa_keys
 ):
 
-    try:
-        srsa1
-    except:
+    if not rsa_keys:
         skip("no rsa available")
         return
 
-    monkeypatch.setattr(api_settings, "JWT_PRIVATE_KEY", srsa1)
-    monkeypatch.setattr(api_settings, "JWT_PUBLIC_KEY", prsa1)
+    monkeypatch.setattr(
+        api_settings, "JWT_PRIVATE_KEY", rsa_keys["secret"]["rsa1"]
+    )
+    monkeypatch.setattr(
+        api_settings, "JWT_PUBLIC_KEY", rsa_keys["public"]["rsa1"]
+    )
 
     for algo in [["HS256", "RS256"], ["RS256", "HS256"]]:
         monkeypatch.setattr(api_settings, "JWT_ALGORITHM", algo)
@@ -263,21 +268,19 @@ def test_multi_algo(
 
 
 def test_kid(
-    monkeypatch, user, call_auth_endpoint
+    monkeypatch, user, call_auth_endpoint, rsa_keys
 ):
 
-    try:
-        srsa1
-    except:
+    if not rsa_keys:
         skip("no rsa available")
         return
 
-    monkeypatch.setattr(api_settings, "JWT_PRIVATE_KEY", { "rsa1": srsa1} )
     monkeypatch.setattr(
-        api_settings, "JWT_PUBLIC_KEY",
-        { "rsa1": prsa1, "rsa2": prsa2 }
+        api_settings, "JWT_PRIVATE_KEY", { "rsa1": rsa_keys["secret"]["rsa1"] }
     )
 
+    monkeypatch.setattr(api_settings, "JWT_PUBLIC_KEY", rsa_keys["public"])
+
     sk = { "hash1": "one", "hash2": "two" }
     # dicts are not ordered in python < 3.7
     if version_info < (3, 7):