Feature/assume role section

README.md

@@ -177,6 +177,14 @@ skip-resource-constraints:
 # Exclude actions from the output by specifying them here. Accepts wildcards, like kms:Delete*
 exclude-actions:
 - ''
+# If this policy needs to include an AssumeRole action
+sts:
+  assume-role:
+    - ''
+  assume-role-with-saml:
+    - ''
+  assume-role-with-web-identity:
+    - ''
 ```
 
 ### Step 2: Copy/paste ARNs

docs/writing-policies/assume-role-actions.md

@@ -0,0 +1,54 @@
+Assume Role Actions
+============
+
+Using CRUD mode, you can add STS assume role actions to your larger CRUD policy by using the `sts` section.
+
+**Input**:
+
+```yaml
+mode: crud
+sts:
+  assume-role:
+    - arn:aws:iam::123456789012:role/demo
+    - arn:aws:iam::123456789012:role/demo2
+  assume-role-with-saml:
+    - arn:aws:iam::123456789012:role/demo3
+  assume-role-with-web-identity:
+    - 'arn:aws:iam::123456789012:role/demo'
+```
+
+**Output**:
+
+```json
+{
+    "Sid": "AssumeRole",
+    "Effect": "Allow",
+    "Action": [
+        "sts:AssumeRole"
+    ],
+    "Resource": [
+        "arn:aws:iam::123456789012:role/demo",
+        "arn:aws:iam::123456789012:role/demo2"
+    ]
+},
+{
+    "Sid": "AssumeRoleWithSAML",
+    "Effect": "Allow",
+    "Action": [
+        "sts:AssumeRoleWithSAML"
+    ],
+    "Resource": [
+        "arn:aws:iam::123456789012:role/demo3"
+    ]
+},
+{
+    "Sid": "AssumeRoleWithWebIdentity",
+    "Effect": "Allow",
+    "Action": [
+        "sts:AssumeRoleWithWebIdentity"
+    ],
+    "Resource": [
+        "arn:aws:iam::123456789012:role/demo"
+    ]
+}
+```

docs/writing-policies/crud-mode.md

@@ -61,6 +61,13 @@ wildcard-only:
   - ''
   service-permissions-management:
   - ''
+sts:
+  assume-role:
+    - ''
+  assume-role-with-saml:
+    - ''
+  assume-role-with-web-identity:
+    - ''
 ```
 
 -   Then just fill it out:
@@ -92,6 +99,13 @@ wildcard-only:
   - ''
   service-permissions-management:
   - ''
+sts:
+  assume-role:
+    - ''
+  assume-role-with-saml:
+    - ''
+  assume-role-with-web-identity:
+    - ''
 ```
 
 -   Run the command:

mkdocs.yml

@@ -34,6 +34,7 @@ nav:
       - "<i>Example</i>: Single actions": 'writing-policies/wildcard-only/example-single-actions.md'
       - "<i>Example</i>: Service-wide": 'writing-policies/wildcard-only/example-service-wide.md'
       - "<i>Example</i>: Combining approaches": 'writing-policies/wildcard-only/combined-approaches.md'
+      - "<i>Example</i>: AssumeRole Actions": 'writing-policies/assume-role-actions.md'
       - "<b>Wildcard-only</b>":
         - Introduction: 'writing-policies/wildcard-only/introduction.md'
         - Single actions: 'writing-policies/wildcard-only/single-actions.md'

policy_sentry/writing/sid_group.py

@@ -108,6 +108,56 @@ def add_skip_resource_constraints(self, skip_resource_constraints_actions):
         else:
             raise Exception("Please provide 'skip_resource_constraints' as a list of IAM actions.")
 
+    def add_sts_actions(self, sts_actions):
+        """
+        To add STS actions to the output from special YAML section
+        """
+        if sts_actions:
+            # Hard coded for this special case
+            service_prefix = "sts"
+            access_level = "Write"
+
+            for action, arns in sts_actions.items():
+                clean_action = action.replace('-','') # Convention to follow adding dashes instead of CamelCase
+                service_action_data = get_action_data(service_prefix, clean_action)
+
+                # Schema validation takes care of this, but just in case. No data returned for the action
+                if not service_action_data:
+                    raise Exception(f"Could not find service action data for {service_prefix} - {clean_action}")
+
+                for row in service_action_data[service_prefix]:
+                    for arn in arns:
+                        if not arn: # skip the - '' situation
+                            continue
+                        if (
+                                does_arn_match(arn, row["resource_arn_format"])
+                                and row["access_level"] == access_level
+                            ):
+                            raw_arn_format = row["resource_arn_format"]
+
+                            # Each action will get its own namespace sts:AssumeRole -> AssumeRole
+                            # -1 index is a neat trick if the colon ever goes away we won't get an index error.
+                            sid_namespace = row["action"].split(':')[-1]
+
+                            temp_sid_dict = {
+                                "arn": [arn],
+                                "service": service_prefix,
+                                "access_level": access_level,
+                                "arn_format": raw_arn_format,
+                                "actions": [row["action"]],
+                                "conditions": [],  # TODO: Add conditions
+                            }
+
+                        # Using a custom namespace and not gathering actions so no need to find
+                        # dependent actions either, though we could do it here    
+
+                        if sid_namespace in self.sids.keys():
+                            # If the ARN already exists there, skip it.
+                            if arn not in self.sids[sid_namespace]["arn"]:
+                                self.sids[sid_namespace]["arn"].append(arn)
+                        else:
+                            self.sids[sid_namespace] = temp_sid_dict
+
     def add_requested_service_wide(self, service_prefixes, access_level):
         """
         When a user requests all wildcard-only actions available under a service at a specific access level
@@ -486,6 +536,14 @@ def process_template(self, cfg, minimize=None):
                         self.add_action_without_resource_constraint(
                             skip_resource_constraints_action, "SkipResourceConstraints"
                         )
+
+            # STS Section
+            if cfg.get("sts"):
+                logger.debug(
+                    f"STS section detected. Building assume role policy statement"
+                )
+                self.add_sts_actions(cfg['sts'])
+
         elif cfg.get("mode") == "actions":
             check_actions_schema(cfg)
             if "actions" in cfg.keys():

policy_sentry/writing/template.py

@@ -42,6 +42,14 @@
 # Exclude actions from the output by specifying them here. Accepts wildcards, like kms:Delete*
 exclude-actions:
 - ''
+# If this policy needs to include an AssumeRole action
+sts:
+  assume-role:
+    - ''
+  assume-role-with-saml:
+    - ''
+  assume-role-with-web-identity:
+    - ''
 """
 
 CRUD_TEMPLATE_DICT = {
@@ -61,7 +69,12 @@
         "service-permissions-management": [],
     },
     "skip-resource-constraints": [],
-    "exclude-actions": []
+    "exclude-actions": [],
+    "sts": {
+        "assume-role": [], 
+        "assume-role-with-saml": [], 
+        "assume-role-with-web-identity": []
+    }
 }
 
 ACTIONS_TEMPLATE_DICT = {"mode": "actions", "name": "", "actions": []}

policy_sentry/writing/validate.py

@@ -1,8 +1,9 @@
 """
 Validation for the Policy Sentry YML Templates.
 """
+from ast import Or
 import logging
-from schema import Optional, Schema, And, Use, SchemaError
+from schema import Optional, Schema, And, Use, Regex, SchemaError
 
 logger = logging.getLogger(__name__)
 
@@ -51,6 +52,7 @@ def check(conf_schema, conf):
         },
         Optional("skip-resource-constraints"): [str],
         Optional("exclude-actions"): [str],
+        Optional("sts"): dict({And(Use(str), Regex(r'^assume-role(-with-)*(saml|web-identity)*$')): [str]}),
     }
 )
 

test/writing/test_sid_group_crud.py

@@ -608,3 +608,34 @@ def test_exclude_actions_empty_sid_from_crud_output(self):
             ]
         }
         self.assertDictEqual(result, expected_result)
+
+
+    def test_write_template_with_sts_actions(self):
+        cfg = {
+            "mode": "crud",
+            "name": "RoleNameWithCRUD",
+            "sts": {"assume-role-with-web-identity": ["arn:aws:iam::123456789012:role/demo"]},
+            "list": [
+                "arn:aws:secretsmanager:us-east-1:123456789012:secret:anothersecret"
+            ],
+        }
+        sid_group = SidGroup()
+        rendered_policy = sid_group.process_template(cfg)
+        desired_output = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Sid": "AssumeRoleWithWebIdentity",
+                    "Effect": "Allow",
+                    "Action": [
+                        "sts:AssumeRoleWithWebIdentity"
+                    ],
+                    "Resource": [
+                        "arn:aws:iam::123456789012:role/demo"
+                    ]
+                }
+            ]
+        }
+
+        # print(json.dumps(rendered_policy, indent=4))
+        self.assertEqual(rendered_policy, desired_output)

test/writing/test_template.py

@@ -47,6 +47,14 @@ def test_crud_template(self):
 # Exclude actions from the output by specifying them here. Accepts wildcards, like kms:Delete*
 exclude-actions:
 - ''
+# If this policy needs to include an AssumeRole action
+sts:
+  assume-role:
+    - ''
+  assume-role-with-saml:
+    - ''
+  assume-role-with-web-identity:
+    - ''
 """
         crud_template = create_crud_template()
         self.maxDiff = None

test/writing/test_validate.py

@@ -42,6 +42,18 @@
     "wrist": ["arn:aws:s3:::example-org-sbx-vmimport",],
 }
 
+valid_crud_with_sts_actions = {
+    "mode": "crud",
+    "name": "RoleNameWithCRUD",
+    "sts": {"assume-role": ["arn:aws:iam::123456789012:role/demo",]}
+}
+
+invalid_crud_with_invalid_sts_action = {
+    "mode": "crud",
+    "name": "RoleNameWithCRUD",
+    "sts": {"assume-role-fake-something": ["arn:aws:iam::123456789012:role/demo",]}
+}
+
 
 class YMLSchemaTestCase(unittest.TestCase):
     def test_actions_schema(self):
@@ -54,5 +66,8 @@ def test_crud_schema(self):
         result = check_crud_schema(valid_cfg_for_crud)
         self.assertTrue(result)
         self.assertTrue(check_crud_schema(valid_crud_with_one_item_only))
+        self.assertTrue(check_crud_schema(valid_crud_with_sts_actions))
         with self.assertRaises(Exception):
             check_crud_schema(invalid_crud_with_mispelled_category)
+        with self.assertRaises(Exception):
+            check_crud_schema(invalid_crud_with_invalid_sts_action)

test/writing/test_write_policy_library_usage.py

@@ -86,6 +86,16 @@
             "Resource": [
                 "arn:aws:kms:us-east-1:123456789012:key/123456"
             ]
+        },
+        {
+            "Sid": "AssumeRole",
+            "Effect": "Allow",
+            "Action": [
+                "sts:AssumeRole"
+            ],
+            "Resource": [
+                "arn:aws:iam::123456789012:role/demo"
+            ]
         }
     ]
 }
@@ -181,6 +191,7 @@ def test_write_crud_policy_with_library_only(self):
             "arn:aws:ssm:us-east-1:123456789012:parameter/test"
         )
         another_crud_template["wildcard-only"]["single-actions"].extend(wildcard_actions_to_add)
+        another_crud_template["sts"]["assume-role"].append("arn:aws:iam::123456789012:role/demo")
         # Modify it
         sid_group = SidGroup()
         # minimize = None
@@ -193,7 +204,8 @@ def test_write_crud_policy_with_library_only(self):
             "SecretsmanagerWriteSecret",
             "S3ListObject",
             "SsmTaggingParameter",
-            "KmsPermissionsmanagementKey"
+            "KmsPermissionsmanagementKey",
+            "AssumeRole"
         ]
         self.maxDiff = None
         print(json.dumps(result, indent=4))