Preload private_key in JwtTokenSource

Author: Roman Tretiak

ydb/oauth2_token_exchange/token_source.py

@@ -10,6 +10,11 @@
 except ImportError:
     jwt = None
 
+try:
+    from cryptography.hazmat.primitives.serialization import load_pem_private_key
+except ImportError:
+    load_pem_private_key = None
+
 
 class Token(abc.ABC):
     def __init__(self, token: str, token_type: str):
@@ -48,6 +53,7 @@ def __init__(
         token_ttl_seconds: int = 3600,
     ):
         assert jwt is not None, "Install pyjwt library to use jwt tokens"
+        assert load_pem_private_key is not None, "Install cryptography library to use jwt tokens"
         self._signing_method = signing_method
         self._key_id = key_id
         if private_key and private_key_file:
@@ -70,6 +76,7 @@ def __init__(
             raise Exception("JWT: no private key specified")
         if self._token_ttl_seconds <= 0:
             raise Exception("JWT: invalid jwt token TTL")
+        self._loaded_private_key = load_pem_private_key(self._private_key.encode(), password=None)
 
     def token(self) -> Token:
         now = time.time()
@@ -96,7 +103,7 @@ def token(self) -> Token:
             headers["kid"] = self._key_id
 
         token = jwt.encode(
-            key=self._private_key,
+            key=self._loaded_private_key,
             algorithm=self._signing_method,
             headers=headers,
             payload=payload,