fixes some bugs and added functionality of fetching public key

guyp-descope · guyp-descope · commit 1a50cee4beea · 2022-05-19T23:22:09.000+03:00
diff --git a/descope/auth.py b/descope/auth.py
@@ -10,7 +10,9 @@
 
 from descope.common import (
     DEFAULT_BASE_URI,
+    DEFAULT_FETCH_PUBLIC_KEY_URI,
     EMAIL_REGEX,
+    GET_KEYS_PATH,
     PHONE_REGEX,
     SIGNIN_OTP_PATH,
     SIGNUP_OTP_PATH,
@@ -22,8 +24,7 @@
 
 
 class AuthClient:
-    def __init__(self, project_id: str, public_key: str):
-
+    def __init__(self, project_id: str, public_key: str = None):
         # validate project id
         if project_id is None or project_id == "":
             # try get the project_id from env
@@ -37,14 +38,15 @@ def __init__(self, project_id: str, public_key: str):
         self.project_id = project_id
 
         if public_key is None or public_key == "":
-            public_key = os.getenv("DESCOPE_PUBLIC_KEY", "")
-            if public_key == "":
-                raise AuthException(
-                    500,
-                    "Init failure",
-                    "Failed to init AuthClient object, public key cannot be found",
-                )
+            public_key = os.getenv("DESCOPE_PUBLIC_KEY", None)
+
+        if public_key is None:
+            self.public_key = None  # public key will be fetch later (on demand)
+        else:
+            self.public_key = self._validate_and_load_public_key(public_key)
 
+    @staticmethod
+    def _validate_and_load_public_key(public_key) -> jwt.PyJWK:
         if isinstance(public_key, str):
             try:
                 public_key = json.loads(public_key)
@@ -64,7 +66,7 @@ def __init__(self, project_id: str, public_key: str):
 
         try:
             # Load and validate public key
-            self.public_key = jwt.PyJWK(public_key)
+            return jwt.PyJWK(public_key)
         except jwt.InvalidKeyError as e:
             raise AuthException(
                 500,
@@ -78,6 +80,40 @@ def __init__(self, project_id: str, public_key: str):
                 f"Failed to init AuthClient object, failed to load public key {e}",
             )
 
+    def _fetch_public_key(self, kid: str) -> None:
+        response = requests.get(
+            f"{DEFAULT_FETCH_PUBLIC_KEY_URI}{GET_KEYS_PATH}/{self.project_id}",
+            headers=self._get_default_headers(),
+        )
+
+        if not response.ok:
+            raise AuthException(
+                401, "public key fetching failed", f"err: {response.reason}"
+            )
+
+        jwks_data = response.text
+        try:
+            jwkeys = json.loads(jwks_data)
+        except Exception as e:
+            raise AuthException(
+                401, "public key fetching failed", f"Failed to load jwks {e}"
+            )
+
+        founded_key = None
+        for key in jwkeys:
+            if key["kid"] == kid:
+                founded_key = key
+                break
+
+        if founded_key:
+            self.public_key = AuthClient._validate_and_load_public_key(founded_key)
+        else:
+            raise AuthException(
+                401,
+                "public key validation failed",
+                "Failed to validate public key, public key not found",
+            )
+
     @staticmethod
     def _verify_delivery_method(method: DeliveryMethod, identifier: str) -> bool:
         if identifier == "" or identifier is None:
@@ -219,23 +255,29 @@ def validate_session_request(self, signed_token):
         """
         DOC
         """
+
         try:
             unverified_header = jwt.get_unverified_header(signed_token)
         except Exception as e:
             raise AuthException(
                 401,
                 "token validation failure",
-                f"Failed to get unverified token header, {e}",
+                f"Failed to parse token header, {e}",
             )
-        token_type = unverified_header.get("typ", None)
-        alg = unverified_header.get("alg", None)
-        if token_type is None or alg is None:
+
+        kid = unverified_header.get("kid", None)
+        if kid is None:
             raise AuthException(
                 401,
                 "token validation failure",
-                f"Token header is missing token type or algorithm, token_type={token_type} alg={alg}",
+                "Token header is missing kid property",
             )
 
+        if self.public_key is None:
+            self._fetch_public_key(
+                kid
+            )  # will set self.public_key or raise exception if failed
+
         try:
             jwt.decode(jwt=signed_token, key=self.public_key.key, algorithms=["ES384"])
         except Exception as e:
diff --git a/descope/common.py b/descope/common.py
@@ -1,6 +1,7 @@
 from enum import Enum
 
 DEFAULT_BASE_URI = "http://localhost:8191"
+DEFAULT_FETCH_PUBLIC_KEY_URI = "http://localhost:8152"  # will use the same base uri as above once gateway will be available
 
 PHONE_REGEX = """^(?:(?:\\(?(?:00|\\+)([1-4]\\d\\d|[1-9]\\d?)\\)?)?[\\-\\.\\ \\\\/]?)?((?:\\(?\\d{1,}\\)?[\\-\\.\\ \\\\/]?){0,})(?:[\\-\\.\\ \\\\/]?(?:#|ext\\.?|extension|x)[\\-\\.\\ \\\\/]?(\\d+))?$"""
 # EMAIL_REGEX = """^(((([a-zA-Z]|\\d|[!#\\$%&'\\*\\+\\-\\/=\\?\\^_`{\\|}~]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])+(\\.([a-zA-Z]|\\d|[!#\\$%&'\\*\\+\\-\\/=\\?\\^_`{\\|}~]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])+)*)|((\\x22)((((\\x20|\\x09)*(\\x0d\\x0a))?(\\x20|\\x09)+)?(([\\x01-\\x08\\x0b\\x0c\\x0e-\\x1f\\x7f]|\\x21|[\\x23-\\x5b]|[\\x5d-\\x7e]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])|(\\([\\x01-\\x09\\x0b\\x0c\\x0d-\\x7f]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}]))))*(((\\x20|\\x09)*(\\x0d\\x0a))?(\\x20|\\x09)+)?(\\x22)))@((([a-zA-Z]|\\d|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])|(([a-zA-Z]|\\d|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])([a-zA-Z]|\\d|-|\\.|_|~|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])*([a-zA-Z]|\\d|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])))\\.)+(([a-zA-Z]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])|(([a-zA-Z]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])([a-zA-Z]|\\d|-|_|~|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])*([a-zA-Z]|[\\x{00A0}-\\x{D7FF}\\x{F900}-\\x{FDCF}\\x{FDF0}-\\x{FFEF}])))\\.?$"""
@@ -9,6 +10,7 @@
 SIGNIN_OTP_PATH = "/v1/auth/signin/otp"
 SIGNUP_OTP_PATH = "/v1/auth/signup/otp"
 VERIFY_CODE_PATH = "/v1/auth/code/verify"
+GET_KEYS_PATH = "/v1/keys"
 
 COOKIE_NAME = "S"
 
diff --git a/samples/sample_app.py b/samples/sample_app.py
@@ -37,16 +37,27 @@ def main():
             )
             logging.info("Code is valid")
             token = cookies.get(COOKIE_NAME)
-        except AuthException:
-            logging.info("Invalid code")
+        except AuthException as e:
+            logging.info(f"Invalid code {e}")
             raise
 
         try:
             logging.info("going to validate session..")
             auth_client.validate_session_request(token)
             logging.info("Session is valid and all is OK")
-        except AuthException:
-            logging.info("Session is not valid")
+        except AuthException as e:
+            logging.info(f"Session is not valid {e}")
+
+        try:
+            old_public_key = auth_client.public_key
+            # fetch and load the public key associated with this project (by kid)
+            auth_client._fetch_public_key(project_id)
+            if old_public_key != auth_client.public_key:
+                logging.info("new public key fetched successfully")
+            else:
+                logging.info("failed to fetch new public_key")
+        except AuthException as e:
+            logging.info(f"failed to fetch public key for this project {e}")
 
     except AuthException:
         raise
diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -25,17 +25,68 @@ def test_auth_client(self):
             AuthException, AuthClient, project_id=None, public_key="dummy"
         )
         self.assertRaises(AuthException, AuthClient, project_id="", public_key="dummy")
-        self.assertRaises(
-            AuthException, AuthClient, project_id="dummy", public_key=None
+        self.assertIsNotNone(
+            AuthException, AuthClient(project_id="dummy", public_key=None)
+        )
+        self.assertIsNotNone(
+            AuthException, AuthClient(project_id="dummy", public_key="")
         )
-        self.assertRaises(AuthException, AuthClient, project_id="dummy", public_key="")
         self.assertRaises(
             AuthException, AuthClient, project_id="dummy", public_key="not dict object"
         )
         self.assertIsNotNone(
             AuthClient(project_id="dummy", public_key=self.public_key_str)
         )
 
+    def test_validate_and_load_public_key(self):
+        # test invalid json
+        self.assertRaises(
+            AuthException,
+            AuthClient._validate_and_load_public_key,
+            public_key="invalid json",
+        )
+        # test not dict object
+        self.assertRaises(
+            AuthException, AuthClient._validate_and_load_public_key, public_key=555
+        )
+        # test invalid dict
+        self.assertRaises(
+            AuthException,
+            AuthClient._validate_and_load_public_key,
+            public_key={"kid": "dummy"},
+        )
+
+    def test_fetch_public_key(self):
+        client = AuthClient(self.dummy_project_id, self.public_key_dict)
+        valid_keys_response = """[
+    {
+        "alg": "ES384",
+        "crv": "P-384",
+        "kid": "299psneX92K3vpbqPMRCnbZKb27",
+        "kty": "EC",
+        "use": "sig",
+        "x": "435yhcD0tqH6z5M8kNFYEcEYXjzBQWiOvIOZO17rOatpXj-MbA6CKrktiblT4xMb",
+        "y": "YMf1EIz68z2_RKBys5byWRUXlqNF_BhO5F0SddkaRtiqZ8M6n7ZnKl65JGN0EEGr"
+    }
+]
+        """
+
+        # Test failed flows
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.ok = False
+            self.assertRaises(AuthException, client._fetch_public_key, "dummy_kid")
+
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.ok = True
+            mock_get.return_value.text = "invalid json"
+            self.assertRaises(AuthException, client._fetch_public_key, "dummy_kid")
+
+        # test success flow
+        with patch("requests.get") as mock_get:
+            mock_get.return_value.ok = True
+            mock_get.return_value.text = valid_keys_response
+            self.assertIsNone(client._fetch_public_key("299psneX92K3vpbqPMRCnbZKb27"))
+
     def test_verify_delivery_method(self):
         self.assertEqual(
             AuthClient._verify_delivery_method(DeliveryMethod.EMAIL, "dummy@dummy.com"),
