Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Preload private_key in JwtTokenSource #519

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Preload private_key in JwtTokenSource #519

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion ydb/oauth2_token_exchange/token_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,15 @@

try:
import jwt
import jwt.utils
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 +54,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"
Comment on lines 56 to +57
Copy link
Preview

Copilot AI Mar 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using assert for dependency checks might be bypassed in optimized mode; consider raising an explicit exception to ensure the cryptography library is available in production.

Suggested change
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"
if jwt is None:
raise ImportError("Install pyjwt library to use jwt tokens")
if load_pem_private_key is None:
raise ImportError("Install cryptography library to use jwt tokens")

Copilot is powered by AI, so mistakes are possible. Review output carefully before use.

self._signing_method = signing_method
self._key_id = key_id
if private_key and private_key_file:
Expand All @@ -57,7 +64,7 @@ def __init__(
self._private_key = private_key
if private_key_file:
private_key_file = os.path.expanduser(private_key_file)
with open(private_key_file, "r") as key_file:
with open(private_key_file, "rb") as key_file:
self._private_key = key_file.read()
self._issuer = issuer
self._subject = subject
Expand All @@ -70,6 +77,10 @@ def __init__(
raise Exception("JWT: no private key specified")
if self._token_ttl_seconds <= 0:
raise Exception("JWT: invalid jwt token TTL")
if isinstance(self._private_key, str):
self._private_key = self._private_key.encode()
if isinstance(self._private_key, bytes) and jwt.utils.is_pem_format(self._private_key):
self._private_key = load_pem_private_key(self._private_key, password=None)

def token(self) -> Token:
now = time.time()
Expand Down
Loading