performance improvements

plugins/aws/fix_plugin_aws/access_edges.py

@@ -1,5 +1,5 @@
 from functools import lru_cache
-from attr import define, frozen
+from attr import frozen, define
 from fix_plugin_aws.resource.base import AwsResource, GraphBuilder
 
 from typing import List, Literal, Set, Optional, Tuple, Union, Pattern
@@ -14,23 +14,24 @@
 )
 from fix_plugin_aws.resource.iam import AwsIamGroup, AwsIamPolicy, AwsIamUser
 from fixlib.baseresources import EdgeType, PolicySourceKind
+from fixlib.json import to_json, to_json_str
 from fixlib.types import Json
 
 from cloudsplaining.scan.policy_document import PolicyDocument
 from cloudsplaining.scan.statement_detail import StatementDetail
-from policy_sentry.querying.actions import get_action_data
+from policy_sentry.querying.actions import get_action_data, get_actions_matching_arn
 from policy_sentry.querying.all import get_all_actions
-from policy_sentry.util.arns import ARN
+from policy_sentry.util.arns import ARN, get_service_from_arn
 import re
 import logging
 
+log = logging.getLogger("fix.plugins.aws")
 
-log = logging.getLogger(__name__)
 
 ALL_ACTIONS = get_all_actions()
 
 
-@define
+@define(slots=True)
 class IamRequestContext:
     principal: AwsResource
     identity_policies: List[Tuple[PolicySource, PolicyDocument]]
@@ -54,29 +55,39 @@ def all_policies(
 IamAction = str
 
 
-@lru_cache(maxsize=1024)
-def find_allowed_action(policy_document: PolicyDocument) -> Set[IamAction]:
+def find_allowed_action(policy_document: PolicyDocument, service_prefix: str) -> Set[IamAction]:
     allowed_actions: Set[IamAction] = set()
     for statement in policy_document.statements:
         if statement.effect_allow:
-            allowed_actions.update(get_expanded_action(statement))
+            allowed_actions.update(get_expanded_action(statement, service_prefix))
 
     return allowed_actions
 
 
-def find_all_allowed_actions(all_involved_policies: List[PolicyDocument]) -> Set[IamAction]:
-    allowed_actions: Set[IamAction] = set()
+def find_all_allowed_actions(all_involved_policies: List[PolicyDocument], resource_arn: str) -> Set[IamAction]:
+    resource_actions = set()
+    try:
+        resource_actions = set(get_actions_matching_arn(resource_arn))
+    except Exception as e:
+        log.warning(f"Error when trying to get actions matching ARN {resource_arn}: {e}")
+
+    service_prefix = ""
+    try:
+        service_prefix = get_service_from_arn(resource_arn)
+    except Exception as e:
+        log.warning(f"Error when trying to get service prefix from ARN {resource_arn}: {e}")
+    policy_actions: Set[IamAction] = set()
     for p in all_involved_policies:
-        allowed_actions.update(find_allowed_action(p))
-    return allowed_actions
+        policy_actions.update(find_allowed_action(p, service_prefix))
+    return policy_actions.intersection(resource_actions)
 
 
-@lru_cache(maxsize=1024)
-def get_expanded_action(statement: StatementDetail) -> Set[str]:
+def get_expanded_action(statement: StatementDetail, service_prefix: str) -> Set[str]:
     actions = set()
     expanded: List[str] = statement.expanded_actions or []
     for action in expanded:
-        actions.add(action)
+        if action.startswith(f"{service_prefix}:"):
+            actions.add(action)
 
     return actions
 
@@ -366,15 +377,15 @@ def check_resource_based_policies(
                 scopes.append(
                     PermissionScope(
                         source=source,
-                        constraints=constraints,
-                        conditions=PermissionCondition(allow=[statement.condition]),
+                        constraints=tuple(constraints),
+                        conditions=PermissionCondition(allow=(to_json_str(statement.condition),)),
                     )
                 )
             else:
                 scopes.append(
                     PermissionScope(
                         source=source,
-                        constraints=constraints,
+                        constraints=tuple(constraints),
                     )
                 )
 
@@ -400,12 +411,11 @@ def check_identity_based_policies(
         for statement, resource_constraints in collect_matching_statements(
             policy=policy, effect="Allow", action=action, resource=resource, principal=None
         ):
-            conditions = []
+            conditions = None
             if statement.condition:
-                conditions.append(statement.condition)
-            scopes.append(
-                PermissionScope(source, resource_constraints, conditions=PermissionCondition(allow=conditions))
-            )
+                conditions = PermissionCondition(allow=(to_json_str(statement.condition),))
+
+            scopes.append(PermissionScope(source, tuple(resource_constraints), conditions=conditions))
 
     return scopes
 
@@ -498,11 +508,11 @@ def check_policies(
         )
         if isinstance(resource_result, FinalAllow):
             scopes = resource_result.scopes
-            final_resource_scopes: List[PermissionScope] = []
+            final_resource_scopes: Set[PermissionScope] = set()
             for scope in scopes:
-                final_resource_scopes.append(scope.with_deny_conditions(deny_conditions))
+                final_resource_scopes.add(scope.with_deny_conditions(deny_conditions))
 
-            return AccessPermission(action=action, level=get_action_level(action), scopes=final_resource_scopes)
+            return AccessPermission(action=action, level=get_action_level(action), scopes=tuple(final_resource_scopes))
         if isinstance(resource_result, Continue):
             scopes = resource_result.scopes
             allowed_scopes.extend(scopes)
@@ -535,17 +545,23 @@ def check_policies(
     # 7. if we reached here, the action is allowed
     level = get_action_level(action)
 
-    final_scopes: List[PermissionScope] = []
+    final_scopes: Set[PermissionScope] = set()
     for scope in allowed_scopes:
         if deny_conditions:
             scope = scope.with_deny_conditions(deny_conditions)
-        final_scopes.append(scope)
+        final_scopes.add(scope)
+
+    # if there is a scope with no conditions, we can ignore everything else
+    for scope in final_scopes:
+        if scope.has_no_condititons():
+            final_scopes = {scope}
+            break
 
     # return the result
     return AccessPermission(
         action=action,
         level=level,
-        scopes=final_scopes,
+        scopes=tuple(final_scopes),
     )
 
 
@@ -555,8 +571,9 @@ def compute_permissions(
     resource_based_policies: List[Tuple[PolicySource, PolicyDocument]],
 ) -> List[AccessPermission]:
 
+    assert resource.arn
     # step 1: find the relevant action to check
-    relevant_actions = find_all_allowed_actions(iam_context.all_policies(resource_based_policies))
+    relevant_actions = find_all_allowed_actions(iam_context.all_policies(resource_based_policies), resource.arn)
 
     all_permissions: List[AccessPermission] = []
 
@@ -610,7 +627,7 @@ def _get_identity_based_policies(self, principal: AwsResource) -> List[Tuple[Pol
                     if doc := to_node.policy_document_json():
                         attached_policies.append(
                             (
-                                PolicySource(kind=PolicySourceKind.Principal, uri=principal.arn or ""),
+                                PolicySource(kind=PolicySourceKind.Principal, uri=to_node.arn or ""),
                                 PolicyDocument(doc),
                             )
                         )
@@ -632,7 +649,7 @@ def _get_identity_based_policies(self, principal: AwsResource) -> List[Tuple[Pol
                             if doc := group_successor.policy_document_json():
                                 group_policies.append(
                                     (
-                                        PolicySource(kind=PolicySourceKind.Group, uri=group.arn or ""),
+                                        PolicySource(kind=PolicySourceKind.Group, uri=group_successor.arn or ""),
                                         PolicyDocument(doc),
                                     )
                                 )
@@ -653,6 +670,12 @@ def add_access_edges(self) -> None:
 
                 permissions = compute_permissions(node, context, resource_policies)
 
+                if not permissions:
+                    continue
+
+                json_permissions = to_json(permissions)
+                reported = {"permissions": json_permissions}
+
                 self.builder.add_edge(
-                    from_node=context.principal, edge_type=EdgeType.access, permissions=permissions, to_node=node
+                    from_node=context.principal, edge_type=EdgeType.access, reported=reported, to_node=node
                 )

plugins/aws/fix_plugin_aws/access_edges_utils.py

@@ -1,11 +1,13 @@
 from abc import ABC
 from attrs import frozen, evolve
 from typing import List, Optional, Tuple, Any
+from fixlib.json import to_json_str
 from fixlib.types import Json
 from fixlib.baseresources import PolicySourceKind
 
 ResourceConstraint = str
 
+ConditionString = str
 
 @frozen
 class PolicySource:
@@ -22,30 +24,39 @@ def resource_policy(self, builder: Any) -> List[Tuple[PolicySource, Json]]:
 @frozen
 class PermissionCondition:
     # if nonempty and any evals to true, access is granted, otherwise implicitly denied
-    allow: Optional[List[Json]] = None
+    allow: Optional[Tuple[ConditionString, ...]] = None
     # if nonempty and any is evals to false, access is implicitly denied
-    boundary: Optional[List[Json]] = None
+    boundary: Optional[Tuple[ConditionString, ...]] = None
     # if nonempty and any evals to true, access is explicitly denied
-    deny: Optional[List[Json]] = None
+    deny: Optional[Tuple[ConditionString, ...]] = None
 
 
 @frozen
 class PermissionScope:
     source: PolicySource
-    constraints: List[ResourceConstraint]  # aka resource constraints
+    constraints: Tuple[ResourceConstraint, ...]  # aka resource constraints
     conditions: Optional[PermissionCondition] = None
 
     def with_deny_conditions(self, deny_conditions: List[Json]) -> "PermissionScope":
         c = self.conditions or PermissionCondition()
-        return evolve(self, conditions=evolve(c, deny=deny_conditions))
+        return evolve(self, conditions=evolve(c, deny=tuple([to_json_str(c) for c in deny_conditions])))
 
     def with_boundary_conditions(self, boundary_conditions: List[Json]) -> "PermissionScope":
         c = self.conditions or PermissionCondition()
-        return evolve(self, conditions=evolve(c, boundary=boundary_conditions))
+        return evolve(self, conditions=evolve(c, boundary=tuple([to_json_str(c) for c in boundary_conditions])))
+
+    def has_no_condititons(self) -> bool:
+        if self.conditions is None:
+            return True
+
+        if self.conditions.allow is None and self.conditions.boundary is None and self.conditions.deny is None:
+            return True
+
+        return False
 
 
 @frozen
 class AccessPermission:
     action: str
     level: str
-    scopes: List[PermissionScope]
+    scopes: Tuple[PermissionScope, ...]

plugins/aws/fix_plugin_aws/collector.py

@@ -256,6 +256,7 @@ def get_last_run() -> Optional[datetime]:
                     raise Exception("Only AWS resources expected")
 
             # add access edges
+            log.info(f"[Aws:{self.account.id}] Create access edges.")
             access_edge_creator = AccessEdgeCreator(global_builder)
             access_edge_creator.add_access_edges()
 

plugins/aws/test/__init__.py

@@ -35,7 +35,7 @@ def builder(aws_client: AwsClient, no_feedback: CoreFeedback) -> Iterator[GraphB
         yield GraphBuilder(
             Graph(),
             Cloud(id="aws"),
-            AwsAccount(id="test"),
+            AwsAccount(id="test", arn="arn:aws:organizations::123456789012:account/o-exampleorgid/123456789012"),
             regions[0],
             {r.id: r for r in regions},
             aws_client,
@@ -51,5 +51,5 @@ def no_feedback() -> CoreFeedback:
 
 @fixture
 def account_collector(aws_config: AwsConfig, no_feedback: CoreFeedback) -> AwsAccountCollector:
-    account = AwsAccount(id="test")
+    account = AwsAccount(id="test", arn="arn:aws:organizations::123456789012:account/o-exampleorgid/123456789012")
     return AwsAccountCollector(aws_config, Cloud(id="aws"), account, ["us-east-1"], no_feedback, {})

plugins/aws/test/acccess_edges_test.py

@@ -16,8 +16,9 @@
     compute_permissions,
 )
 
-from fix_plugin_aws.access_edges_utils import PermissionCondition, PolicySource, PermissionScope
+from fix_plugin_aws.access_edges_utils import PolicySource
 from fixlib.baseresources import PolicySourceKind
+from fixlib.json import to_json_str
 
 
 def test_find_allowed_action() -> None:
@@ -26,11 +27,12 @@ def test_find_allowed_action() -> None:
         "Statement": [
             {"Effect": "Allow", "Action": ["s3:GetObject", "s3:PutObject"], "Resource": ["arn:aws:s3:::bucket/*"]},
             {"Effect": "Allow", "Action": ["s3:ListBuckets"], "Resource": ["*"]},
+            {"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": ["*"]},
             {"Effect": "Deny", "Action": ["s3:DeleteObject"], "Resource": ["arn:aws:s3:::bucket/*"]},
         ],
     }
 
-    allowed_actions = find_allowed_action(PolicyDocument(policy_document))
+    allowed_actions = find_allowed_action(PolicyDocument(policy_document), "s3")
 
     assert allowed_actions == {"s3:GetObject", "s3:PutObject", "s3:ListBuckets"}
 
@@ -394,7 +396,7 @@ def test_compute_permissions_user_inline_policy_allow() -> None:
     s = permissions[0].scopes[0]
     assert s.source.kind == PolicySourceKind.Principal
     assert s.source.uri == user.arn
-    assert s.constraints == ["arn:aws:s3:::my-test-bucket"]
+    assert s.constraints == ("arn:aws:s3:::my-test-bucket",)
 
 
 def test_compute_permissions_user_inline_policy_allow_with_conditions() -> None:
@@ -430,11 +432,12 @@ def test_compute_permissions_user_inline_policy_allow_with_conditions() -> None:
     assert permissions[0].action == "s3:ListBucket"
     assert permissions[0].level == "List"
     assert len(permissions[0].scopes) == 1
-    assert permissions[0].scopes[0] == PermissionScope(
-        PolicySource(kind=PolicySourceKind.Principal, uri=user.arn),
-        ["arn:aws:s3:::my-test-bucket"],
-        conditions=PermissionCondition(allow=[condition]),
-    )
+    s = permissions[0].scopes[0]
+    assert s.source.kind == PolicySourceKind.Principal
+    assert s.source.uri == user.arn
+    assert s.constraints == ("arn:aws:s3:::my-test-bucket",)
+    assert s.conditions
+    assert s.conditions.allow == (to_json_str(condition),)
 
 
 def test_compute_permissions_user_inline_policy_deny() -> None:
@@ -647,9 +650,9 @@ def test_deny_overrides_allow_with_condition() -> None:
     s = p.scopes[0]
     assert s.source.kind == PolicySourceKind.Principal
     assert s.source.uri == user.arn
-    assert s.constraints == ["arn:aws:s3:::my-test-bucket"]
+    assert s.constraints == ("arn:aws:s3:::my-test-bucket",)
     assert s.conditions
-    assert s.conditions.deny == [condition]
+    assert s.conditions.deny == (to_json_str(condition),)
 
 
 def test_compute_permissions_resource_based_policy_allow() -> None:
@@ -690,7 +693,7 @@ def test_compute_permissions_resource_based_policy_allow() -> None:
     s = p.scopes[0]
     assert s.source.kind == PolicySourceKind.Resource
     assert s.source.uri == bucket.arn
-    assert s.constraints == ["arn:aws:s3:::my-test-bucket"]
+    assert s.constraints == ("arn:aws:s3:::my-test-bucket",)
 
 
 def test_compute_permissions_permission_boundary_restrict() -> None:
@@ -747,7 +750,7 @@ def test_compute_permissions_permission_boundary_restrict() -> None:
     s = p.scopes[0]
     assert s.source.kind == PolicySourceKind.Principal
     assert s.source.uri == user.arn
-    assert s.constraints == ["arn:aws:s3:::my-test-bucket"]
+    assert s.constraints == ("arn:aws:s3:::my-test-bucket",)
 
 
 def test_compute_permissions_scp_deny() -> None:
@@ -826,7 +829,7 @@ def test_compute_permissions_user_with_group_policies() -> None:
     s = p.scopes[0]
     assert s.source.kind == PolicySourceKind.Group
     assert s.source.uri == group.arn
-    assert s.constraints == [bucket.arn]
+    assert s.constraints == (bucket.arn,)
 
 
 def test_compute_permissions_implicit_deny() -> None: