From bda5fcbd7cf330ecce8c624344da8fcd0a476e98 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Fri, 15 Mar 2024 16:26:36 +0000 Subject: [PATCH 01/12] OAuth2JwtTokenExchangeCredentials --- ydb/aio/iam.py | 11 ++---- ydb/iam/auth.py | 96 ++++++++++++++++++++++++++++++++++++++----------- 2 files changed, 78 insertions(+), 29 deletions(-) diff --git a/ydb/aio/iam.py b/ydb/aio/iam.py index eab8faff..341dde2e 100644 --- a/ydb/aio/iam.py +++ b/ydb/aio/iam.py @@ -65,17 +65,10 @@ def __init__( iam_channel_credentials=None, ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) - auth.BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key) + auth.BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, "PS256", auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL) def _get_token_request(self): - return iam_token_service_pb2.CreateIamTokenRequest( - jwt=auth.get_jwt( - self._account_id, - self._access_key_id, - self._private_key, - self._jwt_expiration_timeout, - ) - ) + return iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) class YandexPassportOAuthIamCredentials(TokenServiceCredentials): diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index 82e7c9f6..dc313d1e 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -from ydb import credentials, tracing +from ydb import credentials, tracing, issues import grpc import time import abc @@ -23,22 +23,27 @@ DEFAULT_METADATA_URL = "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token" +YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL = "https://iam.api.cloud.yandex.net/iam/v1/tokens" +NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL = "https://token-service.iam.new.nebiuscloud.net" -def get_jwt(account_id, access_key_id, private_key, jwt_expiration_timeout): +def get_jwt(account_id, access_key_id, private_key, jwt_expiration_timeout, algorithm, token_service_url, subject=None): now = time.time() now_utc = datetime.utcfromtimestamp(now) exp_utc = datetime.utcfromtimestamp(now + jwt_expiration_timeout) + payload = { + "iss": account_id, + "aud": token_service_url, + "iat": now_utc, + "exp": exp_utc, + } + if subject is not None: + payload["sub"] = subject return jwt.encode( key=private_key, - algorithm="PS256", - headers={"typ": "JWT", "alg": "PS256", "kid": access_key_id}, - payload={ - "iss": account_id, - "aud": "https://iam.api.cloud.yandex.net/iam/v1/tokens", - "iat": now_utc, - "exp": exp_utc, - }, + algorithm=algorithm, + headers={"typ": "JWT", "alg": algorithm, "kid": access_key_id}, + payload=payload, ) @@ -73,12 +78,16 @@ def _make_token_request(self): class BaseJWTCredentials(abc.ABC): - def __init__(self, account_id, access_key_id, private_key): + def __init__(self, account_id, access_key_id, private_key, algorithm, token_service_url, subject=None): self._account_id = account_id self._jwt_expiration_timeout = 60.0 * 60 self._token_expiration_timeout = 120 self._access_key_id = access_key_id self._private_key = private_key + self._algorithm = algorithm + self._token_service_url = token_service_url + self._subject = subject + def set_token_expiration_timeout(self, value): self._token_expiration_timeout = value @@ -99,6 +108,48 @@ def from_file(cls, key_file, iam_endpoint=None, iam_channel_credentials=None): iam_channel_credentials=iam_channel_credentials, ) + def _get_jwt(self): + return get_jwt( + self._account_id, + self._access_key_id, + self._private_key, + self._jwt_expiration_timeout, + self._algorithm, + self._token_service_url, + self._subject + ) + + +class OAuth2JwtTokenExchangeCredentials(credentials.AbstractExpiringTokenCredentials, BaseJWTCredentials): + def __init__(self, token_exchange_url, account_id, access_key_id, private_key, algorithm, token_service_url, subject=None, tracer=None): + BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, algorithm, token_service_url, subject) + super(OAuth2JwtTokenExchangeCredentials, self).__init__(tracer) + assert requests is not None, "Install requests library to use OAuth 2.0 token exchange credentials provider" + self._token_exchange_url = token_exchange_url + + def _process_response(self, response): + if response.status_code == 403: + raise issues.Unauthenticated(response.content) + if response.status_code >= 500: + raise issues.Unavailable(response.content) + if response.status_code >= 400: + raise issues.BadRequest(response.content) + if response.status_code != 200: + raise issues.Error(response.content) + + @tracing.with_trace() + def _make_token_request(self): + params = { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", + "subject_token": self._get_jwt(), + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + } + headers = {"Content-Type": "application/x-www-form-urlencoded"} + response = requests.post(self._token_exchange_url, data=params, headers=headers) + self._process_response(response) + return {"access_token": response.content, "expires_in": 3600} # TODO + class JWTIamCredentials(TokenServiceCredentials, BaseJWTCredentials): def __init__( @@ -110,17 +161,22 @@ def __init__( iam_channel_credentials=None, ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) - BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key) + BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, "PS256", YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL) def _get_token_request(self): - return self._iam_token_service_pb2.CreateIamTokenRequest( - jwt=get_jwt( - self._account_id, - self._access_key_id, - self._private_key, - self._jwt_expiration_timeout, - ) - ) + return self._iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) + + +def NebiusJWTIamCredentials(OAuth2JwtTokenExchangeCredentials): + def __init__( + self, + account_id, + access_key_id, + private_key, + iam_endpoint=None, + iam_channel_credentials=None, + ): + OAuth2JwtTokenExchangeCredentials.__init__(self, iam_endpoint, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id) class YandexPassportOAuthIamCredentials(TokenServiceCredentials): From fa41446f3f00a2f71aeb7633b3f97b8769ae0a71 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Thu, 21 Mar 2024 13:34:46 +0000 Subject: [PATCH 02/12] NebiusServiceAccountCredentials --- ydb/driver.py | 7 +++++++ ydb/iam/__init__.py | 1 + ydb/iam/auth.py | 37 +++++++++++++++++++++++++++++++------ 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/ydb/driver.py b/ydb/driver.py index 89109b9b..16bba151 100644 --- a/ydb/driver.py +++ b/ydb/driver.py @@ -38,6 +38,13 @@ def credentials_from_env_variables(tracer=None): return ydb.iam.ServiceAccountCredentials.from_file(service_account_key_file) + nebius_service_account_key_file = os.getenv("YDB_NEBIUS_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS") + if nebius_service_account_key_file is not None: + ctx.trace({"credentials.nebius_service_account_key_file": True}) + import ydb.iam + + return ydb.iam.NebiusServiceAccountCredentials.from_file(nebius_service_account_key_file) + anonymous_credetials = os.getenv("YDB_ANONYMOUS_CREDENTIALS", "0") == "1" if anonymous_credetials: ctx.trace({"credentials.anonymous": True}) diff --git a/ydb/iam/__init__.py b/ydb/iam/__init__.py index 7167efe1..cf835769 100644 --- a/ydb/iam/__init__.py +++ b/ydb/iam/__init__.py @@ -1,3 +1,4 @@ # -*- coding: utf-8 -*- from .auth import ServiceAccountCredentials # noqa +from .auth import NebiusServiceAccountCredentials # noqa from .auth import MetadataUrlCredentials # noqa diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index dc313d1e..e96ce2e5 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -127,7 +127,7 @@ def __init__(self, token_exchange_url, account_id, access_key_id, private_key, a assert requests is not None, "Install requests library to use OAuth 2.0 token exchange credentials provider" self._token_exchange_url = token_exchange_url - def _process_response(self, response): + def _process_response_status_code(self, response): if response.status_code == 403: raise issues.Unauthenticated(response.content) if response.status_code >= 500: @@ -137,6 +137,13 @@ def _process_response(self, response): if response.status_code != 200: raise issues.Error(response.content) + def _process_response(self, response): + self._process_response_status_code(response) + response_json = json.loads(response.content) + access_token = response_json["access_token"] + expires_in = response_json["expires_in"] + return {"access_token": access_token, "expires_in": expires_in} + @tracing.with_trace() def _make_token_request(self): params = { @@ -147,8 +154,7 @@ def _make_token_request(self): } headers = {"Content-Type": "application/x-www-form-urlencoded"} response = requests.post(self._token_exchange_url, data=params, headers=headers) - self._process_response(response) - return {"access_token": response.content, "expires_in": 3600} # TODO + return self._process_response(response) class JWTIamCredentials(TokenServiceCredentials, BaseJWTCredentials): @@ -173,10 +179,12 @@ def __init__( account_id, access_key_id, private_key, - iam_endpoint=None, - iam_channel_credentials=None, + token_exchange_url=None, ): - OAuth2JwtTokenExchangeCredentials.__init__(self, iam_endpoint, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id) + url = token_exchange_url + if url is None: + url = "https://auth.new.nebiuscloud.net/oauth2/token/exchange" + OAuth2JwtTokenExchangeCredentials.__init__(self, url, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id) class YandexPassportOAuthIamCredentials(TokenServiceCredentials): @@ -232,3 +240,20 @@ def __init__( iam_endpoint, iam_channel_credentials, ) + + +class NebiusServiceAccountCredentials(NebiusJWTIamCredentials): + def __init__( + self, + service_account_id, + access_key_id, + private_key, + iam_endpoint=None, + iam_channel_credentials=None, + ): + super(NebiusServiceAccountCredentials, self).__init__( + service_account_id, + access_key_id, + private_key, + iam_endpoint, + ) From 35056e58b4b9607471ad49b109c03cf2354cd162 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Fri, 22 Mar 2024 19:37:47 +0000 Subject: [PATCH 03/12] Fix style --- ydb/iam/auth.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index e96ce2e5..b87eb8ff 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -88,7 +88,6 @@ def __init__(self, account_id, access_key_id, private_key, algorithm, token_serv self._token_service_url = token_service_url self._subject = subject - def set_token_expiration_timeout(self, value): self._token_expiration_timeout = value return self @@ -173,7 +172,7 @@ def _get_token_request(self): return self._iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) -def NebiusJWTIamCredentials(OAuth2JwtTokenExchangeCredentials): +class NebiusJWTIamCredentials(OAuth2JwtTokenExchangeCredentials): def __init__( self, account_id, From f0136a6c191fcd71eb752a6f208f4e043d010c07 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Fri, 22 Mar 2024 19:40:23 +0000 Subject: [PATCH 04/12] Format code --- ydb/aio/iam.py | 4 +++- ydb/iam/auth.py | 22 ++++++++++++++++++---- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/ydb/aio/iam.py b/ydb/aio/iam.py index 341dde2e..16a024d2 100644 --- a/ydb/aio/iam.py +++ b/ydb/aio/iam.py @@ -65,7 +65,9 @@ def __init__( iam_channel_credentials=None, ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) - auth.BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, "PS256", auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL) + auth.BaseJWTCredentials.__init__( + self, account_id, access_key_id, private_key, "PS256", auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL + ) def _get_token_request(self): return iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index b87eb8ff..86545f45 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -115,12 +115,22 @@ def _get_jwt(self): self._jwt_expiration_timeout, self._algorithm, self._token_service_url, - self._subject + self._subject, ) class OAuth2JwtTokenExchangeCredentials(credentials.AbstractExpiringTokenCredentials, BaseJWTCredentials): - def __init__(self, token_exchange_url, account_id, access_key_id, private_key, algorithm, token_service_url, subject=None, tracer=None): + def __init__( + self, + token_exchange_url, + account_id, + access_key_id, + private_key, + algorithm, + token_service_url, + subject=None, + tracer=None, + ): BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, algorithm, token_service_url, subject) super(OAuth2JwtTokenExchangeCredentials, self).__init__(tracer) assert requests is not None, "Install requests library to use OAuth 2.0 token exchange credentials provider" @@ -166,7 +176,9 @@ def __init__( iam_channel_credentials=None, ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) - BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, "PS256", YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL) + BaseJWTCredentials.__init__( + self, account_id, access_key_id, private_key, "PS256", YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL + ) def _get_token_request(self): return self._iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) @@ -183,7 +195,9 @@ def __init__( url = token_exchange_url if url is None: url = "https://auth.new.nebiuscloud.net/oauth2/token/exchange" - OAuth2JwtTokenExchangeCredentials.__init__(self, url, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id) + OAuth2JwtTokenExchangeCredentials.__init__( + self, url, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id + ) class YandexPassportOAuthIamCredentials(TokenServiceCredentials): From 4681c11c49858bb52ae32de51951f88d094f3901 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 26 Mar 2024 14:03:48 +0000 Subject: [PATCH 05/12] NebiusServiceAccountCredentials --- ydb/aio/iam.py | 91 +++++++++++++++++++++++++++++++++++++++++++++++-- ydb/iam/auth.py | 20 +++++++---- 2 files changed, 102 insertions(+), 9 deletions(-) diff --git a/ydb/aio/iam.py b/ydb/aio/iam.py index 16a024d2..8439050d 100644 --- a/ydb/aio/iam.py +++ b/ydb/aio/iam.py @@ -1,5 +1,6 @@ import grpc.aio import time +import json import abc import logging @@ -9,11 +10,14 @@ logger = logging.getLogger(__name__) try: - from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc - from yandex.cloud.iam.v1 import iam_token_service_pb2 import jwt except ImportError: jwt = None + +try: + from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc + from yandex.cloud.iam.v1 import iam_token_service_pb2 +except ImportError: iam_token_service_pb2_grpc = None iam_token_service_pb2 = None @@ -55,6 +59,54 @@ async def _make_token_request(self): IamTokenCredentials = TokenServiceCredentials +class OAuth2JwtTokenExchangeCredentials(AbstractExpiringTokenCredentials, auth.BaseJWTCredentials): + def __init__( + self, + token_exchange_url, + account_id, + access_key_id, + private_key, + algorithm, + token_service_url, + subject=None, + ): + super(OAuth2JwtTokenExchangeCredentials, self).__init__() + auth.BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, algorithm, token_service_url, subject) + assert aiohttp is not None, "Install aiohttp library to use OAuth 2.0 token exchange credentials provider" + self._token_exchange_url = token_exchange_url + + def _process_response_status_code(self, response): + if response.status_code == 403: + raise issues.Unauthenticated(response.content) + if response.status_code >= 500: + raise issues.Unavailable(response.content) + if response.status_code >= 400: + raise issues.BadRequest(response.content) + if response.status_code != 200: + raise issues.Error(response.content) + + def _process_response(self, response): + self._process_response_status_code(response) + response_json = json.loads(response.content) + access_token = response_json["access_token"] + expires_in = response_json["expires_in"] + return {"access_token": access_token, "expires_in": expires_in} + + async def _make_token_request(self): + params = { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", + "subject_token": self._get_jwt(), + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + } + headers = {"Content-Type": "application/x-www-form-urlencoded"} + + timeout = aiohttp.ClientTimeout(total=2) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(self._token_exchange_url, data=params, headers=headers) as response: + return self._process_response(response) + + class JWTIamCredentials(TokenServiceCredentials, auth.BaseJWTCredentials): def __init__( self, @@ -66,13 +118,29 @@ def __init__( ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) auth.BaseJWTCredentials.__init__( - self, account_id, access_key_id, private_key, "PS256", auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL + self, account_id, access_key_id, private_key, auth.YANDEX_CLOUD_JWT_ALGORITHM, auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL ) def _get_token_request(self): return iam_token_service_pb2.CreateIamTokenRequest(jwt=self._get_jwt()) +class NebiusJWTIamCredentials(OAuth2JwtTokenExchangeCredentials): + def __init__( + self, + account_id, + access_key_id, + private_key, + token_exchange_url=None, + ): + url = token_exchange_url + if url is None: + url = auth.NEBIUS_CLOUD_IAM_TOKEN_EXCHANGE_URL + OAuth2JwtTokenExchangeCredentials.__init__( + self, url, account_id, access_key_id, private_key, auth.NEBIUS_CLOUD_JWT_ALGORITHM, auth.NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, account_id + ) + + class YandexPassportOAuthIamCredentials(TokenServiceCredentials): def __init__( self, @@ -125,3 +193,20 @@ def __init__( iam_endpoint, iam_channel_credentials, ) + + +class NebiusServiceAccountCredentials(NebiusJWTIamCredentials): + def __init__( + self, + service_account_id, + access_key_id, + private_key, + iam_endpoint=None, + iam_channel_credentials=None, + ): + super(NebiusServiceAccountCredentials, self).__init__( + service_account_id, + access_key_id, + private_key, + iam_endpoint, + ) diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index 86545f45..901e60ca 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -8,11 +8,14 @@ import os try: - from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc - from yandex.cloud.iam.v1 import iam_token_service_pb2 import jwt except ImportError: jwt = None + +try: + from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc + from yandex.cloud.iam.v1 import iam_token_service_pb2 +except ImportError: iam_token_service_pb2_grpc = None iam_token_service_pb2 = None @@ -24,10 +27,15 @@ DEFAULT_METADATA_URL = "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token" YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL = "https://iam.api.cloud.yandex.net/iam/v1/tokens" -NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL = "https://token-service.iam.new.nebiuscloud.net" +NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE = "token-service.iam.new.nebiuscloud.net" +NEBIUS_CLOUD_IAM_TOKEN_EXCHANGE_URL = "https://auth.new.nebiuscloud.net/oauth2/token/exchange" + +YANDEX_CLOUD_JWT_ALGORITHM = "PS256" +NEBIUS_CLOUD_JWT_ALGORITHM = "RS256" def get_jwt(account_id, access_key_id, private_key, jwt_expiration_timeout, algorithm, token_service_url, subject=None): + assert jwt is not None, "Install pyjwt library to use jwt tokens" now = time.time() now_utc = datetime.utcfromtimestamp(now) exp_utc = datetime.utcfromtimestamp(now + jwt_expiration_timeout) @@ -177,7 +185,7 @@ def __init__( ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) BaseJWTCredentials.__init__( - self, account_id, access_key_id, private_key, "PS256", YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL + self, account_id, access_key_id, private_key, YANDEX_CLOUD_JWT_ALGORITHM, YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL ) def _get_token_request(self): @@ -194,9 +202,9 @@ def __init__( ): url = token_exchange_url if url is None: - url = "https://auth.new.nebiuscloud.net/oauth2/token/exchange" + url = NEBIUS_CLOUD_IAM_TOKEN_EXCHANGE_URL OAuth2JwtTokenExchangeCredentials.__init__( - self, url, account_id, access_key_id, private_key, "RS256", NEBIUS_CLOUD_IAM_TOKEN_SERVICE_URL, account_id + self, url, account_id, access_key_id, private_key, NEBIUS_CLOUD_JWT_ALGORITHM, NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, account_id ) From 58a6ac9c8dcec58a428aec3842a5c0e8d3871a37 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 26 Mar 2024 15:45:49 +0000 Subject: [PATCH 06/12] Async NebiusServiceAccountCredentials --- ydb/aio/iam.py | 32 ++++++++++++++------------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/ydb/aio/iam.py b/ydb/aio/iam.py index 8439050d..b9f3f69b 100644 --- a/ydb/aio/iam.py +++ b/ydb/aio/iam.py @@ -6,6 +6,7 @@ import logging from ydb.iam import auth from .credentials import AbstractExpiringTokenCredentials +from ydb import issues logger = logging.getLogger(__name__) @@ -75,23 +76,6 @@ def __init__( assert aiohttp is not None, "Install aiohttp library to use OAuth 2.0 token exchange credentials provider" self._token_exchange_url = token_exchange_url - def _process_response_status_code(self, response): - if response.status_code == 403: - raise issues.Unauthenticated(response.content) - if response.status_code >= 500: - raise issues.Unavailable(response.content) - if response.status_code >= 400: - raise issues.BadRequest(response.content) - if response.status_code != 200: - raise issues.Error(response.content) - - def _process_response(self, response): - self._process_response_status_code(response) - response_json = json.loads(response.content) - access_token = response_json["access_token"] - expires_in = response_json["expires_in"] - return {"access_token": access_token, "expires_in": expires_in} - async def _make_token_request(self): params = { "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", @@ -104,7 +88,19 @@ async def _make_token_request(self): timeout = aiohttp.ClientTimeout(total=2) async with aiohttp.ClientSession(timeout=timeout) as session: async with session.post(self._token_exchange_url, data=params, headers=headers) as response: - return self._process_response(response) + if response.status == 403: + raise issues.Unauthenticated(await response.text()) + if response.status >= 500: + raise issues.Unavailable(await response.text()) + if response.status >= 400: + raise issues.BadRequest(await response.text()) + if response.status != 200: + raise issues.Error(await response.text()) + + response_json = await response.json() + access_token = response_json["access_token"] + expires_in = response_json["expires_in"] + return {"access_token": access_token, "expires_in": expires_in} class JWTIamCredentials(TokenServiceCredentials, auth.BaseJWTCredentials): From 8a36b2bf0a0633b07858e1777325e603adfc2600 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 26 Mar 2024 15:51:09 +0000 Subject: [PATCH 07/12] Style --- ydb/aio/iam.py | 21 +++++++++++++++++---- ydb/iam/auth.py | 9 ++++++++- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/ydb/aio/iam.py b/ydb/aio/iam.py index b9f3f69b..40622f8a 100644 --- a/ydb/aio/iam.py +++ b/ydb/aio/iam.py @@ -1,6 +1,5 @@ import grpc.aio import time -import json import abc import logging @@ -72,7 +71,9 @@ def __init__( subject=None, ): super(OAuth2JwtTokenExchangeCredentials, self).__init__() - auth.BaseJWTCredentials.__init__(self, account_id, access_key_id, private_key, algorithm, token_service_url, subject) + auth.BaseJWTCredentials.__init__( + self, account_id, access_key_id, private_key, algorithm, token_service_url, subject + ) assert aiohttp is not None, "Install aiohttp library to use OAuth 2.0 token exchange credentials provider" self._token_exchange_url = token_exchange_url @@ -114,7 +115,12 @@ def __init__( ): TokenServiceCredentials.__init__(self, iam_endpoint, iam_channel_credentials) auth.BaseJWTCredentials.__init__( - self, account_id, access_key_id, private_key, auth.YANDEX_CLOUD_JWT_ALGORITHM, auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL + self, + account_id, + access_key_id, + private_key, + auth.YANDEX_CLOUD_JWT_ALGORITHM, + auth.YANDEX_CLOUD_IAM_TOKEN_SERVICE_URL, ) def _get_token_request(self): @@ -133,7 +139,14 @@ def __init__( if url is None: url = auth.NEBIUS_CLOUD_IAM_TOKEN_EXCHANGE_URL OAuth2JwtTokenExchangeCredentials.__init__( - self, url, account_id, access_key_id, private_key, auth.NEBIUS_CLOUD_JWT_ALGORITHM, auth.NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, account_id + self, + url, + account_id, + access_key_id, + private_key, + auth.NEBIUS_CLOUD_JWT_ALGORITHM, + auth.NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, + account_id, ) diff --git a/ydb/iam/auth.py b/ydb/iam/auth.py index 901e60ca..852c0c28 100644 --- a/ydb/iam/auth.py +++ b/ydb/iam/auth.py @@ -204,7 +204,14 @@ def __init__( if url is None: url = NEBIUS_CLOUD_IAM_TOKEN_EXCHANGE_URL OAuth2JwtTokenExchangeCredentials.__init__( - self, url, account_id, access_key_id, private_key, NEBIUS_CLOUD_JWT_ALGORITHM, NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, account_id + self, + url, + account_id, + access_key_id, + private_key, + NEBIUS_CLOUD_JWT_ALGORITHM, + NEBIUS_CLOUD_IAM_TOKEN_SERVICE_AUDIENCE, + account_id, ) From 459b88800af0281c2e1abd80365746cc9237533d Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 2 Apr 2024 12:45:05 +0000 Subject: [PATCH 08/12] Test for Yandex jwt --- requirements.txt | 1 + test-requirements.txt | 1 + tests/auth/__init__.py | 0 tests/auth/test_credentials.py | 91 ++++++++++++++++++++++++++++++++++ tests/table/test_tx.py | 11 ---- 5 files changed, 93 insertions(+), 11 deletions(-) create mode 100644 tests/auth/__init__.py create mode 100644 tests/auth/test_credentials.py diff --git a/requirements.txt b/requirements.txt index 3b962207..8d364e41 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ grpcio>=1.42.0 packaging protobuf>=3.13.0,<5.0.0 aiohttp<4 +pyjwt==2.8.0 diff --git a/test-requirements.txt b/test-requirements.txt index 43c61a10..21da70d3 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -46,4 +46,5 @@ pylint-protobuf cython freezegun==1.2.2 pytest-cov +yandexcloud -e . diff --git a/tests/auth/__init__.py b/tests/auth/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py new file mode 100644 index 00000000..90b56e5f --- /dev/null +++ b/tests/auth/test_credentials.py @@ -0,0 +1,91 @@ +import jwt +import concurrent.futures +import grpc +import time + +import ydb.iam + +from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc +from yandex.cloud.iam.v1 import iam_token_service_pb2 + +SERVICE_ACCOUNT_ID = "sa_id" +ACCESS_KEY_ID = "key_id" +PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\nFgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\nM0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\nnPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\nJnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\nR+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\nWYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\nYLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\nuccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\nzrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\nsvlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\nm6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\nrheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\nLxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\nmto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\nJUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\nBjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\nZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\nkFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\nnXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\nFXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\nTP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\ncHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\nof1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\nhL8bFAuNNVrCOp79TNnNIsh7\n-----END PRIVATE KEY-----\n" +PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\nftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\nZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\ny4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\nJLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\njQIDAQAB\n-----END PUBLIC KEY-----\n" + + +def test_credentials(): + credentials = ydb.iam.MetadataUrlCredentials() + raised = False + try: + credentials.auth_metadata() + except Exception: + raised = True + + assert raised + + +class IamTokenServiceForTest(iam_token_service_pb2_grpc.IamTokenServiceServicer): + def Create(self, request, context): + print("IAM token service request: {}".format(request)) + # Validate jwt: + decoded = jwt.decode(request.jwt, key=PUBLIC_KEY, algorithms=["PS256"], audience="https://iam.api.cloud.yandex.net/iam/v1/tokens") + assert decoded["iss"] == SERVICE_ACCOUNT_ID + assert decoded["aud"] == "https://iam.api.cloud.yandex.net/iam/v1/tokens" + assert abs(decoded["iat"] - time.time()) <= 60 + assert abs(decoded["exp"] - time.time()) <= 3600 + + response = iam_token_service_pb2.CreateIamTokenResponse(iam_token="test_token") + response.expires_at.seconds = int(time.time() + 42) + return response + + +class IamTokenServiceTestServer(object): + def __init__(self): + self.server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=2)) + iam_token_service_pb2_grpc.add_IamTokenServiceServicer_to_server(IamTokenServiceForTest(), self.server) + self.server.add_insecure_port(self.get_endpoint()) + self.server.start() + + def stop(self): + self.server.wait_for_termination() + + def get_endpoint(self): + return "[::]:54321" + + +class TestServiceAccountCredentials(ydb.iam.ServiceAccountCredentials): + def __init__( + self, + service_account_id, + access_key_id, + private_key, + iam_endpoint=None, + iam_channel_credentials=None, + ): + super(TestServiceAccountCredentials, self).__init__( + service_account_id, + access_key_id, + private_key, + iam_endpoint, + iam_channel_credentials, + ) + + def _channel_factory(self): + return grpc.insecure_channel( + self._iam_endpoint + ) + + def get_expire_time(self): + return self._expires_in - time.time() + + +def test_service_account_credentials(): + server = IamTokenServiceTestServer() + iam_endpoint = server.get_endpoint() + grpc_channel_creds = grpc.local_channel_credentials() + credentials = TestServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, iam_endpoint, grpc_channel_creds) + credentials.set_token_expiration_timeout(1) + t = credentials.get_auth_token() + assert t == "test_token" + assert credentials.get_expire_time() <= 42 diff --git a/tests/table/test_tx.py b/tests/table/test_tx.py index 750a13bf..067ba9bd 100644 --- a/tests/table/test_tx.py +++ b/tests/table/test_tx.py @@ -38,17 +38,6 @@ def test_tx_begin(driver_sync, database): tx.rollback() -def test_credentials(): - credentials = ydb.iam.MetadataUrlCredentials() - raised = False - try: - credentials.auth_metadata() - except Exception: - raised = True - - assert raised - - def test_tx_snapshot_ro(driver_sync, database): session = driver_sync.table_client.session().create() description = ( From d86cc83f2acdabf6dbdfc858204b5b65d150ada9 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 2 Apr 2024 14:40:27 +0000 Subject: [PATCH 09/12] Test for Nebius credentials --- tests/auth/test_credentials.py | 99 ++++++++++++++++++++++++++-------- 1 file changed, 78 insertions(+), 21 deletions(-) diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 90b56e5f..4c3dc272 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -2,6 +2,10 @@ import concurrent.futures import grpc import time +import http.server +import urllib +import threading +import json import ydb.iam @@ -14,7 +18,7 @@ PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\nftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\nZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\ny4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\nJLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\njQIDAQAB\n-----END PUBLIC KEY-----\n" -def test_credentials(): +def test_metadata_credentials(): credentials = ydb.iam.MetadataUrlCredentials() raised = False try: @@ -48,6 +52,7 @@ def __init__(self): self.server.start() def stop(self): + self.server.stop(1) self.server.wait_for_termination() def get_endpoint(self): @@ -55,22 +60,6 @@ def get_endpoint(self): class TestServiceAccountCredentials(ydb.iam.ServiceAccountCredentials): - def __init__( - self, - service_account_id, - access_key_id, - private_key, - iam_endpoint=None, - iam_channel_credentials=None, - ): - super(TestServiceAccountCredentials, self).__init__( - service_account_id, - access_key_id, - private_key, - iam_endpoint, - iam_channel_credentials, - ) - def _channel_factory(self): return grpc.insecure_channel( self._iam_endpoint @@ -80,12 +69,80 @@ def get_expire_time(self): return self._expires_in - time.time() -def test_service_account_credentials(): +class TestNebiusServiceAccountCredentials(ydb.iam.NebiusServiceAccountCredentials): + def get_expire_time(self): + return self._expires_in - time.time() + + +class NebiusTokenServiceHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + assert self.headers["Content-Type"] == "application/x-www-form-urlencoded" + assert self.path == "/token/exchange" + content_length = int(self.headers["Content-Length"]) + post_data = self.rfile.read(content_length).decode("utf8") + print("NebiusTokenServiceHandler.POST data: {}".format(post_data)) + parsed_request = urllib.parse.parse_qs(str(post_data)) + assert len(parsed_request["grant_type"]) == 1 + assert parsed_request["grant_type"][0] == "urn:ietf:params:oauth:grant-type:token-exchange" + + assert len(parsed_request["requested_token_type"]) == 1 + assert parsed_request["requested_token_type"][0] == "urn:ietf:params:oauth:token-type:access_token" + + assert len(parsed_request["subject_token_type"]) == 1 + assert parsed_request["subject_token_type"][0] == "urn:ietf:params:oauth:token-type:jwt" + + assert len(parsed_request["subject_token"]) == 1 + jwt_token = parsed_request["subject_token"][0] + decoded = jwt.decode(jwt_token, key=PUBLIC_KEY, algorithms=["RS256"], audience="token-service.iam.new.nebiuscloud.net") + assert decoded["iss"] == SERVICE_ACCOUNT_ID + assert decoded["sub"] == SERVICE_ACCOUNT_ID + assert decoded["aud"] == "token-service.iam.new.nebiuscloud.net" + assert abs(decoded["iat"] - time.time()) <= 60 + assert abs(decoded["exp"] - time.time()) <= 3600 + + response = { + "access_token": "test_nebius_token", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "token_type": "Bearer", + "expires_in": 42 + } + + self.send_response(200) + self.send_header("Content-type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response).encode("utf8")) + + +class NebiusTokenServiceForTest(http.server.HTTPServer): + def __init__(self): + http.server.HTTPServer.__init__(self, ("localhost", 54322), NebiusTokenServiceHandler) + + def endpoint(self): + return "http://localhost:54322/token/exchange" + + +def test_yandex_service_account_credentials(): server = IamTokenServiceTestServer() - iam_endpoint = server.get_endpoint() - grpc_channel_creds = grpc.local_channel_credentials() - credentials = TestServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, iam_endpoint, grpc_channel_creds) + credentials = TestServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, server.get_endpoint()) credentials.set_token_expiration_timeout(1) t = credentials.get_auth_token() assert t == "test_token" assert credentials.get_expire_time() <= 42 + server.stop() + + +def test_nebius_service_account_credentials(): + server = NebiusTokenServiceForTest() + + def serve(s): + s.handle_request() + + serve_thread = threading.Thread(target=serve, args=(server,)) + serve_thread.start() + + credentials = TestNebiusServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, server.endpoint()) + t = credentials.get_auth_token() + assert t == "test_nebius_token" + assert credentials.get_expire_time() <= 42 + + serve_thread.join() From 5b5cfe66618afaba6081f41b1436c35628ad85aa Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 2 Apr 2024 16:28:33 +0000 Subject: [PATCH 10/12] Async tests --- tests/aio/test_credentials.py | 50 ++++++++++++++++++++++++++++++++++ tests/auth/test_credentials.py | 5 ++-- 2 files changed, 52 insertions(+), 3 deletions(-) create mode 100644 tests/aio/test_credentials.py diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py new file mode 100644 index 00000000..375e048d --- /dev/null +++ b/tests/aio/test_credentials.py @@ -0,0 +1,50 @@ +import pytest +import time +import grpc +import threading + +import tests.auth.test_credentials +import ydb.aio.iam + + +class TestServiceAccountCredentials(ydb.aio.iam.ServiceAccountCredentials): + def _channel_factory(self): + return grpc.aio.insecure_channel( + self._iam_endpoint + ) + + def get_expire_time(self): + return self._expires_in - time.time() + + +class TestNebiusServiceAccountCredentials(ydb.aio.iam.NebiusServiceAccountCredentials): + def get_expire_time(self): + return self._expires_in - time.time() + + +@pytest.mark.asyncio +async def test_yandex_service_account_credentials(): + server = tests.auth.test_credentials.IamTokenServiceTestServer() + credentials = TestServiceAccountCredentials(tests.auth.test_credentials.SERVICE_ACCOUNT_ID, tests.auth.test_credentials.ACCESS_KEY_ID, tests.auth.test_credentials.PRIVATE_KEY, server.get_endpoint()) + t = (await credentials.auth_metadata())[0][1] + assert t == "test_token" + assert credentials.get_expire_time() <= 42 + server.stop() + + +@pytest.mark.asyncio +async def test_nebius_service_account_credentials(): + server = tests.auth.test_credentials.NebiusTokenServiceForTest() + + def serve(s): + s.handle_request() + + serve_thread = threading.Thread(target=serve, args=(server,)) + serve_thread.start() + + credentials = TestNebiusServiceAccountCredentials(tests.auth.test_credentials.SERVICE_ACCOUNT_ID, tests.auth.test_credentials.ACCESS_KEY_ID, tests.auth.test_credentials.PRIVATE_KEY, server.endpoint()) + t = (await credentials.auth_metadata())[0][1] + assert t == "test_nebius_token" + assert credentials.get_expire_time() <= 42 + + serve_thread.join() diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 4c3dc272..d396bc63 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -14,8 +14,8 @@ SERVICE_ACCOUNT_ID = "sa_id" ACCESS_KEY_ID = "key_id" -PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\nFgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\nM0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\nnPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\nJnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\nR+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\nWYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\nYLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\nuccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\nzrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\nsvlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\nm6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\nrheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\nLxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\nmto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\nJUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\nBjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\nZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\nkFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\nnXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\nFXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\nTP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\ncHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\nof1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\nhL8bFAuNNVrCOp79TNnNIsh7\n-----END PRIVATE KEY-----\n" -PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\nftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\nZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\ny4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\nJLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\njQIDAQAB\n-----END PUBLIC KEY-----\n" +PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\nFgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\nM0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\nnPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\nJnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\nR+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\nWYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\nYLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\nuccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\nzrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\nsvlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\nm6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\nrheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\nLxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\nmto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\nJUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\nBjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\nZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\nkFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\nnXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\nFXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\nTP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\ncHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\nof1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\nhL8bFAuNNVrCOp79TNnNIsh7\n-----END PRIVATE KEY-----\n" # noqa: E501 +PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\nftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\nZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\ny4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\nJLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\njQIDAQAB\n-----END PUBLIC KEY-----\n" # noqa: E501 def test_metadata_credentials(): @@ -124,7 +124,6 @@ def endpoint(self): def test_yandex_service_account_credentials(): server = IamTokenServiceTestServer() credentials = TestServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, server.get_endpoint()) - credentials.set_token_expiration_timeout(1) t = credentials.get_auth_token() assert t == "test_token" assert credentials.get_expire_time() <= 42 From d036843e6da95caed78b4b1a4e8e7f87651e3a92 Mon Sep 17 00:00:00 2001 From: Vasily Gerasimov Date: Tue, 2 Apr 2024 16:31:29 +0000 Subject: [PATCH 11/12] Style --- tests/aio/test_credentials.py | 18 +++++++++++++----- tests/auth/test_credentials.py | 14 ++++++++------ 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index 375e048d..10878a40 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -9,9 +9,7 @@ class TestServiceAccountCredentials(ydb.aio.iam.ServiceAccountCredentials): def _channel_factory(self): - return grpc.aio.insecure_channel( - self._iam_endpoint - ) + return grpc.aio.insecure_channel(self._iam_endpoint) def get_expire_time(self): return self._expires_in - time.time() @@ -25,7 +23,12 @@ def get_expire_time(self): @pytest.mark.asyncio async def test_yandex_service_account_credentials(): server = tests.auth.test_credentials.IamTokenServiceTestServer() - credentials = TestServiceAccountCredentials(tests.auth.test_credentials.SERVICE_ACCOUNT_ID, tests.auth.test_credentials.ACCESS_KEY_ID, tests.auth.test_credentials.PRIVATE_KEY, server.get_endpoint()) + credentials = TestServiceAccountCredentials( + tests.auth.test_credentials.SERVICE_ACCOUNT_ID, + tests.auth.test_credentials.ACCESS_KEY_ID, + tests.auth.test_credentials.PRIVATE_KEY, + server.get_endpoint(), + ) t = (await credentials.auth_metadata())[0][1] assert t == "test_token" assert credentials.get_expire_time() <= 42 @@ -42,7 +45,12 @@ def serve(s): serve_thread = threading.Thread(target=serve, args=(server,)) serve_thread.start() - credentials = TestNebiusServiceAccountCredentials(tests.auth.test_credentials.SERVICE_ACCOUNT_ID, tests.auth.test_credentials.ACCESS_KEY_ID, tests.auth.test_credentials.PRIVATE_KEY, server.endpoint()) + credentials = TestNebiusServiceAccountCredentials( + tests.auth.test_credentials.SERVICE_ACCOUNT_ID, + tests.auth.test_credentials.ACCESS_KEY_ID, + tests.auth.test_credentials.PRIVATE_KEY, + server.endpoint(), + ) t = (await credentials.auth_metadata())[0][1] assert t == "test_nebius_token" assert credentials.get_expire_time() <= 42 diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index d396bc63..1002520f 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -33,7 +33,9 @@ class IamTokenServiceForTest(iam_token_service_pb2_grpc.IamTokenServiceServicer) def Create(self, request, context): print("IAM token service request: {}".format(request)) # Validate jwt: - decoded = jwt.decode(request.jwt, key=PUBLIC_KEY, algorithms=["PS256"], audience="https://iam.api.cloud.yandex.net/iam/v1/tokens") + decoded = jwt.decode( + request.jwt, key=PUBLIC_KEY, algorithms=["PS256"], audience="https://iam.api.cloud.yandex.net/iam/v1/tokens" + ) assert decoded["iss"] == SERVICE_ACCOUNT_ID assert decoded["aud"] == "https://iam.api.cloud.yandex.net/iam/v1/tokens" assert abs(decoded["iat"] - time.time()) <= 60 @@ -61,9 +63,7 @@ def get_endpoint(self): class TestServiceAccountCredentials(ydb.iam.ServiceAccountCredentials): def _channel_factory(self): - return grpc.insecure_channel( - self._iam_endpoint - ) + return grpc.insecure_channel(self._iam_endpoint) def get_expire_time(self): return self._expires_in - time.time() @@ -93,7 +93,9 @@ def do_POST(self): assert len(parsed_request["subject_token"]) == 1 jwt_token = parsed_request["subject_token"][0] - decoded = jwt.decode(jwt_token, key=PUBLIC_KEY, algorithms=["RS256"], audience="token-service.iam.new.nebiuscloud.net") + decoded = jwt.decode( + jwt_token, key=PUBLIC_KEY, algorithms=["RS256"], audience="token-service.iam.new.nebiuscloud.net" + ) assert decoded["iss"] == SERVICE_ACCOUNT_ID assert decoded["sub"] == SERVICE_ACCOUNT_ID assert decoded["aud"] == "token-service.iam.new.nebiuscloud.net" @@ -104,7 +106,7 @@ def do_POST(self): "access_token": "test_nebius_token", "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", "token_type": "Bearer", - "expires_in": 42 + "expires_in": 42, } self.send_response(200) From e4cd3eb20488bbaeb1e600532550323d28d15dbc Mon Sep 17 00:00:00 2001 From: Timofey Koolin Date: Wed, 3 Apr 2024 18:37:20 +0300 Subject: [PATCH 12/12] Update tests/auth/test_credentials.py --- tests/auth/test_credentials.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 1002520f..bd4c9809 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -58,7 +58,7 @@ def stop(self): self.server.wait_for_termination() def get_endpoint(self): - return "[::]:54321" + return "localhost:54321" class TestServiceAccountCredentials(ydb.iam.ServiceAccountCredentials):