Adapt multiple-bucket keys flow to authorize-account api changes

Author: olzhasar-reef

b2sdk/_internal/account_info/abstract.py

@@ -31,8 +31,7 @@ class AbstractAccountInfo(metaclass=B2TraceMetaAbstract):
 
     # The 'allowed' structure to use for old account info that was saved without 'allowed'.
     DEFAULT_ALLOWED = dict(
-        bucketIds=None,
-        bucketNames=None,
+        buckets=None,
         capabilities=ALL_CAPABILITIES,
         namePrefix=None,
     )
@@ -337,28 +336,15 @@ def set_auth_data(
     @classmethod
     def allowed_is_valid(cls, allowed):
         """
-        Make sure that all of the required fields are present, and that
-        bucketIds field is set if bucketNames is.
+        Make sure that all of the required fields are present
 
         If the bucketId is for a bucket that no longer exists, or the
         capabilities do not allow for listBuckets, then we will not have a bucketName.
 
         :param dict allowed: the structure to use for old account info that was saved without 'allowed'
         :rtype: bool
         """
-        return (
-            ('bucketIds' in allowed)
-            and ('bucketNames' in allowed)
-            and (
-                (allowed['bucketIds'] is None and allowed['bucketNames'] is None)
-                or (
-                    allowed['bucketIds']
-                    and (len(allowed['bucketIds']) == len(allowed['bucketNames']))
-                )
-            )
-            and ('capabilities' in allowed)
-            and ('namePrefix' in allowed)
-        )
+        return ('buckets' in allowed) and ('capabilities' in allowed) and ('namePrefix' in allowed)
 
     @abstractmethod
     def _set_auth_data(

b2sdk/_internal/account_info/sqlite_account_info.py

@@ -408,8 +408,10 @@ def _migrate_allowed_to_multi_bucket(self):
         bucket_id = allowed.pop('bucketId')
         bucket_name = allowed.pop('bucketName')
 
-        allowed['bucketIds'] = [bucket_id] if bucket_id is not None else None
-        allowed['bucketNames'] = [bucket_name] if bucket_name is not None else None
+        if bucket_id is not None:
+            allowed['buckets'] = [{'id': bucket_id, 'name': bucket_name}]
+        else:
+            allowed['buckets'] = None
 
         allowed_text = json.dumps(allowed)
         stmt = f"UPDATE account SET allowed = ('{allowed_text}');"
@@ -590,8 +592,7 @@ def get_allowed(self):
         .. code-block:: python
 
             {
-                "bucketIds": null,
-                "bucketNames": null,
+                "buckets": null,
                 "capabilities": [
                     "listKeys",
                     "writeKeys"

b2sdk/_internal/api.py

@@ -636,7 +636,7 @@ def check_bucket_name_restrictions(self, bucket_name: str):
 
         :raises b2sdk.v2.exception.RestrictedBucket: if the account is not allowed to use this bucket
         """
-        self._check_bucket_restrictions('bucketNames', bucket_name)
+        self._check_bucket_restrictions('name', bucket_name)
 
     def check_bucket_id_restrictions(self, bucket_id: str):
         """
@@ -647,14 +647,21 @@ def check_bucket_id_restrictions(self, bucket_id: str):
 
         :raises b2sdk.v2.exception.RestrictedBucket: if the account is not allowed to use this bucket
         """
-        self._check_bucket_restrictions('bucketIds', bucket_id)
+        self._check_bucket_restrictions('id', bucket_id)
 
     def _check_bucket_restrictions(self, key, value):
-        allowed = self.account_info.get_allowed()
-        allowed_bucket_identifier = allowed[key]
+        buckets = self.account_info.get_allowed()['buckets']
 
-        if allowed_bucket_identifier and value not in allowed_bucket_identifier:
-            raise RestrictedBucket(allowed_bucket_identifier)
+        if not buckets:
+            return
+
+        for item in buckets:
+            if item[key] == value:
+                return
+
+        msg = str([b['name'] for b in buckets])
+
+        raise RestrictedBucket(msg)
 
     def _populate_bucket_cache_from_key(self):
         # If the key is restricted to the bucket, pre-populate the cache with it
@@ -663,16 +670,16 @@ def _populate_bucket_cache_from_key(self):
         except MissingAccountData:
             return
 
-        allowed_bucket_ids = allowed.get('bucketIds')
-        if not allowed_bucket_ids:
+        allowed_buckets = allowed.get('buckets')
+        if not allowed_buckets:
             return
 
-        allowed_bucket_names = allowed.get('bucketNames')
-
         # If we have bucketId set we still need to check bucketName. If the bucketName is None,
         # it means that the bucketId belongs to a bucket that was already removed.
-        if None in allowed_bucket_names:
-            raise RestrictedBucketMissing
 
-        for _id, name in zip(allowed_bucket_ids, allowed_bucket_names):
-            self.cache.save_bucket(self.BUCKET_CLASS(self, _id, name=name))
+        for item in allowed_buckets:
+            if item['name'] is None:
+                raise RestrictedBucketMissing
+
+        for item in allowed_buckets:
+            self.cache.save_bucket(self.BUCKET_CLASS(self, item['id'], name=item['name']))

b2sdk/_internal/raw_simulator.py

@@ -99,8 +99,7 @@ def __init__(
         key,
         capabilities,
         expiration_timestamp_or_none,
-        bucket_ids_or_none,
-        bucket_names_or_none,
+        buckets_or_none,
         name_prefix_or_none,
     ):
         self.name = name
@@ -109,15 +108,19 @@ def __init__(
         self.key = key
         self.capabilities = capabilities
         self.expiration_timestamp_or_none = expiration_timestamp_or_none
-        self.bucket_ids_or_none = bucket_ids_or_none
-        self.bucket_names_or_none = bucket_names_or_none
+        self.buckets_or_none = buckets_or_none
         self.name_prefix_or_none = name_prefix_or_none
 
+    def _get_bucket_ids(self):
+        if self.buckets_or_none is None:
+            return None
+
+        return [item['id'] for item in self.buckets_or_none]
+
     def as_key(self):
         return dict(
             accountId=self.account_id,
-            bucketIds=self.bucket_ids_or_none,
-            bucketNames=self.bucket_names_or_none,
+            bucketIds=self._get_bucket_ids(),
             applicationKeyId=self.application_key_id,
             capabilities=self.capabilities,
             expirationTimestamp=self.expiration_timestamp_or_none
@@ -134,15 +137,15 @@ def as_created_key(self):
         """
         result = self.as_key()
         result['applicationKey'] = self.key
+
         return result
 
     def get_allowed(self):
         """
         Return the 'allowed' structure to include in the response from b2_authorize_account.
         """
         return dict(
-            bucketIds=self.bucket_ids_or_none,
-            bucketName=self.bucket_names_or_none,
+            buckets=self.buckets_or_none,
             capabilities=self.capabilities,
             namePrefix=self.name_prefix_or_none,
         )
@@ -1372,8 +1375,7 @@ def create_account(self):
             key=master_key,
             capabilities=ALL_CAPABILITIES,
             expiration_timestamp_or_none=None,
-            bucket_ids_or_none=None,
-            bucket_names_or_none=None,
+            buckets_or_none=None,
             name_prefix_or_none=None,
         )
 
@@ -1401,17 +1403,6 @@ def authorize_account(self, realm_url, application_key_id, application_key):
         self.auth_token_to_key[auth_token] = key_sim
 
         allowed = key_sim.get_allowed()
-        bucketIds = allowed.get('bucketIds')
-
-        if bucketIds is not None:
-            allowed['bucketNames'] = []
-            for _id in bucketIds:
-                if _id in self.bucket_id_to_bucket:
-                    allowed['bucketNames'].append(self.bucket_id_to_bucket[_id].bucket_name)
-                else:
-                    allowed['bucketNames'].append(None)
-        else:
-            allowed['bucketNames'] = None
 
         return dict(
             accountId=key_sim.account_id,
@@ -1425,8 +1416,6 @@ def authorize_account(self, realm_url, application_key_id, application_key):
                     absoluteMinimumPartSize=self.MIN_PART_SIZE,
                     allowed=allowed,
                     s3ApiUrl=self.S3_API_URL,
-                    bucketIds=allowed['bucketIds'],
-                    bucketNames=allowed['bucketNames'],
                     capabilities=allowed['capabilities'],
                     namePrefix=allowed['namePrefix'],
                 ),
@@ -1506,16 +1495,18 @@ def create_key(
         self.app_key_counter += 1
         application_key_id = 'appKeyId%d' % (index,)
         app_key = 'appKey%d' % (index,)
-        bucket_names_or_none = None
+
+        buckets = None
+
         if bucket_ids is not None:
             # It is possible for bucketId to be filled and bucketName to be empty.
             # It can happen when the bucket was deleted.
-            bucket_names_or_none = []
+            buckets = []
             for _id in bucket_ids:
                 try:
-                    bucket_names_or_none.append(self._get_bucket_by_id(_id).bucket_name)
+                    buckets.append({'id': _id, 'name': self._get_bucket_by_id(_id).bucket_name})
                 except NonExistentBucket:
-                    bucket_names_or_none.append(None)
+                    buckets.append({'id': _id, 'name': None})
 
         key_sim = KeySimulator(
             account_id=account_id,
@@ -1524,8 +1515,7 @@ def create_key(
             key=app_key,
             capabilities=capabilities,
             expiration_timestamp_or_none=expiration_timestamp_or_none,
-            bucket_ids_or_none=bucket_ids,
-            bucket_names_or_none=bucket_names_or_none,
+            buckets_or_none=buckets,
             name_prefix_or_none=name_prefix,
         )
         self.key_id_to_key[application_key_id] = key_sim
@@ -2113,8 +2103,17 @@ def _assert_account_auth(
             raise InvalidAuthToken('auth token expired', 'auth_token_expired')
         if capability not in key_sim.capabilities:
             raise Unauthorized('', 'unauthorized')
-        if key_sim.bucket_ids_or_none and bucket_id not in key_sim.bucket_ids_or_none:
-            raise Unauthorized('', 'unauthorized')
+
+        if key_sim.buckets_or_none:
+            found = False
+            for item in key_sim.buckets_or_none:
+                if item['id'] == bucket_id:
+                    found = True
+                    break
+
+            if not found:
+                raise Unauthorized('', 'unauthorized')
+
         if key_sim.name_prefix_or_none is not None:
             if file_name is not None and not file_name.startswith(key_sim.name_prefix_or_none):
                 raise Unauthorized('', 'unauthorized')

b2sdk/_internal/session.py

@@ -143,12 +143,7 @@ def _construct_allowed_dict(self, storage_api_info):
         # `allowed` object has been deprecated in the v3 of the API, but we still
         # construct it artificially to avoid changes in all the reliant parts.
 
-        return {
-            'bucketIds': storage_api_info['bucketIds'],
-            'bucketNames': storage_api_info['bucketNames'],
-            'capabilities': storage_api_info['capabilities'],
-            'namePrefix': storage_api_info['namePrefix'],
-        }
+        return storage_api_info['allowed']
 
     def cancel_large_file(self, file_id):
         return self._wrap_default_token(self.raw_api.cancel_large_file, file_id)
@@ -510,15 +505,17 @@ def _add_app_key_info_to_unauthorized(self, unauthorized):
         # What's allowed?
         allowed = self.account_info.get_allowed()
         capabilities = allowed['capabilities']
-        bucket_name = allowed['bucketName']
         name_prefix = allowed['namePrefix']
 
         # Make a list of messages about the application key restrictions
         key_messages = []
         if set(capabilities) != set(ALL_CAPABILITIES):
             key_messages.append("with capabilities '" + ','.join(capabilities) + "'")
-        if bucket_name is not None:
-            key_messages.append("restricted to bucket '" + bucket_name + "'")
+
+        allowed_buckets_msg = self._get_allowed_buckets_message(allowed)
+        if allowed_buckets_msg:
+            key_messages.append(allowed_buckets_msg)
+
         if name_prefix is not None:
             key_messages.append("restricted to files that start with '" + name_prefix + "'")
         if not key_messages:
@@ -530,6 +527,15 @@ def _add_app_key_info_to_unauthorized(self, unauthorized):
 
         return Unauthorized(new_message, unauthorized.code)
 
+    def _get_allowed_buckets_message(self, allowed) -> str | None:
+        buckets = allowed['buckets']
+        if not buckets:
+            return None
+
+        bucket_names = [b['name'] for b in buckets]
+
+        return f'restricted to buckets {bucket_names}'
+
     def _get_upload_data(self, bucket_id):
         """
         Take ownership of an upload URL / auth token for the bucket and

b2sdk/v2/account_info.py

@@ -77,17 +77,15 @@ def get_allowed(self):
 
         # convert a multi-bucket key to a single bucket
 
-        if 'bucketIds' in allowed:
-            bucket_ids = allowed.pop('bucketIds')
-            if bucket_ids and len(bucket_ids) > 1:
+        if 'buckets' in allowed:
+            buckets = allowed.pop('buckets')
+            if buckets and len(buckets) > 1:
                 raise MissingAccountData(
                     'Multi-bucket keys cannot be used with the current sdk version'
                 )
 
-            allowed['bucketId'] = bucket_ids[0] if bucket_ids else None
-
-            bucket_names = allowed.pop('bucketNames')
-            allowed['bucketName'] = bucket_names[0] if bucket_names else None
+            allowed['bucketId'] = buckets[0]['id'] if buckets else None
+            allowed['bucketName'] = buckets[0]['name'] if buckets else None
 
         return allowed
 

b2sdk/v2/session.py

@@ -89,3 +89,10 @@ def _construct_allowed_dict(self, storage_api_info):
             'capabilities': storage_api_info['capabilities'],
             'namePrefix': storage_api_info['namePrefix'],
         }
+
+    def _get_allowed_buckets_message(self, allowed) -> str | None:
+        bucket_name = allowed['bucketName']
+        if bucket_name is None:
+            return None
+
+        return "restricted to bucket '" + bucket_name + "'"

test/unit/account_info/test_account_info.py

@@ -197,8 +197,7 @@ def allowed(self, apiver):
         return dict(
             capabilities=['readFiles'],
             namePrefix=None,
-            bucketIds=None,
-            bucketNames=None,
+            buckets=None,
         )
 
     def test_set_auth_data_compatibility(self, account_info_default_data, allowed):

test/unit/account_info/test_sqlite_account_info.py

@@ -106,8 +106,7 @@ def test_migrate_to_5_v2(
                     namePrefix=None,
                 ),
                 dict(
-                    bucketIds=['123'],
-                    bucketNames=['bucket1'],
+                    buckets=[{'id': '123', 'name': 'bucket1'}],
                     capabilities=['listBuckets', 'readBuckets'],
                     namePrefix=None,
                 ),
@@ -120,8 +119,7 @@ def test_migrate_to_5_v2(
                     namePrefix=None,
                 ),
                 dict(
-                    bucketIds=None,
-                    bucketNames=None,
+                    buckets=None,
                     capabilities=['listBuckets', 'readBuckets'],
                     namePrefix=None,
                 ),
@@ -153,8 +151,7 @@ def test_multi_bucket_key_error_apiver_v2(
         account_info_default_data,
     ):
         allowed = dict(
-            bucketIds=[1, 2],
-            bucketNames=['bucket1', 'bucket2'],
+            buckets=[{'id': 1, 'name': 'bucket1'}, {'id': 2, 'name': 'bucket2'}],
             capabilities=['listBuckets', 'readBuckets'],
             namePrefix=None,
         )