Skip to content

Commit

Permalink
Preload private_key in JwtTokenSource
Browse files Browse the repository at this point in the history
  • Loading branch information
Roman Tretiak committed Nov 6, 2024
1 parent 92d7119 commit 25e1277
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion ydb/oauth2_token_exchange/token_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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()
Expand All @@ -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,
Expand Down

0 comments on commit 25e1277

Please sign in to comment.