Merge pull request Backblaze#1094 from reef-technologies/multi-bucket-keys

Add multi-bucket key support

Author: olzhasar-reef

b2/_internal/console_tool.py

@@ -1706,7 +1706,8 @@ class KeyCreateBase(Command):
     specified, the key will not expire.
 
     The ``bucket`` is the name of a bucket in the account.  When specified, the key
-    will only allow access to that bucket.
+    will only allow access to that bucket. You can specify multiple buckets by repeating
+    the option, e.g. ``--bucket bucket1 --bucket bucket2``
 
     The ``namePrefix`` restricts file access to files whose names start with the prefix.
 
@@ -1720,7 +1721,7 @@ class KeyCreateBase(Command):
 
     @classmethod
     def _setup_parser(cls, parser):
-        parser.add_argument('--bucket')
+        parser.add_argument('--bucket', action='append')
         add_normalized_argument(parser, '--name-prefix')
         parser.add_argument('--duration', type=int)
         parser.add_argument('keyName')
@@ -1731,11 +1732,21 @@ def _setup_parser(cls, parser):
         super()._setup_parser(parser)
 
     def _run(self, args):
-        # Translate the bucket name into a bucketId
-        if args.bucket is None:
-            bucket_id_or_none = None
+        # Translate a list of bucket names into a list of bucketIds
+        if not args.bucket:
+            bucket_ids_or_none = None
+        elif len(args.bucket) == 1:
+            bucket_ids_or_none = [self.api.get_bucket_by_name(args.bucket[0]).id_]
         else:
-            bucket_id_or_none = self.api.get_bucket_by_name(args.bucket).id_
+            names_seeking = set(args.bucket)
+            bucket_ids_or_none = []
+            for bucket in self.api.list_buckets(use_cache=True):
+                if bucket.name in names_seeking:
+                    names_seeking.remove(bucket.name)
+                    bucket_ids_or_none.append(bucket.id_)
+
+            if names_seeking:
+                raise NonExistentBucket('; '.join(names_seeking))
 
         if args.all_capabilities:
             current_key_caps = set(self.api.account_info.get_allowed()['capabilities'])
@@ -1747,13 +1758,11 @@ def _run(self, args):
                 set(ALL_CAPABILITIES) - preview_feature_caps | current_key_caps
             )
 
-        buckets_ids = [bucket_id_or_none] if bucket_id_or_none else None
-
         application_key = self.api.create_key(
             capabilities=args.capabilities,
             key_name=args.keyName,
             valid_duration_seconds=args.duration,
-            bucket_ids=buckets_ids,
+            bucket_ids=bucket_ids_or_none,
             name_prefix=args.name_prefix,
         )
 
@@ -2318,7 +2327,7 @@ class KeyListBase(Command):
 
     - ID of the application key
     - Name of the application key
-    - Name of the bucket the key is restricted to, or ``-`` for no restriction
+    - Name of the bucket(s) the key is restricted to, or ``-`` for no restriction
     - Date of expiration, or ``-``
     - Time of expiration, or ``-``
     - File name prefix, in single quotes
@@ -2352,35 +2361,35 @@ def _run(self, args):
 
     def print_key(self, key: ApplicationKey, is_long_format: bool):
         if is_long_format:
-            format_str = "{keyId}   {keyName:20s}   {bucketName:20s}   {dateStr:10s}   {timeStr:8s}   '{namePrefix}'   {capabilities}"
+            format_str = "{keyId}   {keyName:20s}   {bucketNames:20s}   {dateStr:10s}   {timeStr:8s}   '{namePrefix}'   {capabilities}"
         else:
             format_str = '{keyId}   {keyName:20s}'
         timestamp_or_none = apply_or_none(int, key.expiration_timestamp_millis)
         (date_str, time_str) = self.timestamp_display(timestamp_or_none)
 
-        bucket_id = key.bucket_ids[0] if key.bucket_ids else None
-
         key_str = format_str.format(
             keyId=key.id_,
             keyName=key.key_name,
-            bucketName=self.bucket_display_name(bucket_id),
+            bucketNames=self.bucket_display_names(key.bucket_ids),
             namePrefix=(key.name_prefix or ''),
             capabilities=','.join(key.capabilities),
             dateStr=date_str,
             timeStr=time_str,
         )
         self._print(key_str)
 
-    def bucket_display_name(self, bucket_id):
-        # Special case for no bucket ID
-        if bucket_id is None:
+    def bucket_display_names(self, bucket_ids):
+        # Special case for no bucket IDs
+        if not bucket_ids:
             return '-'
 
         # Make sure we have the map
         if self.bucket_id_to_bucket_name is None:
             self.bucket_id_to_bucket_name = dict((b.id_, b.name) for b in self.api.list_buckets())
 
-        return self.bucket_id_to_bucket_name.get(bucket_id, 'id=' + bucket_id)
+        display_names = [self.bucket_id_to_bucket_name.get(id_, f'id={id_}') for id_ in bucket_ids]
+
+        return ', '.join(display_names)
 
     def timestamp_display(self, timestamp_or_none):
         """

changelog.d/+key-list-multi-bucket.added.md

@@ -0,0 +1 @@
+Support multi-bucket keys in `key list` subcommand.

changelog.d/1083.added.md

@@ -0,0 +1 @@
+Add multi-bucket keys support to the `key create` subcommand.

test/integration/test_b2_command_line.py

@@ -782,6 +782,45 @@ def test_key_restrictions(b2_tool, bucket_name, sample_file, bucket_factory, b2_
     )
 
 
+def test_multi_bucket_key_restrictions(b2_tool, bucket_factory):
+    bucket_a = bucket_factory()
+    bucket_b = bucket_factory()
+    bucket_c = bucket_factory()
+
+    key_name = 'clt-testKey-01' + random_hex(6)
+
+    created_key_stdout = b2_tool.should_succeed(
+        [
+            'key',
+            'create',
+            '--bucket',
+            bucket_a.name,
+            '--bucket',
+            bucket_b.name,
+            key_name,
+            'listFiles,listBuckets,readFiles',
+        ]
+    )
+
+    mb_key_id, mb_key = created_key_stdout.split()
+
+    b2_tool.should_succeed(
+        ['account', 'authorize', '--environment', b2_tool.realm, mb_key_id, mb_key],
+    )
+
+    b2_tool.should_succeed(
+        ['bucket', 'get', bucket_a.name],
+    )
+    b2_tool.should_succeed(
+        ['bucket', 'get', bucket_b.name],
+    )
+
+    failed_bucket_err = rf"ERROR: Application key is restricted to buckets: \['{bucket_a.name}', '{bucket_b.name}'\]"
+    b2_tool.should_fail(['bucket', 'get', bucket_c.name], failed_bucket_err)
+
+    b2_tool.should_succeed(['key', 'delete', mb_key_id])
+
+
 def test_delete_bucket(b2_tool, bucket_name):
     b2_tool.should_succeed(['bucket', 'delete', bucket_name])
     b2_tool.should_fail(

test/unit/test_console_tool.py

@@ -581,6 +581,62 @@ def test_create_bucket_key_and_authorize_with_it(self):
             0,
         )
 
+    def test_create_multi_bucket_key_and_authorize_with_it(self):
+        # Start with authorizing with the master key
+        self._authorize_account()
+
+        # Make two buckets
+        self._run_command(['bucket', 'create', 'my-bucket-0', 'allPrivate'], 'bucket_0\n', '', 0)
+        self._run_command(['bucket', 'create', 'my-bucket-1', 'allPrivate'], 'bucket_1\n', '', 0)
+
+        # Create a key restricted to those buckets
+        self._run_command(
+            [
+                'key',
+                'create',
+                '--bucket',
+                'my-bucket-0',
+                '--bucket',
+                'my-bucket-1',
+                'key1',
+                'listKeys,listBuckets',
+            ],
+            'appKeyId0 appKey0\n',
+            '',
+            0,
+        )
+
+        # test deprecated command
+        self._run_command(
+            [
+                'create-key',
+                '--bucket',
+                'my-bucket-0',
+                '--bucket',
+                'my-bucket-1',
+                'key2',
+                'listKeys,listBuckets',
+            ],
+            'appKeyId1 appKey1\n',
+            'WARNING: `create-key` command is deprecated. Use `key create` instead.\n',
+            0,
+        )
+
+        # Authorize with the key
+        self._run_command(
+            ['account', 'authorize', 'appKeyId0', 'appKey0'],
+            None,
+            '',
+            0,
+        )
+
+        self._run_command(
+            ['account', 'authorize', 'appKeyId1', 'appKey1'],
+            None,
+            '',
+            0,
+        )
+
     def test_update_bucket_without_lifecycle(self):
         # Start with authorizing with the master key
         self._authorize_account()
@@ -1028,6 +1084,109 @@ def test_keys(self):
             1,
         )
 
+    def test_multi_bucket_keys(self):
+        self._authorize_account()
+
+        self._run_command(['bucket', 'create', 'my-bucket-a', 'allPublic'], 'bucket_0\n', '', 0)
+        self._run_command(['bucket', 'create', 'my-bucket-b', 'allPublic'], 'bucket_1\n', '', 0)
+        self._run_command(['bucket', 'create', 'my-bucket-c', 'allPublic'], 'bucket_2\n', '', 0)
+
+        capabilities = ['readFiles', 'listBuckets']
+        capabilities_with_commas = ','.join(capabilities)
+
+        # Create a multi-bucket key with one of the buckets having invalid name
+        expected_stderr = 'Bucket not found: invalid. If you believe it exists, run `b2 bucket list` to reset cache, then try again.\n'
+
+        self._run_command(
+            [
+                'key',
+                'create',
+                '--bucket',
+                'my-bucket-a',
+                '--bucket',
+                'invalid',
+                'goodKeyName',
+                capabilities_with_commas,
+            ],
+            '',
+            expected_stderr,
+            1,
+        )
+
+        # Create a multi-bucket key
+        self._run_command(
+            [
+                'key',
+                'create',
+                '--bucket',
+                'my-bucket-a',
+                '--bucket',
+                'my-bucket-b',
+                'goodKeyName',
+                capabilities_with_commas,
+            ],
+            'appKeyId0 appKey0\n',
+            '',
+            0,
+        )
+
+        # List keys
+        expected_list_keys_out = 'appKeyId0   goodKeyName\n'
+
+        expected_list_keys_out_long = """
+            appKeyId0   goodKeyName            my-bucket-a, my-bucket-b   -            -          ''   readFiles,listBuckets
+            """
+
+        self._run_command(['key', 'list'], expected_list_keys_out, '', 0)
+        self._run_command(['key', 'list', '--long'], expected_list_keys_out_long, '', 0)
+
+        # authorize and make calls using an application key with bucket restrictions
+        self._run_command(['account', 'authorize', 'appKeyId0', 'appKey0'], None, '', 0)
+
+        self._run_command(
+            ['bucket', 'list'],
+            '',
+            "ERROR: Application key is restricted to buckets: ['my-bucket-a', 'my-bucket-b']\n",
+            1,
+        )
+        self._run_command(
+            ['bucket', 'get', 'my-bucket-c'],
+            '',
+            "ERROR: Application key is restricted to buckets: ['my-bucket-a', 'my-bucket-b']\n",
+            1,
+        )
+
+        def _get_expected_json(bucket_id: str, bucket_name: str):
+            return {
+                'accountId': self.account_id,
+                'bucketId': bucket_id,
+                'bucketInfo': {},
+                'bucketName': bucket_name,
+                'bucketType': 'allPublic',
+                'corsRules': [],
+                'defaultServerSideEncryption': {'mode': None},
+                'lifecycleRules': [],
+                'options': [],
+                'revision': 1,
+            }
+
+        self._run_command(
+            ['bucket', 'get', 'my-bucket-a'],
+            expected_json_in_stdout=_get_expected_json('bucket_0', 'my-bucket-a'),
+        )
+
+        self._run_command(
+            ['bucket', 'get', 'my-bucket-b'],
+            expected_json_in_stdout=_get_expected_json('bucket_1', 'my-bucket-b'),
+        )
+
+        self._run_command(
+            ['ls', '--json', *self.b2_uri_args('my-bucket-c')],
+            '',
+            "ERROR: Application key is restricted to buckets: ['my-bucket-a', 'my-bucket-b']\n",
+            1,
+        )
+
     def test_bucket_info_from_json(self):
         self._authorize_account()
         self._run_command(['bucket', 'create', 'my-bucket', 'allPublic'], 'bucket_0\n', '', 0)