|
| 1 | +# The MIT License (MIT) |
| 2 | +# Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | + |
| 4 | +import base64 |
| 5 | +import json |
| 6 | +import time |
| 7 | +import unittest |
| 8 | +from io import StringIO |
| 9 | + |
| 10 | +import pytest |
| 11 | +from azure.core.credentials import AccessToken |
| 12 | + |
| 13 | +import test_config |
| 14 | +from azure.cosmos import exceptions |
| 15 | +from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy |
| 16 | + |
| 17 | + |
| 18 | +def _remove_padding(encoded_string): |
| 19 | + while encoded_string.endswith("="): |
| 20 | + encoded_string = encoded_string[0:len(encoded_string) - 1] |
| 21 | + |
| 22 | + return encoded_string |
| 23 | + |
| 24 | + |
| 25 | +def get_test_item(num): |
| 26 | + test_item = { |
| 27 | + 'pk': 'pk', |
| 28 | + 'id': 'Item_' + str(num), |
| 29 | + 'test_object': True, |
| 30 | + 'lastName': 'Smith' |
| 31 | + } |
| 32 | + return test_item |
| 33 | + |
| 34 | + |
| 35 | +class CosmosEmulatorCredential(object): |
| 36 | + |
| 37 | + async def get_token(self, *scopes, **kwargs): |
| 38 | + # type: (*str, **Any) -> AccessToken |
| 39 | + """Request an access token for the emulator. Based on Azure Core's Access Token Credential. |
| 40 | +
|
| 41 | + This method is called automatically by Azure SDK clients. |
| 42 | +
|
| 43 | + :param str scopes: desired scopes for the access token. This method requires at least one scope. |
| 44 | + :rtype: :class:`azure.core.credentials.AccessToken` |
| 45 | + :raises CredentialUnavailableError: the credential is unable to attempt authentication because it lacks |
| 46 | + required data, state, or platform support |
| 47 | + :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message`` |
| 48 | + attribute gives a reason. |
| 49 | + """ |
| 50 | + aad_header_cosmos_emulator = "{\"typ\":\"JWT\",\"alg\":\"RS256\",\"x5t\":\"" \ |
| 51 | + "CosmosEmulatorPrimaryMaster\",\"kid\":\"CosmosEmulatorPrimaryMaster\"}" |
| 52 | + |
| 53 | + aad_claim_cosmos_emulator_format = {"aud": "https://localhost.localhost", |
| 54 | + "iss": "https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984", |
| 55 | + "iat": int(time.time()), "nbf": int(time.time()), |
| 56 | + "exp": int(time.time() + 7200), "aio": "", "appid": "localhost", |
| 57 | + "appidacr": "1", "idp": "https://localhost:8081/", |
| 58 | + "oid": "96313034-4739-43cb-93cd-74193adbe5b6", "rh": "", "sub": "localhost", |
| 59 | + "tid": "EmulatorFederation", "uti": "", "ver": "1.0", |
| 60 | + "scp": "user_impersonation", |
| 61 | + "groups": ["7ce1d003-4cb3-4879-b7c5-74062a35c66e", |
| 62 | + "e99ff30c-c229-4c67-ab29-30a6aebc3e58", |
| 63 | + "5549bb62-c77b-4305-bda9-9ec66b85d9e4", |
| 64 | + "c44fd685-5c58-452c-aaf7-13ce75184f65", |
| 65 | + "be895215-eab5-43b7-9536-9ef8fe130330"]} |
| 66 | + |
| 67 | + emulator_key = test_config.TestConfig.masterKey |
| 68 | + |
| 69 | + first_encoded_bytes = base64.urlsafe_b64encode(aad_header_cosmos_emulator.encode("utf-8")) |
| 70 | + first_encoded_padded = str(first_encoded_bytes, "utf-8") |
| 71 | + first_encoded = _remove_padding(first_encoded_padded) |
| 72 | + |
| 73 | + str_io_obj = StringIO() |
| 74 | + json.dump(aad_claim_cosmos_emulator_format, str_io_obj) |
| 75 | + aad_claim_cosmos_emulator_format_string = str(str_io_obj.getvalue()).replace(" ", "") |
| 76 | + second = aad_claim_cosmos_emulator_format_string |
| 77 | + second_encoded_bytes = base64.urlsafe_b64encode(second.encode("utf-8")) |
| 78 | + second_encoded_padded = str(second_encoded_bytes, "utf-8") |
| 79 | + second_encoded = _remove_padding(second_encoded_padded) |
| 80 | + |
| 81 | + emulator_key_encoded_bytes = base64.urlsafe_b64encode(emulator_key.encode("utf-8")) |
| 82 | + emulator_key_encoded_padded = str(emulator_key_encoded_bytes, "utf-8") |
| 83 | + emulator_key_encoded = _remove_padding(emulator_key_encoded_padded) |
| 84 | + |
| 85 | + return AccessToken(first_encoded + "." + second_encoded + "." + emulator_key_encoded, int(time.time() + 7200)) |
| 86 | + |
| 87 | + |
| 88 | +@pytest.mark.cosmosEmulator |
| 89 | +class TestAADAsync(unittest.TestCase): |
| 90 | + client: CosmosClient = None |
| 91 | + database: DatabaseProxy = None |
| 92 | + container: ContainerProxy = None |
| 93 | + configs = test_config.TestConfig |
| 94 | + host = configs.host |
| 95 | + masterKey = configs.masterKey |
| 96 | + credential = CosmosEmulatorCredential() if configs.is_emulator else configs.credential_async |
| 97 | + |
| 98 | + @classmethod |
| 99 | + async def setUpClass(cls): |
| 100 | + cls.client = CosmosClient(cls.host, cls.credential) |
| 101 | + cls.database = cls.client.get_database_client(cls.configs.TEST_DATABASE_ID) |
| 102 | + cls.container = cls.database.get_container_client(cls.configs.TEST_SINGLE_PARTITION_CONTAINER_ID) |
| 103 | + |
| 104 | + async def test_aad_credentials_async(self): |
| 105 | + # Do any R/W data operations with your authorized AAD client |
| 106 | + |
| 107 | + print("Container info: " + str(self.container.read())) |
| 108 | + await self.container.create_item(get_test_item(0)) |
| 109 | + print("Point read result: " + str(await self.container.read_item(item='Item_0', partition_key='pk'))) |
| 110 | + query_results = list(await self.container.query_items(query='select * from c', partition_key='pk')) |
| 111 | + assert len(query_results) == 1 |
| 112 | + print("Query result: " + str(query_results[0])) |
| 113 | + await self.container.delete_item(item='Item_0', partition_key='pk') |
| 114 | + |
| 115 | + # Attempting to do management operations will return a 403 Forbidden exception |
| 116 | + try: |
| 117 | + await self.client.delete_database(self.configs.TEST_DATABASE_ID) |
| 118 | + except exceptions.CosmosHttpResponseError as e: |
| 119 | + assert e.status_code == 403 |
| 120 | + print("403 error assertion success") |
| 121 | + |
| 122 | + |
| 123 | +if __name__ == "__main__": |
| 124 | + unittest.main() |
0 commit comments