Allow customizing which actions are being guessed #22

tests/iam/action_test.py

@@ -42,7 +42,6 @@ def test_base_action(test_input, expected):
         Action('autoscaling', 'DescribeAutoScalingGroups'),
     ]),
     (Action('ec2', 'DetachVolume'), [
-        # Is this the best result? Aren't we usually just interested in Attach-Detach?
         Action('ec2', 'CreateVolume'),
         Action('ec2', 'DeleteVolume'),
         Action('ec2', 'AttachVolume'),
@@ -58,10 +57,18 @@ def test_base_action(test_input, expected):
         Action('s3', 'ListObjects'),
     ]),
 ])
-def test_find_matching_actions(test_input, expected):
-    assert test_input.matching_actions() == expected
+def test_find_matching_actions_without_filtering(test_input, expected):
+    assert test_input.matching_actions(allowed_prefixes=None) == expected
 
 
+@pytest.mark.parametrize("test_input,expected,allowed_prefixes", [
+    (Action('ec2', 'DetachVolume'), [
+        Action('ec2', 'AttachVolume'),
+    ], ['Attach']),
+])
+def test_find_matching_actions_with_filtering(test_input, expected, allowed_prefixes):
+    assert test_input.matching_actions(allowed_prefixes=allowed_prefixes) == expected
+
     # TODO:
     # * Encrypt/Decrypt/GenerateDataKey?
     # * ecr:BatchDeleteImage, ecr:BatchGetImage

tests/integration/cli_guess_test.py

@@ -61,3 +61,44 @@ def test_should_guess_all_matching_statements():
     result = runner.invoke(cli.root_group, args=["guess"], input=input_policy.to_json())
     assert result.exit_code == 0
     assert parse_policy_document(result.output) == expected_output
+
+
+def test_should_guess_only_specific_actions_and_fix_upper_lowercase():
+    input_policy = PolicyDocument(
+        Version="2012-10-17",
+        Statement=[
+            Statement(
+                Effect="Allow",
+                Action=[
+                    Action('ec2', 'DetachVolume'),
+                ],
+                Resource=["*"]
+            ),
+        ]
+    )
+
+    expected_output = PolicyDocument(
+        Version="2012-10-17",
+        Statement=[
+            Statement(
+                Effect="Allow",
+                Action=[
+                    Action('ec2', 'DetachVolume'),
+                ],
+                Resource=["*"]
+            ),
+            Statement(
+                Effect="Allow",
+                Action=[
+                    Action('ec2', 'AttachVolume'),
+                    Action('ec2', 'DescribeVolumes'),
+                ],
+                Resource=["*"]
+            ),
+        ]
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli.root_group, args=["guess", "--only", "Attach", "--only", "describe"], input=input_policy.to_json())
+    assert result.exit_code == 0
+    assert parse_policy_document(result.output) == expected_output

trailscraper/cli.py

@@ -107,11 +107,16 @@ def generate():
 
 
 @click.command("guess")
-def guess():
+@click.option("--only", multiple=True,
+              help='Only guess actions with the given prefix, e.g. Describe (can be passed multiple times)')
+def guess(only):
     """Extend a policy passed in through STDIN by guessing related actions"""
     stdin = click.get_text_stream('stdin')
     policy = parse_policy_document(stdin)
-    policy = guess_statements(policy)
+
+    allowed_prefixes = [s.title() for s in only]
+
+    policy = guess_statements(policy, allowed_prefixes)
     click.echo(policy.to_json())
 
 

trailscraper/guess.py

@@ -2,13 +2,13 @@
 from trailscraper.iam import PolicyDocument, Statement
 
 
-def _guess_actions(actions):
+def _guess_actions(actions, allowed_prefixes):
     return [item for action in actions
-            for item in action.matching_actions()]
+            for item in action.matching_actions(allowed_prefixes)]
 
 
-def _extend_statement(statement):
-    extended_actions = _guess_actions(statement.Action)
+def _extend_statement(statement, allowed_prefixes):
+    extended_actions = _guess_actions(statement.Action, allowed_prefixes)
     if extended_actions:
         return [statement, Statement(Action=extended_actions,
                                      Effect=statement.Effect,
@@ -17,9 +17,9 @@ def _extend_statement(statement):
     return [statement]
 
 
-def guess_statements(policy):
+def guess_statements(policy, allowed_prefixes):
     """Guess additional create actions"""
     extended_statements = [item for statement in policy.Statement
-                           for item in _extend_statement(statement)]
+                           for item in _extend_statement(statement, allowed_prefixes)]
 
     return PolicyDocument(Version=policy.Version, Statement=extended_statements)

trailscraper/iam.py

@@ -9,6 +9,8 @@
 from toolz.curried import groupby as groupbyz
 from toolz.curried import map as mapz
 
+BASE_ACTION_PREFIXES = ["Describe", "Create", "Delete", "Update", "Detach", "Attach", "List", "Put", "Get", ]
+
 
 class BaseElement(object):
     """Base Class for all IAM Policy classes"""
@@ -36,18 +38,6 @@ def __repr__(self):
 class Action(BaseElement):
     """Action in an IAM Policy."""
 
-    BASE_ACTION_PREFIXES = [
-        "Describe",
-        "Create",
-        "Delete",
-        "Update",
-        "Detach",
-        "Attach",
-        "List",
-        "Put",
-        "Get",
-    ]
-
     def __init__(self, prefix, action):
         self.action = action
         self.prefix = prefix
@@ -57,20 +47,24 @@ def json_repr(self):
 
     def _base_action(self):
         without_prefix = self.action
-        for prefix in self.BASE_ACTION_PREFIXES:
+        for prefix in BASE_ACTION_PREFIXES:
             without_prefix = re.sub(prefix, "", without_prefix)
 
         without_plural = re.sub(r"s$", "", without_prefix)
 
         return without_plural
 
-    def matching_actions(self):
+    def matching_actions(self, allowed_prefixes):
         """Return a matching create action for this Action"""
-        potential_matches = [Action(prefix=self.prefix, action=action_prefix+ self._base_action())
-                             for action_prefix in self.BASE_ACTION_PREFIXES]
 
-        potential_matches.append(Action(prefix=self.prefix, action="Describe" + self._base_action()+"s"))
-        potential_matches.append(Action(prefix=self.prefix, action="List" + self._base_action()+"s"))
+        if not allowed_prefixes:
+            allowed_prefixes = BASE_ACTION_PREFIXES
+
+        potential_matches = [Action(prefix=self.prefix, action=action_prefix + self._base_action())
+                             for action_prefix in allowed_prefixes]
+
+        potential_matches += [Action(prefix=self.prefix, action=action_prefix + self._base_action() + "s")
+                              for action_prefix in allowed_prefixes]
 
         return [potential_match
                 for potential_match in potential_matches