Test for Yandex jwt

Author: UgnineSirdis

requirements.txt

@@ -2,3 +2,4 @@ grpcio>=1.42.0
 packaging
 protobuf>=3.13.0,<5.0.0
 aiohttp<4
+pyjwt==2.8.0

test-requirements.txt

@@ -46,4 +46,5 @@ pylint-protobuf
 cython
 freezegun==1.2.2
 pytest-cov
+yandexcloud
 -e .

tests/auth/__init__.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

tests/auth/test_credentials.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 = (