From 4ce89ff1d196c3a5372aebd1eeea575676cb1dce Mon Sep 17 00:00:00 2001 From: hiranya911 Date: Thu, 26 Mar 2020 15:37:56 -0700 Subject: [PATCH 1/2] Added tenant-scoped custom token support --- firebase_admin/_token_gen.py | 4 +++- firebase_admin/auth.py | 3 ++- tests/test_tenant_mgt.py | 29 +++++++++++++++++++++++++++++ tests/test_token_gen.py | 9 +++++++-- 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/firebase_admin/_token_gen.py b/firebase_admin/_token_gen.py index 4234bfa7b..2e40414cd 100644 --- a/firebase_admin/_token_gen.py +++ b/firebase_admin/_token_gen.py @@ -130,7 +130,7 @@ def signing_provider(self): 'details on creating custom tokens.'.format(error, url)) return self._signing_provider - def create_custom_token(self, uid, developer_claims=None): + def create_custom_token(self, uid, developer_claims=None, tenant_id=None): """Builds and signs a Firebase custom auth token.""" if developer_claims is not None: if not isinstance(developer_claims, dict): @@ -161,6 +161,8 @@ def create_custom_token(self, uid, developer_claims=None): 'iat': now, 'exp': now + MAX_TOKEN_LIFETIME_SECONDS, } + if tenant_id: + payload['tenant_id'] = tenant_id if developer_claims is not None: payload['claims'] = developer_claims diff --git a/firebase_admin/auth.py b/firebase_admin/auth.py index 4b80bc05e..e42dd4980 100644 --- a/firebase_admin/auth.py +++ b/firebase_admin/auth.py @@ -556,7 +556,8 @@ def tenant_id(self): return self._tenant_id def create_custom_token(self, uid, developer_claims=None): - return self._token_generator.create_custom_token(uid, developer_claims) + return self._token_generator.create_custom_token( + uid, developer_claims, tenant_id=self.tenant_id) def verify_id_token(self, id_token, check_revoked=False): """Verifies the signature and data for the provided ID token.""" diff --git a/tests/test_tenant_mgt.py b/tests/test_tenant_mgt.py index 4a765769f..da14c954f 100644 --- a/tests/test_tenant_mgt.py +++ b/tests/test_tenant_mgt.py @@ -21,6 +21,7 @@ import firebase_admin from firebase_admin import auth +from firebase_admin import credentials from firebase_admin import exceptions from firebase_admin import tenant_mgt from tests import testutils @@ -731,6 +732,34 @@ def test_invalid_tenant_id(self, tenant_mgt_app): assert excinfo.value.http_response is None +@pytest.fixture(scope='module') +def tenant_aware_custom_token_app(): + cred = credentials.Certificate(testutils.resource_filename('service_account.json')) + app = firebase_admin.initialize_app(cred, name='tenantAwareCustomToken') + yield app + firebase_admin.delete_app(app) + + +class TestCreateCustomToken: + + def test_custom_token(self, tenant_aware_custom_token_app): + client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app) + + custom_token = client.create_custom_token('user1') + + test_token_gen.verify_custom_token( + custom_token, expected_claims=None, tenant_id='test-tenant') + + def test_custom_token_with_claims(self, tenant_aware_custom_token_app): + client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app) + claims = {'admin': True} + + custom_token = client.create_custom_token('user1', claims) + + test_token_gen.verify_custom_token( + custom_token, expected_claims=claims, tenant_id='test-tenant') + + def _assert_tenant(tenant, tenant_id='tenant-id'): assert isinstance(tenant, tenant_mgt.Tenant) assert tenant.tenant_id == tenant_id diff --git a/tests/test_token_gen.py b/tests/test_token_gen.py index bbb6ef4e4..e320a1aaf 100644 --- a/tests/test_token_gen.py +++ b/tests/test_token_gen.py @@ -66,7 +66,7 @@ def _merge_jwt_claims(defaults, overrides): del defaults[key] return defaults -def _verify_custom_token(custom_token, expected_claims): +def verify_custom_token(custom_token, expected_claims, tenant_id=None): assert isinstance(custom_token, bytes) token = google.oauth2.id_token.verify_token( custom_token, @@ -75,6 +75,11 @@ def _verify_custom_token(custom_token, expected_claims): assert token['uid'] == MOCK_UID assert token['iss'] == MOCK_SERVICE_ACCOUNT_EMAIL assert token['sub'] == MOCK_SERVICE_ACCOUNT_EMAIL + if tenant_id is None: + assert token.get('tenant_id') is None + else: + assert token['tenant_id'] == tenant_id + header = jwt.decode_header(custom_token) assert header.get('typ') == 'JWT' assert header.get('alg') == 'RS256' @@ -198,7 +203,7 @@ class TestCreateCustomToken: def test_valid_params(self, auth_app, values): user, claims = values custom_token = auth.create_custom_token(user, claims, app=auth_app) - _verify_custom_token(custom_token, claims) + verify_custom_token(custom_token, claims) @pytest.mark.parametrize('values', invalid_args.values(), ids=list(invalid_args)) def test_invalid_params(self, auth_app, values): From 6674428bb98e08e786c8d014a0b62628cc127294 Mon Sep 17 00:00:00 2001 From: hiranya911 Date: Thu, 26 Mar 2020 15:42:10 -0700 Subject: [PATCH 2/2] Fixed a lint error; Improved test assertion --- tests/test_token_gen.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_token_gen.py b/tests/test_token_gen.py index e320a1aaf..072de3f16 100644 --- a/tests/test_token_gen.py +++ b/tests/test_token_gen.py @@ -76,9 +76,9 @@ def verify_custom_token(custom_token, expected_claims, tenant_id=None): assert token['iss'] == MOCK_SERVICE_ACCOUNT_EMAIL assert token['sub'] == MOCK_SERVICE_ACCOUNT_EMAIL if tenant_id is None: - assert token.get('tenant_id') is None + assert 'tenant_id' not in token else: - assert token['tenant_id'] == tenant_id + assert token['tenant_id'] == tenant_id header = jwt.decode_header(custom_token) assert header.get('typ') == 'JWT'