Merge pull request #417 from UgnineSirdis/oauth2-token-exchange Implement OAuth 2.0 Token Exchange credentials provider in Python SDK

Author: rekby

examples/oauth2-token-exchange-credentials/main.py

@@ -0,0 +1,78 @@
+import argparse
+import ydb
+import ydb.oauth2_token_exchange
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        usage="%(prog)s options", description="Oauth 2.0 token exchange credentials example"
+    )
+    parser.add_argument("-d", "--database", required=True, help="Name of the database to use")
+    parser.add_argument("-e", "--endpoint", required=True, help="Endpoint url to use")
+    parser.add_argument("--token-endpoint", required=True, help="Token endpoint url to use")
+    parser.add_argument("--private-key-file", required=True, help="Private key file path in pem format")
+    parser.add_argument("--key-id", help="Key id")
+    parser.add_argument("--audience", help="Audience")
+    parser.add_argument("--issuer", help="Jwt token issuer")
+    parser.add_argument("--subject", help="Jwt token subject")
+
+    return parser.parse_args()
+
+
+def execute_query(session):
+    # Create the transaction and execute the `select 1` query.
+    # All transactions must be committed using the `commit_tx` flag in the last
+    # statement. The either way to commit transaction is using `commit` method of `TxContext` object, which is
+    # not recommended.
+    return session.transaction().execute(
+        "select 1 as cnt;",
+        commit_tx=True,
+        settings=ydb.BaseRequestSettings().with_timeout(3).with_operation_timeout(2),
+    )
+
+
+def main():
+    args = parse_args()
+
+    # Example demonstrates how to initializate driver instance
+    # using the oauth 2.0 token exchange credentials provider.
+    driver = ydb.Driver(
+        endpoint=args.endpoint,
+        database=args.database,
+        root_certificates=ydb.load_ydb_root_certificate(),
+        credentials=ydb.oauth2_token_exchange.Oauth2TokenExchangeCredentials(
+            token_endpoint=args.token_endpoint,
+            audience=args.audience,
+            subject_token_source=ydb.oauth2_token_exchange.JwtTokenSource(
+                signing_method="RS256",
+                private_key_file=args.private_key_file,
+                key_id=args.key_id,
+                issuer=args.issuer,
+                subject=args.subject,
+                audience=args.audience,
+            ),
+        ),
+    )
+
+    # Start driver context manager.
+    # The recommended way of using Driver object is using `with`
+    # clause, because the context manager automatically stops the driver.
+    with driver:
+        # wait until driver become initialized
+        driver.wait(fail_fast=True, timeout=5)
+
+        # Initialize the session pool instance and enter the context manager.
+        # The context manager automatically stops the session pool.
+        # On the session pool termination all YDB sessions are closed.
+        with ydb.SessionPool(driver) as pool:
+
+            # Execute the query with the `retry_operation_helper` the.
+            # The `retry_operation_sync` helper used to help developers
+            # to retry YDB specific errors like locks invalidation.
+            # The first argument of the `retry_operation_sync` is a function to retry.
+            # This function must have session as the first argument.
+            result = pool.retry_operation_sync(execute_query)
+            assert result[0].rows[0].cnt == 1
+
+
+main()

test-requirements.txt

@@ -30,6 +30,7 @@ pytest-asyncio==0.21.0
 pytest-docker-compose==3.2.1
 python-dotenv==0.18.0
 PyYAML==5.3.1
+pyjwt==2.0.0
 requests==2.31.0
 texttable==1.6.4
 toml==0.10.2
@@ -46,4 +47,5 @@ pylint-protobuf
 cython
 freezegun==1.2.2
 pytest-cov
+yandexcloud
 -e .

tests/aio/test_credentials.py

@@ -0,0 +1,62 @@
+import pytest
+import time
+import grpc
+import threading
+
+import tests.auth.test_credentials
+import tests.oauth2_token_exchange
+import tests.oauth2_token_exchange.test_token_exchange
+import ydb.aio.iam
+import ydb.aio.oauth2_token_exchange
+import ydb.oauth2_token_exchange.token_source
+
+
+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 TestOauth2TokenExchangeCredentials(ydb.aio.oauth2_token_exchange.Oauth2TokenExchangeCredentials):
+    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_oauth2_token_exchange_credentials():
+    server = tests.oauth2_token_exchange.test_token_exchange.Oauth2TokenExchangeServiceForTest(40124)
+
+    def serve(s):
+        s.handle_request()
+
+    serve_thread = threading.Thread(target=serve, args=(server,))
+    serve_thread.start()
+
+    credentials = TestOauth2TokenExchangeCredentials(
+        server.endpoint(),
+        ydb.oauth2_token_exchange.token_source.FixedTokenSource("test_src_token", "test_token_type"),
+        audience=["a1", "a2"],
+        scope=["s1", "s2"],
+    )
+    t = (await credentials.auth_metadata())[0][1]
+    assert t == "Bearer test_dst_token"
+    assert credentials.get_expire_time() <= 42
+
+    serve_thread.join()

tests/auth/__init__.py

@@ -0,0 +1,74 @@
+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"  # 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():
+    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.stop(1)
+        self.server.wait_for_termination()
+
+    def get_endpoint(self):
+        return "localhost:54321"
+
+
+class TestServiceAccountCredentials(ydb.iam.ServiceAccountCredentials):
+    def _channel_factory(self):
+        return grpc.insecure_channel(self._iam_endpoint)
+
+    def get_expire_time(self):
+        return self._expires_in - time.time()
+
+
+def test_yandex_service_account_credentials():
+    server = IamTokenServiceTestServer()
+    credentials = TestServiceAccountCredentials(SERVICE_ACCOUNT_ID, ACCESS_KEY_ID, PRIVATE_KEY, server.get_endpoint())
+    t = credentials.get_auth_token()
+    assert t == "test_token"
+    assert credentials.get_expire_time() <= 42
+    server.stop()