From 6ee932ea80f9202f28cb454e780130b00d768cb6 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 2 Nov 2021 19:30:43 -0700
Subject: [PATCH] s3-credentials list-buckets command, closes #9

---
 README.md                    |  8 ++++++++
 s3_credentials/cli.py        | 19 +++++++++++++++++++
 tests/test_s3_credentials.py | 22 ++++++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/README.md b/README.md
index 5a1c7a9..09a050f 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,14 @@ Add `--nl` to collapse these to single lines as valid newline-delimited JSON.
 
 Add `--array` to output a valid JSON array of objects instead.
 
+### list-buckets
+
+Shows a list of all buckets in your AWS account.
+
+    s3-credentials list-buckets
+
+Accepts the same `--nl` and `--array` options as `list-users`.
+
 ### list-user-policies
 
 To see a list of inline policies belonging to users:
diff --git a/s3_credentials/cli.py b/s3_credentials/cli.py
index d5b1478..9a89078 100644
--- a/s3_credentials/cli.py
+++ b/s3_credentials/cli.py
@@ -214,3 +214,22 @@ def list_user_policies(usernames):
                 click.echo(
                     json.dumps(policy_response["PolicyDocument"], indent=4, default=str)
                 )
+
+
+@cli.command()
+@click.option("--array", help="Output a valid JSON array", is_flag=True)
+@click.option("--nl", help="Output newline-delimited JSON", is_flag=True)
+def list_buckets(array, nl):
+    "List all buckets"
+    s3 = boto3.client("s3")
+    gathered = []
+    for bucket in s3.list_buckets()["Buckets"]:
+        if array:
+            gathered.append(bucket)
+        else:
+            if nl:
+                click.echo(json.dumps(bucket, default=str))
+            else:
+                click.echo(json.dumps(bucket, indent=4, default=str))
+    if gathered:
+        click.echo(json.dumps(gathered, indent=4, default=str))
diff --git a/tests/test_s3_credentials.py b/tests/test_s3_credentials.py
index 987e7d3..1b85775 100644
--- a/tests/test_s3_credentials.py
+++ b/tests/test_s3_credentials.py
@@ -39,6 +39,28 @@ def test_list_users(mocker, option, expected):
         assert result.output == expected
 
 
+@pytest.mark.parametrize(
+    "option,expected",
+    (
+        ("", '{\n    "name": "one"\n}\n{\n    "name": "two"\n}\n'),
+        (
+            "--array",
+            '[\n    {\n        "name": "one"\n    },\n'
+            '    {\n        "name": "two"\n    }\n]\n',
+        ),
+        ("--nl", '{"name": "one"}\n{"name": "two"}\n'),
+    ),
+)
+def test_list_buckets(mocker, option, expected):
+    boto3 = mocker.patch("boto3.client")
+    boto3().list_buckets.return_value = {"Buckets": [{"name": "one"}, {"name": "two"}]}
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(cli, ["list-buckets"] + ([option] if option else []))
+        assert result.exit_code == 0
+        assert result.output == expected
+
+
 def test_create(mocker):
     boto3 = mocker.patch("boto3.client")
     boto3.return_value = Mock()