Better Error Handling for Missing Entity Keys When Using Single Purpose Keys (#154)

* Updated JOSE auth decrypt error message to accurately present correct error when entity keys are not found.
- changes to step to match new error message
- added new EntityKeyNotFound error
- logic changes to account for missing entity keys and raising new error

Co-authored-by: Branden Jordan <[email protected]>

Author: jnpblrdrgz

Pipfile.lock

@@ -1,7 +1,7 @@
 from launchkey import LAUNCHKEY_PRODUCTION
 from launchkey.factories import OrganizationFactory
 from behave import given, when, then
-from hamcrest import assert_that, equal_to
+from hamcrest import assert_that, equal_to, string_contains_in_order
 
 from managers import DirectoryManager
 
@@ -64,7 +64,8 @@ def attempt_basic_api_call(context):
 
 @then("no valid key will be available to decrypt response")
 def no_key_exists_to_decrypt_response(context):
-    assert_that(str(context.current_exception), equal_to("Incorrect decryption."),
-                "Expected an Incorrect Decryption error but received a different error.")
 
-    context.execute_steps("Then a ValueError error occurs")
+    assert_that(str(context.current_exception),
+                string_contains_in_order("The key id:", "could not be found in the entities available keys."))
+
+    context.execute_steps("Then a EntityKeyNotFound error occurs")

features/steps/key_steps.py

@@ -165,6 +165,11 @@ class NoIssuerKey(LaunchKeyAPIException):
     """Issuer key was not loaded"""
 
 
+class EntityKeyNotFound(LaunchKeyAPIException):
+    """The key id with this current response could not be found within the
+    entities available keys"""
+
+
 class InvalidJWTResponse(LaunchKeyAPIException):
     """JWT Response is not in a valid format"""
 

launchkey/exceptions/__init__.py

@@ -23,7 +23,7 @@
 from ..exceptions import InvalidEntityID, InvalidPrivateKey, \
     InvalidIssuer, InvalidAlgorithm, LaunchKeyAPIException, \
     JWTValidationFailure, UnexpectedAPIResponse, NoIssuerKey, \
-    InvalidJWTResponse, UnexpectedKeyID
+    InvalidJWTResponse, UnexpectedKeyID, EntityKeyNotFound
 from .. import VALID_JWT_ISSUER_LIST, API_CACHE_TIME, \
     JOSE_SUPPORTED_CONTENT_HASH_ALGS, JOSE_SUPPORTED_JWE_ALGS, \
     JOSE_SUPPORTED_JWE_ENCS, JOSE_SUPPORTED_JWT_ALGS, \
@@ -486,11 +486,17 @@ def decrypt_response(self, response):
         :return: Decrypted string
         """
         package = JWEnc().unpack(response)
-        keys = list(self.issuer_private_keys)
+        keys = list()
         if 'kid' in package.headers:
             for key in self.issuer_private_keys:
                 if key.kid == package.headers['kid']:
                     keys = [key]
+            if len(keys) == 0:
+                raise EntityKeyNotFound("The key id: %s could not be found in "
+                                        "the entities available keys."
+                                        % package.headers["kid"])
+        else:
+            keys = list(self.issuer_private_keys)
         return JWE().decrypt(response, keys=keys).decode('utf-8')
 
     def decrypt_rsa_response(self, response, key_id):

launchkey/transports/jose_auth.py

@@ -12,7 +12,7 @@
 from launchkey.exceptions import InvalidAlgorithm, UnexpectedAPIResponse, \
     InvalidEntityID, InvalidIssuer, InvalidPrivateKey, NoIssuerKey, \
     InvalidJWTResponse, JWTValidationFailure, LaunchKeyAPIException, \
-    UnexpectedKeyID
+    UnexpectedKeyID, EntityKeyNotFound
 from launchkey import JOSE_SUPPORTED_JWT_ALGS, JOSE_SUPPORTED_JWE_ALGS, JOSE_SUPPORTED_JWE_ENCS, \
     JOSE_SUPPORTED_CONTENT_HASH_ALGS, API_CACHE_TIME, VALID_JWT_ISSUER_LIST, JOSE_JWT_LEEWAY, SDK_VERSION
 
@@ -161,9 +161,9 @@ def setUp(self):
         self._transport.get.return_value = public_key
         self._transport._server_time_difference = 0, time()
 
-        self._jwt_patch = patch("launchkey.transports.jose_auth.JWT", return_value=MagicMock(spec=JWT)).start()
-        self._jwt_patch.return_value.unpack.return_value.headers = faux_jwt_headers
-
+        self._jwenc_patch = patch("launchkey.transports.jose_auth.JWEnc",
+                                  return_value=MagicMock(spec=JWT)).start()
+        self._jwenc_patch.return_value.unpack.return_value.headers = faux_jwt_headers
         self.addCleanup(patch.stopall)
 
     def _encrypt_decrypt(self):
@@ -192,6 +192,28 @@ def test_supported_jwe_encryptions_success(self):
             self._transport.jwe_claims_encryption = enc
             self._encrypt_decrypt()
 
+    def test_no_valid_keys_returns_error(self):
+        with self.assertRaises(EntityKeyNotFound):
+            self._jwenc_patch.return_value.unpack.return_value.headers = \
+                {"alg": "RS512",
+                 "typ": "JWT",
+                 "kid": "ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff:ff"
+                 }
+            self._encrypt_decrypt()
+
+    @patch("launchkey.transports.jose_auth.JWE")
+    def test_no_kid_in_headers_uses_all_keys_to_decrypt(self, jwe_patch):
+        self._jwenc_patch.return_value.unpack.return_value.headers = \
+            {"alg": "RS512",
+             "typ": "JWT",
+             }
+        jwe_patch.return_value.decrypt.return_value = b'{"tobe": "encrypted"}'
+        jwe_patch.return_value.encrypt.return_value = \
+            "this.value.is.encrypted.correctly"
+        self._encrypt_decrypt()
+        jwe_patch.return_value.decrypt.assert_called_once_with(
+            ANY, keys=self._transport.issuer_private_keys)
+
 
 class TestHashlibSupportedAlgs(unittest.TestCase):