10000 Added tenant-scoped custom token support (#433) · firebase/firebase-admin-python@d35cd6d · GitHub
[go: up one dir, main page]

Skip to content

Commit d35cd6d

Browse files
authored
Added tenant-scoped custom token support (#433)
* Added tenant-scoped custom token support * Fixed a lint error; Improved test assertion
1 parent 19628b7 commit d35cd6d

File tree

4 files changed

+41
-4
lines changed

4 files changed

+41
-4
lines changed

firebase_admin/_token_gen.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ def signing_provider(self):
130130
'details on creating custom tokens.'.format(error, url))
131131
return self._signing_provider
132132

133-
def create_custom_token(self, uid, developer_claims=None):
133+
def create_custom_token(self, uid, developer_claims=None, tenant_id=None):
134134
"""Builds and signs a Firebase custom auth token."""
135135
if developer_claims is not None:
136136
if not isinstance(developer_claims, dict):
@@ -161,6 +161,8 @@ def create_custom_token(self, uid, developer_claims=None):
161161
'iat': now,
162162
'exp': now + MAX_TOKEN_LIFETIME_SECONDS,
163163
}
164+
if tenant_id:
165+
payload['tenant_id'] = tenant_id
164166

165167
if developer_claims is not None:
166168
payload['claims'] = developer_claims

firebase_admin/auth.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,8 @@ def tenant_id(self):
556556
return self._tenant_id
557557

558558
def create_custom_token(self, uid, developer_claims=None):
559-
return self._token_generator.create_custom_token(uid, developer_claims)
559+
return self._token_generator.create_custom_token(
560+
uid, developer_claims, tenant_id=self.tenant_id)
560561

561562
def verify_id_token(self, id_token, check_revoked=False):
562563
"""Verifies the signature and data for the provided ID token."""

tests/test_tenant_mgt.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import firebase_admin
2323
from firebase_admin import auth
24+
from firebase_admin import credentials
2425
from firebase_admin import exceptions
2526
from firebase_admin import tenant_mgt
2627
from tests import testutils
@@ -731,6 +732,34 @@ def test_invalid_tenant_id(self, tenant_mgt_app):
731732
assert excinfo.value.http_response is None
732733

733734

735+
@pytest.fixture(scope='module')
736+
def tenant_aware_custom_token_app():
737+
cred = credentials.Certificate(testutils.resource_filename('service_account.json'))
738+
app = firebase_admin.initialize_app(cred, name='tenantAwareCustomToken')
739+
yield app
740+
firebase_admin.delete_app(app)
741+
742+
743+
class TestCreateCustomToken:
744+
745+
def test_custom_token(self, tenant_aware_custom_token_app):
746+
client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app)
747+
748+
custom_token = client.create_custom_token('user1')
749+
750+
test_token_gen.verify_custom_token(
751+
custom_token, expected_claims=None, tenant_id='test-tenant')
752+
753+
def test_custom_token_with_claims(self, tenant_aware_custom_token_app):
754+
client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app)
755+
claims = {'admin': True}
756+
757+
custom_token = client.create_custom_token('user1', claims)
758+
759+
test_token_gen.verify_custom_token(
760+
custom_token, expected_claims=claims, tenant_id='test-tenant')
761+
762+
734763
def _assert_tenant(tenant, tenant_id='tenant-id'):
735764
assert isinstance(tenant, tenant_mgt.Tenant)
736765
assert tenant.tenant_id == tenant_id

tests/test_token_gen.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def _merge_jwt_claims(defaults, overrides):
6666
del defaults[key]
6767
return defaults
6868

69-
def _verify_custom_token(custom_token, expected_claims):
69+
def verify_custom_token(custom_token, expected_claims, tenant_id=None):
7070
assert isinstance(custom_token, bytes)
7171
token = google.oauth2.id_token.verify_token(
7272
custom_token,
@@ -75,6 +75,11 @@ def _verify_custom_token(custom_token, expected_claims):
7575
assert token['uid'] == MOCK_UID
7676
assert token['iss'] == MOCK_SERVICE_ACCOUNT_EMAIL
7777
assert token['sub'] == MOCK_SERVICE_ACCOUNT_EMAIL
78+
if tenant_id is None:
79+
assert 'tenant_id' not in token
80+
else:
81+
assert token['tenant_id'] == tenant_id
82+
7883
header = jwt.decode_header(custom_token)
7984
assert header.get('typ') == 'JWT'
8085
assert header.get('alg') == 'RS256'
@@ -198,7 +203,7 @@ class TestCreateCustomToken:
198203
def test_valid_params(self, auth_app, values):
199204
user, claims = values
200205
custom_token = auth.create_custom_token(user, claims, app=auth_app)
201-
_verify_custom_token(custom_token, claims)
206+
verify_custom_token(custom_token, claims)
202207

203208
@pytest.mark.parametrize('values', invalid_args.values(), ids=list(invalid_args))
204209
def test_invalid_params(self, auth_app, values):

0 commit comments

Comments
 (0)
0