aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python
diff options
context:
space:
mode:
authorrobot-piglet <robot-piglet@yandex-team.com>2024-04-24 09:46:23 +0300
committerrobot-piglet <robot-piglet@yandex-team.com>2024-04-24 09:57:59 +0300
commitf123a9ad03e3fed5450470ab4a68741025fdc2e6 (patch)
tree39277834e5fcaaf398d456af85f2e78bd5291e6b /contrib/python
parentf61ee5c5cae99c042ad96c01ab6719589dce0304 (diff)
downloadydb-f123a9ad03e3fed5450470ab4a68741025fdc2e6.tar.gz
Intermediate changes
Diffstat (limited to 'contrib/python')
-rw-r--r--contrib/python/ydb/py3/.dist-info/METADATA2
-rw-r--r--contrib/python/ydb/py3/ya.make6
-rw-r--r--contrib/python/ydb/py3/ydb/aio/iam.py25
-rw-r--r--contrib/python/ydb/py3/ydb/aio/oauth2_token_exchange.py47
-rw-r--r--contrib/python/ydb/py3/ydb/iam/auth.py60
-rw-r--r--contrib/python/ydb/py3/ydb/oauth2_token_exchange/__init__.py4
-rw-r--r--contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_exchange.py135
-rw-r--r--contrib/python/ydb/py3/ydb/oauth2_token_exchange/token_source.py104
-rw-r--r--contrib/python/ydb/py3/ydb/ydb_version.py2
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"