diff options
author | robot-piglet <robot-piglet@yandex-team.com> | 2024-04-24 09:46:23 +0300 |
---|---|---|
committer | robot-piglet <robot-piglet@yandex-team.com> | 2024-04-24 09:57:59 +0300 |
commit | f123a9ad03e3fed5450470ab4a68741025fdc2e6 (patch) | |
tree | 39277834e5fcaaf398d456af85f2e78bd5291e6b /contrib/python | |
parent | f61ee5c5cae99c042ad96c01ab6719589dce0304 (diff) | |
download | ydb-f123a9ad03e3fed5450470ab4a68741025fdc2e6.tar.gz |
Intermediate changes
Diffstat (limited to 'contrib/python')
-rw-r--r-- | contrib/python/ydb/py3/.dist-info/METADATA | 2 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ya.make | 6 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/aio/iam.py | 25 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/aio/oauth2_token_exchange.py | 47 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/iam/auth.py | 60 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/oauth2_token_exchange/__init__.py | 4 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_exchange.py | 135 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_source.py | 104 | ||||
-rw-r--r-- | contrib/python/ydb/py3/ydb/ydb_version.py | 2 |
9 files changed, 350 insertions, 35 deletions
diff --git a/contrib/python/ydb/py3/.dist-info/METADATA b/contrib/python/ydb/py3/.dist-info/METADATA index e7397b376d..6db8a9f068 100644 --- a/contrib/python/ydb/py3/.dist-info/METADATA +++ b/contrib/python/ydb/py3/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.1 Name: ydb -Version: 3.10.0 +Version: 3.11.1 Summary: YDB Python SDK Home-page: http://github.com/ydb-platform/ydb-python-sdk Author: Yandex LLC diff --git a/contrib/python/ydb/py3/ya.make b/contrib/python/ydb/py3/ya.make index df4b531e51..bc2089bef2 100644 --- a/contrib/python/ydb/py3/ya.make +++ b/contrib/python/ydb/py3/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(3.10.0) +VERSION(3.11.1) LICENSE(Apache-2.0) @@ -53,6 +53,7 @@ PY_SRCS( ydb/aio/credentials.py ydb/aio/driver.py ydb/aio/iam.py + ydb/aio/oauth2_token_exchange.py ydb/aio/pool.py ydb/aio/resolver.py ydb/aio/scheme.py @@ -76,6 +77,9 @@ PY_SRCS( ydb/import_client.py ydb/interceptor.py ydb/issues.py + ydb/oauth2_token_exchange/__init__.py + ydb/oauth2_token_exchange/token_exchange.py + ydb/oauth2_token_exchange/token_source.py ydb/operation.py ydb/pool.py ydb/resolver.py diff --git a/contrib/python/ydb/py3/ydb/aio/iam.py b/contrib/python/ydb/py3/ydb/aio/iam.py index eab8faffe0..5a2a29f602 100644 --- a/contrib/python/ydb/py3/ydb/aio/iam.py +++ b/contrib/python/ydb/py3/ydb/aio/iam.py @@ -9,11 +9,14 @@ from .credentials import AbstractExpiringTokenCredentials 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 @@ -65,17 +68,17 @@ class JWTIamCredentials(TokenServiceCredentials, auth.BaseJWTCredentials): 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, + auth.YANDEX_CLOUD_JWT_ALGORITHM, + 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/contrib/python/ydb/py3/ydb/aio/oauth2_token_exchange.py b/contrib/python/ydb/py3/ydb/aio/oauth2_token_exchange.py new file mode 100644 index 0000000000..8b02c58a2c --- /dev/null +++ b/contrib/python/ydb/py3/ydb/aio/oauth2_token_exchange.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- +from .credentials import AbstractExpiringTokenCredentials +from ydb.oauth2_token_exchange.token_source import TokenSource +from ydb.oauth2_token_exchange.token_exchange import Oauth2TokenExchangeCredentialsBase +import typing + +try: + import aiohttp +except ImportError: + aiohttp = None + + +class Oauth2TokenExchangeCredentials(AbstractExpiringTokenCredentials, Oauth2TokenExchangeCredentialsBase): + def __init__( + self, + token_endpoint: str, + subject_token_source: typing.Optional[TokenSource] = None, + actor_token_source: typing.Optional[TokenSource] = None, + audience: typing.Union[typing.List[str], str, None] = None, + scope: typing.Union[typing.List[str], str, None] = None, + resource: typing.Optional[str] = None, + grant_type: str = "urn:ietf:params:oauth:grant-type:token-exchange", + requested_token_type: str = "urn:ietf:params:oauth:token-type:access_token", + ): + assert aiohttp is not None, "Install aiohttp library to use OAuth 2.0 token exchange credentials provider" + super(Oauth2TokenExchangeCredentials, self).__init__() + Oauth2TokenExchangeCredentialsBase.__init__( + self, + token_endpoint, + subject_token_source, + actor_token_source, + audience, + scope, + resource, + grant_type, + requested_token_type, + ) + + async def _make_token_request(self): + params = self._make_token_request_params() + 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_endpoint, data=params, headers=headers) as response: + self._process_response_status_code(await response.text(), response.status) + return self._process_response_json(await response.json()) diff --git a/contrib/python/ydb/py3/ydb/iam/auth.py b/contrib/python/ydb/py3/ydb/iam/auth.py index 82e7c9f6c8..7b4fa4e8d9 100644 --- a/contrib/python/ydb/py3/ydb/iam/auth.py +++ b/contrib/python/ydb/py3/ydb/iam/auth.py @@ -8,11 +8,14 @@ import json 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 @@ -23,22 +26,28 @@ except ImportError: 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" +YANDEX_CLOUD_JWT_ALGORITHM = "PS256" -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): + 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) + 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 +82,15 @@ class TokenServiceCredentials(credentials.AbstractExpiringTokenCredentials): 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 +111,17 @@ class BaseJWTCredentials(abc.ABC): 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 JWTIamCredentials(TokenServiceCredentials, BaseJWTCredentials): def __init__( @@ -110,17 +133,12 @@ class JWTIamCredentials(TokenServiceCredentials, BaseJWTCredentials): 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, YANDEX_CLOUD_JWT_ALGORITHM, 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()) class YandexPassportOAuthIamCredentials(TokenServiceCredentials): diff --git a/contrib/python/ydb/py3/ydb/oauth2_token_exchange/__init__.py b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/__init__.py new file mode 100644 index 0000000000..b43386eab3 --- /dev/null +++ b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/__init__.py @@ -0,0 +1,4 @@ +# -*- coding: utf-8 -*- +from .token_source import FixedTokenSource # noqa +from .token_source import JwtTokenSource # noqa +from .token_exchange import Oauth2TokenExchangeCredentials # noqa diff --git a/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_exchange.py b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_exchange.py new file mode 100644 index 0000000000..e4d1db9498 --- /dev/null +++ b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_exchange.py @@ -0,0 +1,135 @@ +# -*- coding: utf-8 -*- +import typing +import json +import abc + +try: + import requests +except ImportError: + requests = None + +from ydb import credentials, tracing, issues +from .token_source import TokenSource + + +class Oauth2TokenExchangeCredentialsBase(abc.ABC): + def __init__( + self, + token_endpoint: str, + subject_token_source: typing.Optional[TokenSource] = None, + actor_token_source: typing.Optional[TokenSource] = None, + audience: typing.Union[typing.List[str], str, None] = None, + scope: typing.Union[typing.List[str], str, None] = None, + resource: typing.Optional[str] = None, + grant_type: str = "urn:ietf:params:oauth:grant-type:token-exchange", + requested_token_type: str = "urn:ietf:params:oauth:token-type:access_token", + ): + self._token_endpoint = token_endpoint + self._subject_token_source = subject_token_source + self._actor_token_source = actor_token_source + self._audience = audience + self._scope = scope + self._resource = resource + self._grant_type = grant_type + self._requested_token_type = requested_token_type + + if not self._token_endpoint: + raise Exception("Oauth2 token exchange: no token endpoint specified") + + def _process_response_status_code(self, content: str, status_code: int): + if status_code == 403: + raise issues.Unauthenticated(content) + if status_code >= 500: + raise issues.Unavailable(content) + if status_code >= 400: + raise issues.BadRequest(content) + if status_code != 200: + raise issues.Error(content) + + def _process_response_json(self, response_json): + access_token = response_json["access_token"] + expires_in = response_json["expires_in"] + token_type = response_json["token_type"] + scope = response_json.get("scope") + if token_type.lower() != "bearer": + raise Exception("Oauth2 token exchange: unsupported token type: {}".format(token_type)) + if expires_in <= 0: + raise Exception("Oauth2 token exchange: incorrect expiration time: {}".format(expires_in)) + if scope and scope != self._get_scope_param(): + raise Exception( + 'Oauth2 token exchange: different scope. Expected: "{}", but got: "{}"'.format( + self._get_scope_param(), scope + ) + ) + return {"access_token": "Bearer " + access_token, "expires_in": expires_in} + + def _get_scope_param(self) -> typing.Optional[str]: + if self._scope is None: + return None + if isinstance(self._scope, str): + return self._scope + # list + return " ".join(self._scope) + + def _make_token_request_params(self): + params = { + "grant_type": self._grant_type, + "requested_token_type": self._requested_token_type, + } + if self._resource: + params["resource"] = self._resource + if self._audience: + params["audience"] = self._audience + scope = self._get_scope_param() + if scope: + params["scope"] = scope + if self._subject_token_source is not None: + t = self._subject_token_source.token() + params["subject_token"] = t.token + params["subject_token_type"] = t.token_type + if self._actor_token_source is not None: + t = self._actor_token_source.token() + params["actor_token"] = t.token + params["actor_token_type"] = t.token_type + + return params + + +class Oauth2TokenExchangeCredentials(credentials.AbstractExpiringTokenCredentials, Oauth2TokenExchangeCredentialsBase): + def __init__( + self, + token_endpoint: str, + subject_token_source: typing.Optional[TokenSource] = None, + actor_token_source: typing.Optional[TokenSource] = None, + audience: typing.Union[typing.List[str], str, None] = None, + scope: typing.Union[typing.List[str], str, None] = None, + resource: typing.Optional[str] = None, + grant_type: str = "urn:ietf:params:oauth:grant-type:token-exchange", + requested_token_type: str = "urn:ietf:params:oauth:token-type:access_token", + tracer=None, + ): + super(Oauth2TokenExchangeCredentials, self).__init__(tracer) + Oauth2TokenExchangeCredentialsBase.__init__( + self, + token_endpoint, + subject_token_source, + actor_token_source, + audience, + scope, + resource, + grant_type, + requested_token_type, + ) + + @tracing.with_trace() + def _make_token_request(self): + assert ( + requests is not None + ), "Install requests library to use Oauth2TokenExchangeCredentials credentials provider" + + params = self._make_token_request_params() + headers = {"Content-Type": "application/x-www-form-urlencoded"} + response = requests.post(self._token_endpoint, data=params, headers=headers) + self._process_response_status_code(response.content, response.status_code) + response_json = json.loads(response.content) + return self._process_response_json(response_json) diff --git a/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_source.py b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_source.py new file mode 100644 index 0000000000..f33e329b5c --- /dev/null +++ b/contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_source.py @@ -0,0 +1,104 @@ +# -*- coding: utf-8 -*- +import abc +import typing +import time +import os +from datetime import datetime + +try: + import jwt +except ImportError: + jwt = None + + +class Token(abc.ABC): + def __init__(self, token: str, token_type: str): + self.token = token + self.token_type = token_type + + +class TokenSource(abc.ABC): + @abc.abstractmethod + def token(self) -> Token: + """ + :return: A Token object ready for exchange + """ + pass + + +class FixedTokenSource(TokenSource): + def __init__(self, token: str, token_type: str): + self._token = Token(token, token_type) + + def token(self) -> Token: + return self._token + + +class JwtTokenSource(TokenSource): + def __init__( + self, + signing_method: str, + private_key: typing.Optional[str] = None, + private_key_file: typing.Optional[str] = None, + key_id: typing.Optional[str] = None, + issuer: typing.Optional[str] = None, + subject: typing.Optional[str] = None, + audience: typing.Union[typing.List[str], str, None] = None, + id: typing.Optional[str] = None, + token_ttl_seconds: int = 3600, + ): + assert jwt is not None, "Install pyjwt library to use jwt tokens" + self._signing_method = signing_method + self._key_id = key_id + if private_key and private_key_file: + raise Exception("JWT: both private_key and private_key_file are set") + self._private_key = "" + if private_key: + 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: + self._private_key = key_file.read() + self._issuer = issuer + self._subject = subject + self._audience = audience + self._id = id + self._token_ttl_seconds = token_ttl_seconds + if not self._signing_method: + raise Exception("JWT: no signing method specified") + if not self._private_key: + raise Exception("JWT: no private key specified") + if self._token_ttl_seconds <= 0: + raise Exception("JWT: invalid jwt token TTL") + + def token(self) -> Token: + now = time.time() + now_utc = datetime.utcfromtimestamp(now) + exp_utc = datetime.utcfromtimestamp(now + self._token_ttl_seconds) + payload = { + "iat": now_utc, + "exp": exp_utc, + } + if self._audience: + payload["aud"] = self._audience + if self._issuer: + payload["iss"] = self._issuer + if self._subject: + payload["sub"] = self._subject + if self._id: + payload["jti"] = self._id + + headers = { + "alg": self._signing_method, + "typ": "JWT", + } + if self._key_id: + headers["kid"] = self._key_id + + token = jwt.encode( + key=self._private_key, + algorithm=self._signing_method, + headers=headers, + payload=payload, + ) + return Token(token, "urn:ietf:params:oauth:token-type:jwt") diff --git a/contrib/python/ydb/py3/ydb/ydb_version.py b/contrib/python/ydb/py3/ydb/ydb_version.py index 71960e35b5..03dc246c4e 100644 --- a/contrib/python/ydb/py3/ydb/ydb_version.py +++ b/contrib/python/ydb/py3/ydb/ydb_version.py @@ -1 +1 @@ -VERSION = "3.10.0" +VERSION = "3.11.1" |