From c93c2ceaf5154ec48d8e58a75d38b3d39b389000 Mon Sep 17 00:00:00 2001 From: Nikita Melkozerov Date: Thu, 28 Nov 2024 14:51:03 +0100 Subject: [PATCH] make the linter happy --- plugins/aws/fix_plugin_aws/access_edges.py | 119 +++++++------- plugins/aws/fix_plugin_aws/collector.py | 10 +- plugins/aws/test/acccess_edges_test.py | 171 ++++++++++----------- 3 files changed, 136 insertions(+), 164 deletions(-) diff --git a/plugins/aws/fix_plugin_aws/access_edges.py b/plugins/aws/fix_plugin_aws/access_edges.py index d348779a3d..32f6482b1e 100644 --- a/plugins/aws/fix_plugin_aws/access_edges.py +++ b/plugins/aws/fix_plugin_aws/access_edges.py @@ -92,9 +92,9 @@ class ActionToCheck: class ArnResourceValueKind(enum.Enum): - Static = 1 # the segment is a fixed value, e.g. "s3", "vpc/vpc-0e9801d129EXAMPLE", - Pattern = 2 # the segment is a pattern, e.g. "my_corporate_bucket/*", - Any = 3 # the segment is missing, e.g. "::" or it is a wildcard, e.g. "*" + Static = 1 # the segment is a fixed value, e.g. "s3", "vpc/vpc-0e9801d129EXAMPLE", + Pattern = 2 # the segment is a pattern, e.g. "my_corporate_bucket/*", + Any = 3 # the segment is missing, e.g. "::" or it is a wildcard, e.g. "*" @staticmethod def from_str(value: str) -> "ArnResourceValueKind": @@ -104,6 +104,7 @@ def from_str(value: str) -> "ArnResourceValueKind": return ArnResourceValueKind.Pattern return ArnResourceValueKind.Static + @frozen(slots=True) class ArnResource: value: str @@ -121,19 +122,16 @@ def matches(self, segment: str) -> bool: case ArnResourceValueKind.Static: _match = segment == self.value - if self.not_resource: _match = not _match - return _match - @frozen(slots=True) class ArnAccountId: value: str - wildcard: bool # if the account is a wildcard, e.g. "*" or "::" + wildcard: bool # if the account is a wildcard, e.g. "*" or "::" principal_arns: Set[str] children: List[ArnResource] @@ -144,7 +142,7 @@ def matches(self, segment: str) -> bool: @frozen(slots=True) class ArnRegion: value: str - wildcard: bool # if the region is a wildcard, e.g. "*" or "::" + wildcard: bool # if the region is a wildcard, e.g. "*" or "::" principal_arns: Set[str] children: List[ArnAccountId] @@ -160,12 +158,12 @@ class ArnService: def matches(self, segment: str) -> bool: return self.value == segment - + @frozen(slots=True) class ArnPartition: value: str - wildcard: bool # for the cases like "Allow": "*" on all resources + wildcard: bool # for the cases like "Allow": "*" on all resources principal_arns: Set[str] children: List[ArnService] @@ -181,7 +179,6 @@ class PrincipalTree: def __init__(self) -> None: self.partitions: List[ArnPartition] = [] - def _add_allow_all_wildcard(self, principal_arn: str) -> None: partition = next((p for p in self.partitions if p.value == "*"), None) if not partition: @@ -195,7 +192,6 @@ def _add_resource(self, resource_constraint: str, principal_arn: str, nr: bool = _add resource will add the principal arn at the resource level """ - try: arn = ARN(resource_constraint) # Find existing or create partition @@ -221,11 +217,15 @@ def _add_resource(self, resource_constraint: str, principal_arn: str, nr: bool = account_wildcard = arn.account == "*" or not arn.account account = next((a for a in region.children if a.value == (arn.account or "*")), None) if not account: - account = ArnAccountId(value=arn.account or "*", wildcard=account_wildcard, principal_arns=set(), children=[]) + account = ArnAccountId( + value=arn.account or "*", wildcard=account_wildcard, principal_arns=set(), children=[] + ) region.children.append(account) # Add resource - resource = next((r for r in account.children if r.value == arn.resource_string and r.not_resource == nr), None) + resource = next( + (r for r in account.children if r.value == arn.resource_string and r.not_resource == nr), None + ) if not resource: if arn.resource_string == "*": resource_kind = ArnResourceValueKind.Any @@ -233,7 +233,9 @@ def _add_resource(self, resource_constraint: str, principal_arn: str, nr: bool = resource_kind = ArnResourceValueKind.Pattern else: resource_kind = ArnResourceValueKind.Static - resource = ArnResource(value=arn.resource_string, principal_arns=set(), kind=resource_kind, not_resource=nr) + resource = ArnResource( + value=arn.resource_string, principal_arns=set(), kind=resource_kind, not_resource=nr + ) account.children.append(resource) resource.principal_arns.add(principal_arn) @@ -242,7 +244,6 @@ def _add_resource(self, resource_constraint: str, principal_arn: str, nr: bool = log.error(f"Error parsing ARN {principal_arn}: {e}") pass - def _add_service(self, service_prefix: str, principal_arn: str) -> None: # Find existing or create partition partition = next((p for p in self.partitions if p.value == "*"), None) @@ -258,11 +259,9 @@ def _add_service(self, service_prefix: str, principal_arn: str) -> None: service.principal_arns.add(principal_arn) - - def add_principal(self, principal_arn: str, policy_documents: List[FixPolicyDocument]) -> None: """ - This method iterates over every policy statement and adds corresponding arns to principal tree. + This method iterates over every policy statement and adds corresponding arns to principal tree. """ for policy_doc in policy_documents: @@ -276,30 +275,30 @@ def add_principal(self, principal_arn: str, policy_documents: List[FixPolicyDocu self._add_resource(resource, principal_arn) for not_resource in statement.not_resource: self._add_resource(not_resource, principal_arn, nr=True) - + if has_wildcard_resource or (not statement.resources and not statement.not_resource): for ap in statement.actions_patterns: if ap.kind == WildcardKind.any: self._add_allow_all_wildcard(principal_arn) self._add_service(ap.service, principal_arn) - def list_principals(self, resource_arn: ARN) -> Set[str]: """ this will be called for every resource and it must be fast """ - principals = set() + principals: Set[str] = set() matching_partitions = [p for p in self.partitions if p.value if p.matches(resource_arn.partition)] if not matching_partitions: return principals - matching_services = [s for p in matching_partitions for s in p.children if s.matches(resource_arn.service_prefix)] + matching_services = [ + s for p in matching_partitions for s in p.children if s.matches(resource_arn.service_prefix) + ] if not matching_services: return principals principals.update([arn for s in matching_services for arn in s.principal_arns]) - matching_regions = [r for s in matching_services for r in s.children if r.matches(resource_arn.region)] if not matching_regions: return principals @@ -310,20 +309,21 @@ def list_principals(self, resource_arn: ARN) -> Set[str]: return principals principals.update([arn for a in matching_account_ids for arn in a.principal_arns]) - matching_resources = [r for a in matching_account_ids for r in a.children if r.matches(resource_arn.resource_string)] + matching_resources = [ + r for a in matching_account_ids for r in a.children if r.matches(resource_arn.resource_string) + ] if not matching_resources: return principals - + principals.update([arn for r in matching_resources for arn in r.principal_arns]) return principals - @frozen(slots=True) class ResourceWildcardPattern: raw_value: str - partition: str | None # None in case the whole string is "*" + partition: str | None # None in case the whole string is "*" service: str region: str region_value_kind: ArnResourceValueKind @@ -332,22 +332,21 @@ class ResourceWildcardPattern: resource: str resource_value_kind: ArnResourceValueKind - @staticmethod def from_str(value: str) -> "ResourceWildcardPattern": if value == "*": return ResourceWildcardPattern( raw_value=value, - partition=None, - service="*", + partition=None, + service="*", region="*", - region_value_kind=ArnResourceValueKind.Any, + region_value_kind=ArnResourceValueKind.Any, account="*", - account_value_kind=ArnResourceValueKind.Any, - resource="*", - resource_value_kind=ArnResourceValueKind.Any + account_value_kind=ArnResourceValueKind.Any, + resource="*", + resource_value_kind=ArnResourceValueKind.Any, ) - + try: splitted = value.split(":", 5) if len(splitted) != 6: @@ -363,12 +362,13 @@ def from_str(value: str) -> "ResourceWildcardPattern": account=account, account_value_kind=ArnResourceValueKind.from_str(account), resource=resource, - resource_value_kind=ArnResourceValueKind.from_str(resource) + resource_value_kind=ArnResourceValueKind.from_str(resource), ) except Exception as e: log.error(f"Error parsing resource pattern {value}: {e}") raise e + @frozen(slots=True) class IamRequestContext: principal: AwsResource @@ -378,7 +378,6 @@ class IamRequestContext: # starting from the root, then all org units, then the account service_control_policy_levels: Tuple[Tuple[FixPolicyDocument, ...], ...] - def all_policies( self, resource_based_policies: Optional[Tuple[Tuple[PolicySource, FixPolicyDocument], ...]] = None ) -> List[FixPolicyDocument]: @@ -392,6 +391,7 @@ def all_policies( IamAction = str + @lru_cache(maxsize=4096) def find_allowed_action(policy_document: FixPolicyDocument, service_prefix: str) -> Set[IamAction]: allowed_actions: Set[IamAction] = set() @@ -592,29 +592,28 @@ def match_pattern(resource_segment: str, wildcard_segment: str, wildcard_segment return resource_segment == wildcard_segment - def expand_arn_wildcards_and_match(identifier: ARN, wildcard_string: ResourceWildcardPattern) -> bool: # if wildard is *, we can shortcut here if wildcard_string.partition is None: return True - + # go through the ARN segments and match them if not wildcard_string.partition == identifier.partition: return False - + if not wildcard_string.service == identifier.service_prefix: return False - + if not match_pattern(identifier.region, wildcard_string.region, wildcard_string.region_value_kind): return False - + if not match_pattern(identifier.account, wildcard_string.account, wildcard_string.account_value_kind): return False - + if not match_pattern(identifier.resource_string, wildcard_string.resource, wildcard_string.resource_value_kind): return False - + return True @@ -966,12 +965,12 @@ def check_identity_policies_closure(resource: ARN) -> List[PermissionScope]: return check_identity_policies_closure + @lru_cache(maxsize=4096) def check_permission_boundaries( request_context: IamRequestContext, action: ActionToCheck ) -> Callable[[ARN], Union[Literal["Denied", "NextStep"], List[Json]]]: - matching_fns = [] # ignore policy sources and resource constraints because permission boundaries @@ -995,7 +994,7 @@ def check_permission_boundaries_closure(resource: ARN) -> Union[Literal["Denied" # no matching permission boundaries that allow access return "Denied" - + return check_permission_boundaries_closure @@ -1185,7 +1184,6 @@ def check_non_resource_policies_closure(resource: ARN) -> Optional[AccessPermiss # comes from the resource based policies and identity based policies allowed_scopes: List[PermissionScope] = [] - # 1. check for explicit deny. If denied, we can abort immediately result = explicit_deny_fn(resource) if result == "Denied": @@ -1197,14 +1195,14 @@ def check_non_resource_policies_closure(resource: ARN) -> Optional[AccessPermiss # satisfying any of the conditions above will deny the action deny_conditions.append(c) - - # 2. check for organization SCPs # todo: move it outside the loop - if len(request_context.service_control_policy_levels) > 0 and not is_service_linked_role(request_context.principal): + if len(request_context.service_control_policy_levels) > 0 and not is_service_linked_role( + request_context.principal + ): org_scp_allowed = scp_allowed(request_context, action, resource) if not org_scp_allowed: return None - + # 3. skip resource based policies because the resource has none # 4. to make it a bit simpler, we check the permission boundaries before checking identity based policies @@ -1252,7 +1250,6 @@ def check_non_resource_policies_closure(resource: ARN) -> Optional[AccessPermiss scopes=tuple(final_scopes), ) - return check_non_resource_policies_closure @@ -1379,7 +1376,6 @@ def _build_principal_tree(self) -> PrincipalTree: return tree - def _compute_actions_for_resource(self) -> Dict[str, set[IamAction]]: actions_for_resource: Dict[str, set[IamAction]] = {} @@ -1490,14 +1486,13 @@ def add_access_edges(self) -> None: assert node.arn resource_arn = ARN(node.arn) - if not isinstance(node, HasResourcePolicy): # here we have identity-based policies only and can prune some principals for arn in self.principal_tree.list_principals(resource_arn): context = self.arn_to_context.get(arn) if not context: raise ValueError(f"Principal {arn} not found in the context") - + permissions = compute_permissions( resource_arn, context, tuple(), self.actions_for_resource.get(node.arn, set()) ) @@ -1509,7 +1504,9 @@ def add_access_edges(self) -> None: for permission in permissions: access[permission.level] = True reported = to_json({"permissions": permissions} | access, strip_nulls=True) - self.builder.add_edge(from_node=context.principal, edge_type=EdgeType.iam, reported=reported, node=node) + self.builder.add_edge( + from_node=context.principal, edge_type=EdgeType.iam, reported=reported, node=node + ) else: # here we have resource-based policies and must check all principals. @@ -1529,13 +1526,13 @@ def add_access_edges(self) -> None: if not permissions: continue - access: Dict[PermissionLevel, bool] = {} + access = {} for permission in permissions: access[permission.level] = True reported = to_json({"permissions": permissions} | access, strip_nulls=True) - self.builder.add_edge(from_node=context.principal, edge_type=EdgeType.iam, reported=reported, node=node) - - + self.builder.add_edge( + from_node=context.principal, edge_type=EdgeType.iam, reported=reported, node=node + ) all_principal_arns = {p.principal.arn for p in self.principals if p.principal.arn} diff --git a/plugins/aws/fix_plugin_aws/collector.py b/plugins/aws/fix_plugin_aws/collector.py index c0fa052482..03a08346c5 100644 --- a/plugins/aws/fix_plugin_aws/collector.py +++ b/plugins/aws/fix_plugin_aws/collector.py @@ -2,7 +2,6 @@ from collections import defaultdict from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone -import os from typing import List, Type, Optional, Union, cast, Any from fix_plugin_aws.access_edges import AccessEdgeCreator @@ -69,7 +68,6 @@ from fixlib.threading import ExecutorQueue, GatherFutures from fixlib.types import Json from .utils import global_region_by_partition -from pyinstrument import Profiler log = logging.getLogger("fix.plugins.aws") @@ -264,18 +262,12 @@ def get_last_run() -> Optional[datetime]: log.warning(f"Unexpected node type {node} in graph") raise Exception("Only AWS resources expected") - access_edge_collection_enabled = True + access_edge_collection_enabled = False if access_edge_collection_enabled and global_builder.config.collect_access_edges: # add access edges - profiler = Profiler() - profiler.start() log.info(f"[Aws:{self.account.id}] Create access edges.") access_edge_creator = AccessEdgeCreator(global_builder) access_edge_creator.add_access_edges() - profiler.stop() - html_output = profiler.output_html() - with open(f"profiler_{self.account.id}.html", "w") as f: - f.write(html_output) # final hook when the graph is complete for node, data in list(self.graph.nodes(data=True)): diff --git a/plugins/aws/test/acccess_edges_test.py b/plugins/aws/test/acccess_edges_test.py index d59c897a2d..cd15c1342d 100644 --- a/plugins/aws/test/acccess_edges_test.py +++ b/plugins/aws/test/acccess_edges_test.py @@ -19,7 +19,7 @@ ActionToCheck, get_actions_matching_arn, PrincipalTree, - ArnResourceValueKind + ArnResourceValueKind, ) from fixlib.baseresources import PolicySourceKind, PolicySource, PermissionLevel @@ -1025,9 +1025,9 @@ def test_principal_tree_add_allow_all_wildcard() -> None: """Test adding wildcard (*) permission to the principal tree.""" tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" - + tree._add_allow_all_wildcard(principal_arn) - + # Verify the wildcard partition exists assert len(tree.partitions) == 1 partition = tree.partitions[0] @@ -1041,32 +1041,32 @@ def test_principal_tree_add_resource() -> None: tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" resource_arn = "arn:aws:s3:::my-bucket/my-object" - + tree._add_resource(resource_arn, principal_arn) - + # Verify the partition structure assert len(tree.partitions) == 1 partition = tree.partitions[0] assert partition.value == "aws" assert not partition.wildcard - + # Verify service level assert len(partition.children) == 1 service = partition.children[0] assert service.value == "s3" - + # Verify region level assert len(service.children) == 1 region = service.children[0] assert region.value == "*" assert region.wildcard - + # Verify account level assert len(region.children) == 1 account = region.children[0] assert account.value == "*" assert account.wildcard - + # Verify resource level assert len(account.children) == 1 resource = account.children[0] @@ -1081,16 +1081,16 @@ def test_principal_tree_add_resource_with_wildcard() -> None: tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" resource_arn = "arn:aws:s3:::my-bucket/*" - + tree._add_resource(resource_arn, principal_arn) - + # Verify the resource level has correct wildcard pattern partition = tree.partitions[0] service = partition.children[0] region = service.children[0] account = region.children[0] resource = account.children[0] - + assert resource.value == "my-bucket/*" assert resource.kind == ArnResourceValueKind.Pattern assert principal_arn in resource.principal_arns @@ -1101,9 +1101,9 @@ def test_principal_tree_add_not_resource() -> None: tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" resource_arn = "arn:aws:s3:::my-bucket/private/*" - + tree._add_resource(resource_arn, principal_arn, nr=True) - + # Verify the NotResource flag is set correctly through the tree partition = tree.partitions[0] service = partition.children[0] @@ -1118,14 +1118,14 @@ def test_principal_tree_add_service() -> None: tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" service_prefix = "s3" - + tree._add_service(service_prefix, principal_arn) - + # Verify service is added under wildcard partition assert len(tree.partitions) == 1 partition = tree.partitions[0] assert partition.value == "*" - + assert len(partition.children) == 1 service = partition.children[0] assert service.value == "s3" @@ -1136,37 +1136,30 @@ def test_principal_tree_add_principal_policy() -> None: """Test adding a principal with policy documents to the principal tree.""" tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" - + policy_json = { "Version": "2012-10-17", "Statement": [ - { - "Effect": "Allow", - "Action": ["s3:GetObject"], - "Resource": "arn:aws:s3:::my-bucket/*" - }, - { - "Effect": "Allow", - "Action": ["s3:ListAllMyBuckets"], - "Resource": "*" - } - ] + {"Effect": "Allow", "Action": ["s3:GetObject"], "Resource": "arn:aws:s3:::my-bucket/*"}, + {"Effect": "Allow", "Action": ["s3:ListAllMyBuckets"], "Resource": "*"}, + ], } - + policy_doc = FixPolicyDocument(policy_json) tree.add_principal(principal_arn, [policy_doc]) - + # Verify both the specific resource and wildcard permissions are added assert any( - p.value == "aws" and - any(s.value == "s3" and - any(r.value == "*" and - any(a.value == "*" and - any(res.value == "my-bucket/*" - for res in a.children) - for a in r.children) - for r in s.children) - for s in p.children) + p.value == "aws" + and any( + s.value == "s3" + and any( + r.value == "*" + and any(a.value == "*" and any(res.value == "my-bucket/*" for res in a.children) for a in r.children) + for r in s.children + ) + for s in p.children + ) for p in tree.partitions ) @@ -1176,33 +1169,29 @@ def test_principal_tree_list_principals() -> None: tree = PrincipalTree() principal1 = "arn:aws:iam::123456789012:user/test-user1" principal2 = "arn:aws:iam::123456789012:user/test-user2" - + # Add different types of permissions - policy_doc1 = FixPolicyDocument({ - "Version": "2012-10-17", - "Statement": [{ - "Effect": "Allow", - "Action": ["s3:GetObject"], - "Resource": "arn:aws:s3:::my-bucket/*" - }] - }) - - policy_doc2 = FixPolicyDocument({ - "Version": "2012-10-17", - "Statement": [{ - "Effect": "Allow", - "Action": ["s3:ListAllMyBuckets"], - "Resource": "*" - }] - }) - + policy_doc1 = FixPolicyDocument( + { + "Version": "2012-10-17", + "Statement": [{"Effect": "Allow", "Action": ["s3:GetObject"], "Resource": "arn:aws:s3:::my-bucket/*"}], + } + ) + + policy_doc2 = FixPolicyDocument( + { + "Version": "2012-10-17", + "Statement": [{"Effect": "Allow", "Action": ["s3:ListAllMyBuckets"], "Resource": "*"}], + } + ) + tree.add_principal(principal1, [policy_doc1]) tree.add_principal(principal2, [policy_doc2]) - + # Test specific resource access resource_arn = ARN("arn:aws:s3:::my-bucket/test.txt") matching_principals = tree.list_principals(resource_arn) - + assert principal1 in matching_principals # Has specific access assert principal2 in matching_principals # Has wildcard access @@ -1211,29 +1200,23 @@ def test_principal_tree_add_multiple_statements() -> None: """Test adding multiple statements for the same principal.""" tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" - - policy_doc = FixPolicyDocument({ - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Action": ["s3:GetObject"], - "Resource": "arn:aws:s3:::bucket1/*" - }, - { - "Effect": "Allow", - "Action": ["s3:PutObject"], - "Resource": "arn:aws:s3:::bucket2/*" - } - ] - }) - + + policy_doc = FixPolicyDocument( + { + "Version": "2012-10-17", + "Statement": [ + {"Effect": "Allow", "Action": ["s3:GetObject"], "Resource": "arn:aws:s3:::bucket1/*"}, + {"Effect": "Allow", "Action": ["s3:PutObject"], "Resource": "arn:aws:s3:::bucket2/*"}, + ], + } + ) + tree.add_principal(principal_arn, [policy_doc]) - + # Test access to both buckets bucket1_arn = ARN("arn:aws:s3:::bucket1/test.txt") bucket2_arn = ARN("arn:aws:s3:::bucket2/test.txt") - + assert principal_arn in tree.list_principals(bucket1_arn) assert principal_arn in tree.list_principals(bucket2_arn) @@ -1242,28 +1225,28 @@ def test_principal_tree_not_resource() -> None: """Test NotResource handling in the principal tree.""" tree = PrincipalTree() principal_arn = "arn:aws:iam::123456789012:user/test-user" - - policy_doc = FixPolicyDocument({ - "Version": "2012-10-17", - "Statement": [{ - "Effect": "Allow", - "Action": ["s3:GetObject"], - "NotResource": ["arn:aws:s3:::private-bucket/*"] - }] - }) - + + policy_doc = FixPolicyDocument( + { + "Version": "2012-10-17", + "Statement": [ + {"Effect": "Allow", "Action": ["s3:GetObject"], "NotResource": ["arn:aws:s3:::private-bucket/*"]} + ], + } + ) + tree.add_principal(principal_arn, [policy_doc]) - + # Test access is denied to private bucket private_arn = ARN("arn:aws:s3:::private-bucket/secret.txt") public_arn = ARN("arn:aws:s3:::public-bucket/public.txt") ec2 = ARN("arn:aws:ec2:us-east-1:123456789012:instance/i-1234567890abcdef0") - + matching_principals = tree.list_principals(private_arn) assert principal_arn not in matching_principals - + matching_principals = tree.list_principals(public_arn) assert principal_arn in matching_principals matching_principals = tree.list_principals(ec2) - assert len(matching_principals) == 0 \ No newline at end of file + assert len(matching_principals) == 0