Intermediate changes

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"