Custom Claims and List Users Support (#86)

* Implemented set_custom_user_claims() function

* Unit and integration tests for the new functionality

* Implemented list_users() function

* Tests for list users API

* Updated tests

* Implemented user iteration using Python's native iterable API

* Cleaned up the test code

* Moved user iteration logic to _user_mgt

* Code cleanup

* Updated the list users API by adding the ListUsersPage class

* Updated error message

Author: hiranya911

firebase_admin/_user_mgt.py

@@ -14,6 +14,7 @@
 
 """Firebase user management sub module."""
 
+import json
 import re
 
 from google.auth import transport
@@ -29,9 +30,17 @@
 USER_CREATE_ERROR = 'USER_CREATE_ERROR'
 USER_UPDATE_ERROR = 'USER_UPDATE_ERROR'
 USER_DELETE_ERROR = 'USER_DELETE_ERROR'
+USER_DOWNLOAD_ERROR = 'LIST_USERS_ERROR'
 
 ID_TOOLKIT_URL = 'https://www.googleapis.com/identitytoolkit/v3/relyingparty/'
 
+MAX_LIST_USERS_RESULTS = 1000
+MAX_CLAIMS_PAYLOAD_SIZE = 1000
+RESERVED_CLAIMS = set([
+    'acr', 'amr', 'at_hash', 'aud', 'auth_time', 'azp', 'cnf', 'c_hash', 'exp', 'iat',
+    'iss', 'jti', 'nbf', 'nonce', 'sub', 'firebase',
+])
+
 
 class _Validator(object):
     """A collectoin of data validation utilities.
@@ -118,6 +127,37 @@ def validate_delete_list(cls, delete_attr):
                 'Invalid delete list: "{0}". Delete list must be a '
                 'non-empty list.'.format(delete_attr))
 
+    @classmethod
+    def validate_custom_claims(cls, custom_claims):
+        """Validates the specified custom claims.
+
+        Custom claims must be specified as a JSON string.The string must not exceed 1000
+        characters, and the parsed JSON payload must not contain reserved JWT claims.
+        """
+        if not isinstance(custom_claims, six.string_types) or not custom_claims:
+            raise ValueError(
+                'Invalid custom claims: "{0}". Custom claims must be a non-empty JSON '
+                'string.'.format(custom_claims))
+
+        if len(custom_claims) > MAX_CLAIMS_PAYLOAD_SIZE:
+            raise ValueError(
+                'Custom claims payload must not exceed {0} '
+                'characters.'.format(MAX_CLAIMS_PAYLOAD_SIZE))
+        try:
+            parsed = json.loads(custom_claims)
+        except Exception:
+            raise ValueError('Failed to parse custom claims string as JSON.')
+        else:
+            if not isinstance(parsed, dict):
+                raise ValueError('Custom claims must be parseable as a JSON object.')
+            invalid_claims = RESERVED_CLAIMS.intersection(set(parsed.keys()))
+            if len(invalid_claims) > 1:
+                joined = ', '.join(sorted(invalid_claims))
+                raise ValueError('Claims "{0}" are reserved, and must not be set.'.format(joined))
+            elif len(invalid_claims) == 1:
+                raise ValueError(
+                    'Claim "{0}" is reserved, and must not be set.'.format(invalid_claims.pop()))
+
 
 class ApiCallError(Exception):
     """Represents an Exception encountered while invoking the Firebase user management API."""
@@ -132,6 +172,7 @@ class UserManager(object):
     """Provides methods for interacting with the Google Identity Toolkit."""
 
     _VALIDATORS = {
+        'customAttributes' : _Validator.validate_custom_claims,
         'deleteAttribute' : _Validator.validate_delete_list,
         'deleteProvider' : _Validator.validate_delete_list,
         'disabled' : _Validator.validate_disabled,
@@ -163,7 +204,8 @@ class UserManager(object):
         'phone_number' : 'phoneNumber',
         'photo_url' : 'photoUrl',
         'password' : 'password',
-        'disabled' : 'disabled',
+        'disabled' : 'disableUser',
+        'custom_claims' : 'customAttributes',
     }
 
     _REMOVABLE_FIELDS = {
@@ -207,6 +249,26 @@ def get_user(self, **kwargs):
                     'No user record found for the provided {0}: {1}.'.format(key_type, key))
             return response['users'][0]
 
+    def list_users(self, page_token=None, max_results=MAX_LIST_USERS_RESULTS):
+        """Retrieves a batch of users."""
+        if page_token is not None:
+            if not isinstance(page_token, six.string_types) or not page_token:
+                raise ValueError('Page token must be a non-empty string.')
+        if not isinstance(max_results, int):
+            raise ValueError('Max results must be an integer.')
+        elif max_results < 1 or max_results > MAX_LIST_USERS_RESULTS:
+            raise ValueError(
+                'Max results must be a positive integer less than '
+                '{0}.'.format(MAX_LIST_USERS_RESULTS))
+
+        payload = {'maxResults': max_results}
+        if page_token:
+            payload['nextPageToken'] = page_token
+        try:
+            return self._request('post', 'downloadAccount', json=payload)
+        except requests.exceptions.RequestException as error:
+            self._handle_http_error(USER_DOWNLOAD_ERROR, 'Failed to download user accounts.', error)
+
     def create_user(self, **kwargs):
         """Creates a new user account with the specified properties."""
         payload = self._init_payload('create_user', UserManager._CREATE_USER_FIELDS, **kwargs)
@@ -236,9 +298,12 @@ def update_user(self, uid, **kwargs):
         if 'phoneNumber' in payload and payload['phoneNumber'] is None:
             payload['deleteProvider'] = ['phone']
             del payload['phoneNumber']
-        if 'disabled' in payload:
-            payload['disableUser'] = payload['disabled']
-            del payload['disabled']
+        if 'customAttributes' in payload:
+            custom_claims = payload['customAttributes']
+            if custom_claims is None:
+                custom_claims = {}
+            if isinstance(custom_claims, dict):
+                payload['customAttributes'] = json.dumps(custom_claims)
 
         self._validate(payload, self._VALIDATORS, 'update user')
         try:
@@ -306,3 +371,35 @@ def _request(self, method, urlpath, **kwargs):
         resp = self._session.request(method, ID_TOOLKIT_URL + urlpath, **kwargs)
         resp.raise_for_status()
         return resp.json()
+
+
+class UserIterator(object):
+    """An iterator that allows iterating over user accounts, one at a time.
+
+    This implementation loads a page of users into memory, and iterates on them. When the whole
+    page has been traversed, it loads another page. This class never keeps more than one page
+    of entries in memory.
+    """
+
+    def __init__(self, current_page):
+        if not current_page:
+            raise ValueError('Current page must not be None.')
+        self._current_page = current_page
+        self._index = 0
+
+    def next(self):
+        if self._index == len(self._current_page.users):
+            if self._current_page.has_next_page:
+                self._current_page = self._current_page.get_next_page()
+                self._index = 0
+        if self._index < len(self._current_page.users):
+            result = self._current_page.users[self._index]
+            self._index += 1
+            return result
+        raise StopIteration
+
+    def __next__(self):
+        return self.next()
+
+    def __iter__(self):
+        return self

firebase_admin/auth.py

@@ -19,6 +19,7 @@
 creating and managing user accounts in Firebase projects.
 """
 
+import json
 import time
 
 from google.auth import jwt
@@ -163,6 +164,36 @@ def get_user_by_phone_number(phone_number, app=None):
     except _user_mgt.ApiCallError as error:
         raise AuthError(error.code, str(error), error.detail)
 
+def list_users(page_token=None, max_results=_user_mgt.MAX_LIST_USERS_RESULTS, app=None):
+    """Retrieves a page of user accounts from a Firebase project.
+
+    The ``page_token`` argument governs the starting point of the page. The ``max_results``
+    argument governs the maximum number of user accounts that may be included in the returned page.
+    This function never returns None. If there are no user accounts in the Firebase project, this
+    returns an empty page.
+
+    Args:
+        page_token: A non-empty page token string, which indicates the starting point of the page
+            (optional). Defaults to ``None``, which will retrieve the first page of users.
+        max_results: A positive integer indicating the maximum number of users to include in the
+            returned page (optional). Defaults to 1000, which is also the maximum number allowed.
+        app: An App instance (optional).
+
+    Returns:
+        ListUsersPage: A ListUsersPage instance.
+
+    Raises:
+        ValueError: If max_results or page_token are invalid.
+        AuthError: If an error occurs while retrieving the user accounts.
+    """
+    user_manager = _get_auth_service(app).user_manager
+    def download(page_token, max_results):
+        try:
+            return user_manager.list_users(page_token, max_results)
+        except _user_mgt.ApiCallError as error:
+            raise AuthError(error.code, str(error), error.detail)
+    return ListUsersPage(download, page_token, max_results)
+
 
 def create_user(**kwargs):
     """Creates a new user account with the specified properties.
@@ -195,11 +226,12 @@ def create_user(**kwargs):
         raise AuthError(error.code, str(error), error.detail)
 
 
-def update_user(uid, **kwargs): # pylint: disable=missing-param-doc
+def update_user(uid, **kwargs):
     """Updates an existing user account with the specified properties.
 
     Args:
         uid: A user ID string.
+        kwargs: A series of keyword arguments (optional).
 
     Keyword Args:
         display_name: The user's display name (optional). Can be removed by explicitly passing
@@ -212,7 +244,8 @@ def update_user(uid, **kwargs): # pylint: disable=missing-param-doc
         photo_url: The user's photo URL (optional). Can be removed by explicitly passing None.
         password: The user's raw, unhashed password. (optional).
         disabled: A boolean indicating whether or not the user account is disabled (optional).
-        app: An App instance (optional).
+        custom_claims: A dictionary or a JSON string contining the custom claims to be set on the
+            user account (optional).
 
     Returns:
         UserRecord: An updated UserRecord instance for the user.
@@ -229,6 +262,31 @@ def update_user(uid, **kwargs): # pylint: disable=missing-param-doc
     except _user_mgt.ApiCallError as error:
         raise AuthError(error.code, str(error), error.detail)
 
+def set_custom_user_claims(uid, custom_claims, app=None):
+    """Sets additional claims on an existing user account.
+
+     Custom claims set via this function can be used to define user roles and privilege levels.
+     These claims propagate to all the devices where the user is already signed in (after token
+     expiration or when token refresh is forced), and next time the user signs in. The claims
+     can be accessed via the user's ID token JWT. If a reserved OIDC claim is specified (sub, iat,
+     iss, etc), an error is thrown. Claims payload must also not be larger then 1000 characters
+     when serialized into a JSON string.
+
+     Args:
+         uid: A user ID string.
+         custom_claims: A dictionary or a JSON string of custom claims. Pass None to unset any
+             claims set previously.
+         app: An App instance (optional).
+
+    Raises:
+        ValueError: If the specified user ID or the custom claims are invalid.
+        AuthError: If an error occurs while updating the user account.
+    """
+    user_manager = _get_auth_service(app).user_manager
+    try:
+        user_manager.update_user(uid, custom_claims=custom_claims)
+    except _user_mgt.ApiCallError as error:
+        raise AuthError(error.code, str(error), error.detail)
 
 def delete_user(uid, app=None):
     """Deletes the user identified by the specified user ID.
@@ -393,6 +451,20 @@ def provider_data(self):
         providers = self._data.get('providerUserInfo', [])
         return [_ProviderUserInfo(entry) for entry in providers]
 
+    @property
+    def custom_claims(self):
+        """Returns any custom claims set on this user account.
+
+        Returns:
+          dict: A dictionary of claims or None.
+        """
+        claims = self._data.get('customAttributes')
+        if claims:
+            parsed = json.loads(claims)
+            if parsed != {}:
+                return parsed
+        return None
+
 
 class UserMetadata(object):
     """Contains additional metadata associated with a user account."""
@@ -414,6 +486,85 @@ def last_sign_in_timestamp(self):
             return int(self._data['lastLoginAt'])
         return None
 
+class ExportedUserRecord(UserRecord):
+    """Contains metadata associated with a user including password hash and salt."""
+
+    def __init__(self, data):
+        super(ExportedUserRecord, self).__init__(data)
+
+    @property
+    def password_hash(self):
+        """The user's password hash as a base64-encoded string.
+
+        If the Firebase Auth hashing algorithm (SCRYPT) was used to create the user account, this
+        will be the base64-encoded password hash of the user. If a different hashing algorithm was
+        used to create this user, as is typical when migrating from another Auth system, this
+        will be an empty string. If no password is set, this will be None.
+        """
+        return self._data.get('passwordHash')
+
+    @property
+    def password_salt(self):
+        """The user's password salt as a base64-encoded string.
+
+        If the Firebase Auth hashing algorithm (SCRYPT) was used to create the user account, this
+        will be the base64-encoded password salt of the user. If a different hashing algorithm was
+        used to create this user, as is typical when migrating from another Auth system, this will
+        be an empty string. If no password is set, this will be None.
+        """
+        return self._data.get('salt')
+
+
+class ListUsersPage(object):
+    """Represents a page of user records exported from a Firebase project.
+
+    Provides methods for traversing the user accounts included in this page, as well as retrieving
+    subsequent pages of users. The iterator returned by ``iterate_all()`` can be used to iterate
+    through all users in the Firebase project starting from this page.
+    """
+
+    def __init__(self, download, page_token, max_results):
+        self._download = download
+        self._max_results = max_results
+        self._current = download(page_token, max_results)
+
+    @property
+    def users(self):
+        """A list of ``ExportedUserRecord`` instances available in this page."""
+        return [ExportedUserRecord(user) for user in self._current.get('users', [])]
+
+    @property
+    def next_page_token(self):
+        """Page token string for the next page (empty string indicates no more pages)."""
+        return self._current.get('nextPageToken', '')
+
+    @property
+    def has_next_page(self):
+        """A boolean indicating whether more pages are available."""
+        return bool(self.next_page_token)
+
+    def get_next_page(self):
+        """Retrieves the next page of user accounts, if available.
+
+        Returns:
+            ListUsersPage: Next page of users, or None if this is the last page.
+        """
+        if self.has_next_page:
+            return ListUsersPage(self._download, self.next_page_token, self._max_results)
+        return None
+
+    def iterate_all(self):
+        """Retrieves an iterator for user accounts.
+
+        Returned iterator will iterate through all the user accounts in the Firebase project
+        starting from this page. The iterator will never buffer more than one page of users
+        in memory at a time.
+
+        Returns:
+            iterator: An iterator of ExportedUserRecord instances.
+        """
+        return _user_mgt.UserIterator(self)
+
 
 class _ProviderUserInfo(UserInfo):
     """Contains metadata regarding how a user is known by a particular identity provider."""