--statement option for create and policy commands, refs #72

simonw · simonw · commit 54f810acfc16 · 2022-06-30T12:50:00.000-07:00
diff --git a/README.md b/README.md
@@ -93,6 +93,7 @@ The `create` command has a number of options:
 - `--read-only`: The user should only be allowed to read files from the bucket.
 - `--write-only`: The user should only be allowed to write files to the bucket, but not read them. This can be useful for logging and backups.
 - `--policy filepath-or-string`: A custom policy document (as a file path, literal JSON string or `-` for standard input) - see below.
+- `--statement json-statement`: Custom JSON statement block to be added to the generated policy.
 - `--bucket-region`: If creating buckets, the region in which they should be created.
 - `--silent`: Don't output details of what is happening, just output the JSON for the created access credentials at the end.
 - `--dry-run`: Output details of AWS changes that would have been made without applying them.
@@ -164,6 +165,14 @@ You can also pass `-` to read from standard input, or you can pass the literal J
   ]
 }'
 ```
+You can also specify one or more extra statement blocks that should be added to the generated policy, using `--statement JSON`. This example enables the AWS `textract:` APIs for the generated credentials, useful for using with the [s3-ocr](https://datasette.io/tools/s3-ocr) tool:
+```
+% s3-credentials create my-s3-bucket --statement '{
+  "Effect": "Allow",
+  "Action": "textract:*",
+  "Resource": "*"
+}'
+```
 
 ## Other commands
 
@@ -174,6 +183,7 @@ You can use the `s3-credentials policy` command to generate the JSON policy docu
 - `--read-only` - generate a read-only policy
 - `--write-only` - generate a write-only policy
 - `--prefix` - policy should be restricted to keys in the bucket that start with this prefix
+- `--statement json-statement`: Custom JSON statement block
 - `--public-bucket` - generate a bucket policy for a public bucket
 
 With none of these options it defaults to a read-write policy.
diff --git a/help.md b/help.md
@@ -80,6 +80,7 @@ Options:
   --policy POLICY                 Path to a policy.json file, or literal JSON
                                   string - $!BUCKET_NAME!$ will be replaced with
                                   the name of the bucket
+  --statement STATEMENT           JSON statement to add to the policy
   --bucket-region TEXT            Region in which to create buckets
   --silent                        Don't show performed steps
   --dry-run                       Show steps without executing them
@@ -301,11 +302,12 @@ Usage: s3-credentials policy [OPTIONS] BUCKETS...
       s3-credentials policy my-bucket --read-only
 
 Options:
-  --read-only      Only allow reading from the bucket
-  --write-only     Only allow writing to the bucket
-  --prefix TEXT    Restrict to keys starting with this prefix
-  --public-bucket  Bucket policy for allowing public access
-  --help           Show this message and exit.
+  --read-only            Only allow reading from the bucket
+  --write-only           Only allow writing to the bucket
+  --prefix TEXT          Restrict to keys starting with this prefix
+  --statement STATEMENT  JSON statement to add to the policy
+  --public-bucket        Bucket policy for allowing public access
+  --help                 Show this message and exit.
 ```
 ### s3-credentials put-object --help
 
diff --git a/s3_credentials/cli.py b/s3_credentials/cli.py
@@ -130,6 +130,27 @@ def convert(self, value, param, ctx):
         return integer
 
 
+class StatementParam(click.ParamType):
+    "Ensures statement is valid JSON with required fields"
+    name = "statement"
+
+    def convert(self, statement, param, ctx):
+        try:
+            data = json.loads(statement)
+        except ValueError:
+            self.fail("Invalid JSON string")
+        if not isinstance(data, dict):
+            self.fail("JSON must be an object")
+        missing_keys = {"Effect", "Action", "Resource"} - data.keys()
+        if missing_keys:
+            self.fail(
+                "Statement JSON missing required keys: {}".format(
+                    ", ".join(sorted(missing_keys))
+                )
+            )
+        return data
+
+
 @cli.command()
 @click.argument(
     "buckets",
@@ -141,12 +162,19 @@ def convert(self, value, param, ctx):
 @click.option(
     "--prefix", help="Restrict to keys starting with this prefix", default="*"
 )
+@click.option(
+    "extra_statements",
+    "--statement",
+    multiple=True,
+    type=StatementParam(),
+    help="JSON statement to add to the policy",
+)
 @click.option(
     "--public-bucket",
     help="Bucket policy for allowing public access",
     is_flag=True,
 )
-def policy(buckets, read_only, write_only, prefix, public_bucket):
+def policy(buckets, read_only, write_only, prefix, extra_statements, public_bucket):
     """
     Output generated JSON policy for one or more buckets
 
@@ -183,6 +211,8 @@ def policy(buckets, read_only, write_only, prefix, public_bucket):
             statements.extend(policies.write_only_statements(bucket, prefix))
     else:
         assert False, "Unknown permission: {}".format(permission)
+    if extra_statements:
+        statements.extend(extra_statements)
     bucket_access_policy = policies.wrap_policy(statements)
     click.echo(json.dumps(bucket_access_policy, indent=4))
 
@@ -229,6 +259,13 @@ def policy(buckets, read_only, write_only, prefix, public_bucket):
     type=PolicyParam(),
     help="Path to a policy.json file, or literal JSON string - $!BUCKET_NAME!$ will be replaced with the name of the bucket",
 )
+@click.option(
+    "extra_statements",
+    "--statement",
+    multiple=True,
+    type=StatementParam(),
+    help="JSON statement to add to the policy",
+)
 @click.option("--bucket-region", help="Region in which to create buckets")
 @click.option("--silent", help="Don't show performed steps", is_flag=True)
 @click.option("--dry-run", help="Show steps without executing them", is_flag=True)
@@ -252,6 +289,7 @@ def create(
     read_only,
     write_only,
     policy,
+    extra_statements,
     bucket_region,
     user_permissions_boundary,
     silent,
@@ -278,6 +316,7 @@ def create(
         raise click.ClickException(
             "Cannot use --read-only and --write-only at the same time"
         )
+    extra_statements = list(extra_statements)
 
     def log(message):
         if not silent:
@@ -348,7 +387,6 @@ def log(message):
                         log("Attached bucket policy allowing public access")
     # At this point the buckets definitely exist - create the inline policy for assume_role()
     assume_role_policy = {}
-    bucket_access_policy = {}
     if policy:
         assume_role_policy = json.loads(policy.replace("$!BUCKET_NAME!$", bucket))
     else:
@@ -364,6 +402,7 @@ def log(message):
                 statements.extend(policies.write_only_statements(bucket, prefix))
         else:
             assert False, "Unknown permission: {}".format(permission)
+        statements.extend(extra_statements)
         assume_role_policy = policies.wrap_policy(statements)
 
     if duration:
@@ -382,7 +421,7 @@ def log(message):
             credentials_response = sts.assume_role(
                 RoleArn=s3_role_arn,
                 RoleSessionName="s3.{permission}.{buckets}".format(
-                    permission="custom" if policy else permission,
+                    permission="custom" if (policy or extra_statements) else permission,
                     buckets=",".join(buckets),
                 ),
                 Policy=json.dumps(assume_role_policy),
@@ -410,7 +449,8 @@ def log(message):
     if not username:
         # Default username is "s3.read-write.bucket1,bucket2"
         username = "s3.{permission}.{buckets}".format(
-            permission="custom" if policy else permission, buckets=",".join(buckets)
+            permission="custom" if (policy or extra_statements) else permission,
+            buckets=",".join(buckets),
         )
     if dry_run or (not user_exists(iam, username)):
         kwargs = {"UserName": username}
@@ -443,18 +483,18 @@ def log(message):
     user_policy = {}
     for bucket in buckets:
         policy_name = "s3.{permission}.{bucket}".format(
-            permission="custom" if policy else permission,
+            permission="custom" if (policy or extra_statements) else permission,
             bucket=bucket,
         )
         if policy:
             user_policy = json.loads(policy.replace("$!BUCKET_NAME!$", bucket))
         else:
             if permission == "read-write":
-                user_policy = policies.read_write(bucket, prefix)
+                user_policy = policies.read_write(bucket, prefix, extra_statements)
             elif permission == "read-only":
-                user_policy = policies.read_only(bucket, prefix)
+                user_policy = policies.read_only(bucket, prefix, extra_statements)
             elif permission == "write-only":
-                user_policy = policies.write_only(bucket, prefix)
+                user_policy = policies.write_only(bucket, prefix, extra_statements)
             else:
                 assert False, "Unknown permission: {}".format(permission)
 
diff --git a/s3_credentials/policies.py b/s3_credentials/policies.py
@@ -1,5 +1,8 @@
-def read_write(bucket, prefix="*"):
-    return wrap_policy(read_write_statements(bucket, prefix=prefix))
+def read_write(bucket, prefix="*", extra_statements=None):
+    statements = read_write_statements(bucket, prefix=prefix)
+    if extra_statements:
+        statements.extend(extra_statements)
+    return wrap_policy(statements)
 
 
 def read_write_statements(bucket, prefix="*"):
@@ -15,8 +18,11 @@ def read_write_statements(bucket, prefix="*"):
     ]
 
 
-def read_only(bucket, prefix="*"):
-    return wrap_policy(read_only_statements(bucket, prefix))
+def read_only(bucket, prefix="*", extra_statements=None):
+    statements = read_only_statements(bucket, prefix=prefix)
+    if extra_statements:
+        statements.extend(extra_statements)
+    return wrap_policy(statements)
 
 
 def read_only_statements(bucket, prefix="*"):
@@ -70,8 +76,11 @@ def read_only_statements(bucket, prefix="*"):
     ]
 
 
-def write_only(bucket, prefix="*"):
-    return wrap_policy(write_only_statements(bucket, prefix))
+def write_only(bucket, prefix="*", extra_statements=None):
+    statements = write_only_statements(bucket, prefix=prefix)
+    if extra_statements:
+        statements.extend(extra_statements)
+    return wrap_policy(statements)
 
 
 def write_only_statements(bucket, prefix="*"):
diff --git a/tests/test_dry_run.py b/tests/test_dry_run.py
@@ -6,7 +6,7 @@
 
 
 def assert_match_with_wildcards(pattern, input):
-    # Pattern language is simple: '*' becomes '*?'
+    # Pattern language is simple: '*' becomes '.*?'
     bits = pattern.split("*")
     regex = "^{}$".format(".*?".join(re.escape(bit) for bit in bits))
     print(regex)
@@ -61,6 +61,17 @@ def assert_match_with_wildcards(pattern, input):
 Would call create access key for user 's3.read-write.my-bucket'"""
             ),
         ),
+        (
+            [
+                "--statement",
+                '{"Effect": "Allow", "Action": "textract:*", "Resource": "*"}',
+            ],
+            (
+                """Would create bucket: 'my-bucket'
+Would create user: 's3.custom.my-bucket' with permissions boundary: 'arn:aws:iam::aws:policy/AmazonS3FullAccess'
+*"Action": "textract:*"""
+            ),
+        ),
     ),
 )
 def test_dry_run(options, expected):
diff --git a/tests/test_s3_credentials.py b/tests/test_s3_credentials.py
@@ -282,16 +282,24 @@ def test_list_buckets_details(stub_s3):
 READ_ONLY_POLICY = '{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": ["s3:ListBucket", "s3:GetBucketLocation"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1"]}, {"Effect": "Allow", "Action": ["s3:GetObject", "s3:GetObjectAcl", "s3:GetObjectLegalHold", "s3:GetObjectRetention", "s3:GetObjectTagging"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/*"]}]}'
 WRITE_ONLY_POLICY = '{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": ["s3:PutObject"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/*"]}]}'
 PREFIX_POLICY = '{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": ["s3:GetBucketLocation"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1"]}, {"Effect": "Allow", "Action": ["s3:ListBucket"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1"], "Condition": {"StringLike": {"s3:prefix": ["my-prefix/*"]}}}, {"Effect": "Allow", "Action": ["s3:GetObject", "s3:GetObjectAcl", "s3:GetObjectLegalHold", "s3:GetObjectRetention", "s3:GetObjectTagging"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/my-prefix/*"]}, {"Effect": "Allow", "Action": ["s3:PutObject", "s3:DeleteObject"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/my-prefix/*"]}]}'
+EXTRA_STATEMENTS_POLICY = '{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": ["s3:ListBucket", "s3:GetBucketLocation"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1"]}, {"Effect": "Allow", "Action": ["s3:GetObject", "s3:GetObjectAcl", "s3:GetObjectLegalHold", "s3:GetObjectRetention", "s3:GetObjectTagging"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/*"]}, {"Effect": "Allow", "Action": ["s3:PutObject", "s3:DeleteObject"], "Resource": ["arn:aws:s3:::pytest-bucket-simonw-1/*"]}, {"Effect": "Allow", "Action": "textract:*", "Resource": "*"}]}'
 
 # Used by both test_create and test_create_duration
 CREATE_TESTS = (
+    # options,use_policy_stdin,expected_policy,expected_name_fragment
     ([], False, READ_WRITE_POLICY, "read-write"),
     (["--read-only"], False, READ_ONLY_POLICY, "read-only"),
     (["--write-only"], False, WRITE_ONLY_POLICY, "write-only"),
     (["--prefix", "my-prefix/"], False, PREFIX_POLICY, "read-write"),
     (["--policy", "POLICYFILEPATH"], False, CUSTOM_POLICY, "custom"),
     (["--policy", "-"], True, CUSTOM_POLICY, "custom"),
     (["--policy", CUSTOM_POLICY], False, CUSTOM_POLICY, "custom"),
+    (
+        ["--statement", '{"Effect": "Allow", "Action": "textract:*", "Resource": "*"}'],
+        False,
+        EXTRA_STATEMENTS_POLICY,
+        "custom",
+    ),
 )
 
 
@@ -344,6 +352,22 @@ def test_create(
         ]
 
 
+@pytest.mark.parametrize(
+    "statement,expected_error",
+    (
+        ("", "Invalid JSON string"),
+        ("{}", "missing required keys: Action, Effect, Resource"),
+        ('{"Action": 1}', "missing required keys: Effect, Resource"),
+        ('{"Action": 1, "Effect": 2}', "missing required keys: Resource"),
+    ),
+)
+def test_create_statement_error(statement, expected_error):
+    runner = CliRunner()
+    result = runner.invoke(cli, ["create", "--statement", statement])
+    assert result.exit_code == 2
+    assert expected_error in result.output
+
+
 @pytest.fixture
 def mocked_for_duration(mocker):
     boto3 = mocker.patch("boto3.client")
@@ -820,6 +844,13 @@ def test_auth_option_errors(extra_option):
         (["--read-only"], READ_ONLY_POLICY),
         (["--write-only"], WRITE_ONLY_POLICY),
         (["--prefix", "my-prefix/"], PREFIX_POLICY),
+        (
+            [
+                "--statement",
+                '{"Effect": "Allow", "Action": "textract:*", "Resource": "*"}',
+            ],
+            EXTRA_STATEMENTS_POLICY,
+        ),
     ),
 )
 def test_policy(options, expected):
