diff --git a/core/utils_aws.py b/core/utils_aws.py index be68639..cb151a4 100644 --- a/core/utils_aws.py +++ b/core/utils_aws.py @@ -16,32 +16,86 @@ AWS_RETRY_CONFIG = Config(retries={"mode": "adaptive", "max_attempts": 8}) +# Upper bound on the items collected for a single resource type. Some list +# operations (e.g. describe_images, describe_snapshots) return every publicly +# shared resource in the region when called without an owner filter, which would +# otherwise stall the scan and blow up the raw-data file. +MAX_ITEMS_PER_RESOURCE_TYPE = 50_000 + + +def extract_result_path(container: Any, result_path: list) -> list: + if isinstance(result_path, str): + result_path = [result_path] + value = container + for key in result_path: + if not isinstance(value, dict): + return [] + value = value.get(key) + return value if isinstance(value, list) else [] + def paginate( client: Any, operation_name: str, - result_key: str, + result_path: list, + max_items: int = MAX_ITEMS_PER_RESOURCE_TYPE, **kwargs: Any, ) -> list: items: list = [] for page in client.get_paginator(operation_name).paginate(**kwargs): - items.extend(page.get(result_key, [])) + items.extend(extract_result_path(page, result_path)) + if len(items) >= max_items: + return items[:max_items] return items def paginate_or_call( client: Any, operation_name: str, - result_key: str, + result_path: list, + max_items: int = MAX_ITEMS_PER_RESOURCE_TYPE, **kwargs: Any, ) -> list: """paginate() when boto3 supports it for this operation, else a single call.""" if client.can_paginate(operation_name): - return paginate(client, operation_name, result_key, **kwargs) + return paginate(client, operation_name, result_path, max_items, **kwargs) response = getattr(client, operation_name)(**kwargs) - if not isinstance(response, dict): - return [] - return response.get(result_key, []) + return extract_result_path(response, result_path)[:max_items] + + +def parse_resource_type_params(raw_params: Any) -> dict: + if isinstance(raw_params, dict): + return raw_params + try: + parsed = json.loads(raw_params or "{}") + except (TypeError, ValueError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def resolve_aws_call_spec(resource_type_code: str, params: dict) -> dict | None: + service = params.get("service") + operation = params.get("operation") + if service and operation: + return { + "source": "params", + "service": service, + "operation": operation, + "result_path": list(params.get("result_path") or []), + "kwargs": dict(params.get("kwargs") or {}), + } + + parts = resource_type_code.split(".") + if len(parts) >= 4 and parts[0] == "AWS": + return { + "source": "code", + "service": parts[1], + "operation": parts[2], + "result_path": [part.strip() for part in parts[3:]], + "kwargs": {}, + } + + return None def convert_datetime(obj: Any) -> Any: @@ -78,7 +132,11 @@ def build_aws_resource_inventory( # Load the ResourceType mapping resource_type_mapping = { - item["code"]: {"id": item["id"], "name": item["name"]} + item["code"]: { + "id": item["id"], + "name": item["name"], + "params": item.get("params"), + } for item in load_data("resourcetype") if item["csp"] == 2 and item["status"] == "t" } @@ -89,19 +147,49 @@ def build_aws_resource_inventory( # Aggregate resources by type and location aggregated_resources = defaultdict(int) + # How each catalogue row was resolved, summarised into run.log once the + # sweep is done. Without it a scan result cannot be traced back to the + # master data that produced it. + spec_sources = defaultdict(int) + # Iterate through each resource type in the JSON for idx, (resource_type_code, resource_info) in enumerate( resource_type_mapping.items(), start=1 ): - parts = resource_type_code.split(".") - if len(parts) != 4 or parts[0] != "AWS": - # logger.warning(f"Invalid resource type format: {resource_type_code}. Skipping.") + params = parse_resource_type_params(resource_info.get("params")) + spec = resolve_aws_call_spec(resource_type_code, params) + if spec is None: + if len(resource_type_code.split(".")) == 2: + # Service-level placeholder (e.g. AWS.iam) with nothing to + # call yet -- intentional, so keep it off the console. + spec_sources["placeholder"] += 1 + logger.debug( + "No call spec for placeholder resource type %s. Skipping.", + resource_type_code, + ) + else: + spec_sources["invalid"] += 1 + logger.warning( + "Invalid resource type format: %s. Skipping.", + resource_type_code, + ) continue - # Extract service name, operation name, and result key - service_name, operation_name, result_key = parts[1], parts[2], parts[3] + service_name = spec["service"] + operation_name = spec["operation"] + result_path = spec["result_path"] + call_kwargs = spec["kwargs"] - # logger.info(f"Processing service {service_name} with operation {operation_name}") + spec_sources[spec["source"]] += 1 + logger.debug( + "Resolved %s from %s -> %s.%s result_path=%s kwargs=%s", + resource_type_code, + spec["source"], + service_name, + operation_name, + result_path, + call_kwargs, + ) try: client = session.client( @@ -111,7 +199,20 @@ def build_aws_resource_inventory( # logger.error(f"Operation {operation_name} does not exist for service {service_name}") continue - resources = paginate_or_call(client, operation_name, result_key.strip()) + resources = paginate_or_call( + client, + operation_name, + result_path, + MAX_ITEMS_PER_RESOURCE_TYPE, + **call_kwargs, + ) + if len(resources) >= MAX_ITEMS_PER_RESOURCE_TYPE: + logger.warning( + "Item cap of %d reached for %s.%s; results truncated.", + MAX_ITEMS_PER_RESOURCE_TYPE, + service_name, + operation_name, + ) # Aggregate the resources for resource in resources: @@ -141,6 +242,17 @@ def build_aws_resource_inventory( ) continue + logger.info( + "Resource type resolution: %d of %d rows scanned (%d from params, " + "%d from code), %d placeholders skipped, %d invalid.", + spec_sources["params"] + spec_sources["code"], + len(resource_type_mapping), + spec_sources["params"], + spec_sources["code"], + spec_sources["placeholder"], + spec_sources["invalid"], + ) + # Save raw data to a JSON file raw_data = convert_datetime(raw_data) diff --git a/tests/test_utils_aws.py b/tests/test_utils_aws.py index 8ccf9de..7aa3112 100644 --- a/tests/test_utils_aws.py +++ b/tests/test_utils_aws.py @@ -10,9 +10,12 @@ from core.utils_aws import ( convert_datetime, + extract_result_path, get_missing_months_aws, paginate, paginate_or_call, + parse_resource_type_params, + resolve_aws_call_spec, ) @@ -270,6 +273,369 @@ def test_failed_service_is_skipped_and_logged_at_debug( self.assertFalse(any(r.levelno >= logging.WARNING for r in cm.records)) +class ParseResourceTypeParamsTests(unittest.TestCase): + def test_parses_json_string(self): + self.assertEqual( + parse_resource_type_params('{"service": "s3"}'), {"service": "s3"} + ) + + def test_empty_object_string(self): + self.assertEqual(parse_resource_type_params("{}"), {}) + + def test_none_and_empty_string_degrade_to_empty_dict(self): + self.assertEqual(parse_resource_type_params(None), {}) + self.assertEqual(parse_resource_type_params(""), {}) + + def test_malformed_json_degrades_to_empty_dict(self): + self.assertEqual(parse_resource_type_params("{not json"), {}) + + def test_non_object_json_degrades_to_empty_dict(self): + self.assertEqual(parse_resource_type_params("[1, 2]"), {}) + + def test_dict_passes_through(self): + self.assertEqual( + parse_resource_type_params({"service": "ec2"}), {"service": "ec2"} + ) + + +class ResolveAwsCallSpecTests(unittest.TestCase): + def test_four_part_code_with_no_params(self): + spec = resolve_aws_call_spec("AWS.s3.list_buckets.Buckets", {}) + self.assertEqual( + spec, + { + "source": "code", + "service": "s3", + "operation": "list_buckets", + "result_path": ["Buckets"], + "kwargs": {}, + }, + ) + + def test_five_part_code_with_no_params_keeps_full_result_path(self): + spec = resolve_aws_call_spec( + "AWS.cloudfront.list_distributions.DistributionList.Items", {} + ) + self.assertEqual(spec["result_path"], ["DistributionList", "Items"]) + self.assertEqual(spec["service"], "cloudfront") + + def test_params_take_precedence_over_code(self): + spec = resolve_aws_call_spec( + "AWS.wrong.wrong_op.Wrong", + { + "service": "cloudfront", + "operation": "list_distributions", + "result_path": ["DistributionList", "Items"], + }, + ) + self.assertEqual( + spec, + { + "source": "params", + "service": "cloudfront", + "operation": "list_distributions", + "result_path": ["DistributionList", "Items"], + "kwargs": {}, + }, + ) + + def test_params_with_kwargs(self): + spec = resolve_aws_call_spec( + "AWS.ec2.describe_snapshots.Snapshots", + { + "service": "ec2", + "operation": "describe_snapshots", + "result_path": ["Snapshots"], + "kwargs": {"OwnerIds": ["self"]}, + }, + ) + self.assertEqual(spec["kwargs"], {"OwnerIds": ["self"]}) + + def test_unknown_params_keys_are_ignored(self): + spec = resolve_aws_call_spec( + "AWS.s3.list_buckets.Buckets", + { + "service": "s3", + "operation": "list_buckets", + "result_path": ["Buckets"], + "future_key": "whatever", + }, + ) + self.assertNotIn("future_key", spec) + self.assertEqual(spec["service"], "s3") + + def test_params_missing_result_path_defaults_to_empty(self): + spec = resolve_aws_call_spec( + "AWS.iam", {"service": "iam", "operation": "list_users"} + ) + self.assertEqual(spec["result_path"], []) + + def test_other_cloud_params_fall_back_to_code(self): + spec = resolve_aws_call_spec( + "AWS.s3.list_buckets.Buckets", {"kind": "functionapp"} + ) + self.assertEqual(spec["service"], "s3") + + def test_partial_params_fall_back_to_code(self): + spec = resolve_aws_call_spec("AWS.s3.list_buckets.Buckets", {"service": "s3"}) + self.assertEqual(spec["operation"], "list_buckets") + + def test_source_records_which_branch_resolved_the_row(self): + from_code = resolve_aws_call_spec("AWS.s3.list_buckets.Buckets", {}) + from_params = resolve_aws_call_spec( + "AWS.s3.list_buckets.Buckets", + {"service": "s3", "operation": "list_buckets"}, + ) + self.assertEqual(from_code["source"], "code") + self.assertEqual(from_params["source"], "params") + + def test_two_part_placeholder_is_unresolvable(self): + self.assertIsNone(resolve_aws_call_spec("AWS.iam", {})) + + def test_malformed_code_is_unresolvable(self): + self.assertIsNone(resolve_aws_call_spec("AWS.ec2.describe_instances", {})) + self.assertIsNone(resolve_aws_call_spec("Azure.a.b.c", {})) + self.assertIsNone(resolve_aws_call_spec("", {})) + + +class ExtractResultPathTests(unittest.TestCase): + def test_single_key(self): + self.assertEqual(extract_result_path({"Items": [1, 2]}, ["Items"]), [1, 2]) + + def test_nested_key(self): + self.assertEqual( + extract_result_path({"A": {"B": ["x"]}}, ["A", "B"]), + ["x"], + ) + + def test_missing_intermediate_key_returns_empty(self): + self.assertEqual(extract_result_path({"A": {}}, ["A", "B"]), []) + self.assertEqual(extract_result_path({}, ["A", "B"]), []) + + def test_non_dict_intermediate_returns_empty(self): + self.assertEqual(extract_result_path({"A": "scalar"}, ["A", "B"]), []) + + def test_non_list_final_value_returns_empty(self): + self.assertEqual(extract_result_path({"A": {"B": 5}}, ["A", "B"]), []) + self.assertEqual(extract_result_path({"A": {"B": {}}}, ["A", "B"]), []) + + def test_empty_path_returns_empty(self): + self.assertEqual(extract_result_path({"Items": [1]}, []), []) + + def test_non_dict_container_returns_empty(self): + self.assertEqual(extract_result_path(None, ["Items"]), []) + + +class BuildAwsResourceInventorySpecTests(unittest.TestCase): + """The scanner drives the resolved spec, not the raw code string.""" + + def _row(self, id_, code, name, params="{}"): + return { + "code": code, + "id": id_, + "name": name, + "csp": 2, + "status": "t", + "params": params, + } + + def _run(self, rows, poc_return=None): + """Run the scanner with load_data/boto3/paginate_or_call mocked out. + + Returns (paginate_or_call mock, captured log records). A plain handler + is used instead of assertLogs because a clean scan logs nothing at all. + """ + records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record): + records.append(record) + + aws_logger = logging.getLogger("core.engine.aws") + handler = _Capture() + previous_level = aws_logger.level + aws_logger.addHandler(handler) + aws_logger.setLevel(logging.DEBUG) + + try: + with tempfile.TemporaryDirectory() as tmp: + report_path = os.path.join(tmp, "report") + raw_data_path = os.path.join(tmp, "raw") + os.makedirs(os.path.join(report_path, "data"), exist_ok=True) + os.makedirs(raw_data_path, exist_ok=True) + + with ( + patch("core.utils_aws.load_data", return_value=rows), + patch("core.utils_aws.boto3.Session"), + patch("core.utils_aws.connect") as mock_connect, + patch("core.utils_aws.paginate_or_call") as mock_poc, + ): + mock_connect.return_value.__enter__.return_value = MagicMock() + mock_poc.return_value = [] if poc_return is None else poc_return + + from core.utils_aws import build_aws_resource_inventory + + build_aws_resource_inventory( + 2, + {"accessKey": "AK", "secretKey": "SK", "region": "us-east-1"}, + report_path, + raw_data_path, + ) + return mock_poc, records + finally: + aws_logger.removeHandler(handler) + aws_logger.setLevel(previous_level) + + def test_params_drive_the_call_for_nested_result_path(self): + rows = [ + self._row( + 1, + "AWS.cloudfront.list_distributions.DistributionList.Items", + "CloudFront", + '{"service": "cloudfront", "operation": "list_distributions", ' + '"result_path": ["DistributionList", "Items"]}', + ) + ] + mock_poc, records = self._run(rows, poc_return=[{"Id": "E1"}]) + + args, kwargs = mock_poc.call_args + self.assertEqual(args[1], "list_distributions") + self.assertEqual(args[2], ["DistributionList", "Items"]) + self.assertEqual(kwargs, {}) + self.assertFalse(any(r.levelno >= logging.WARNING for r in records)) + + def test_params_kwargs_are_forwarded_to_the_call(self): + rows = [ + self._row( + 1, + "AWS.ec2.describe_snapshots.Snapshots", + "EBS Snapshot", + '{"service": "ec2", "operation": "describe_snapshots", ' + '"result_path": ["Snapshots"], "kwargs": {"OwnerIds": ["self"]}}', + ) + ] + mock_poc, _ = self._run(rows) + + _, kwargs = mock_poc.call_args + self.assertEqual(kwargs, {"OwnerIds": ["self"]}) + + def test_code_is_used_when_params_are_empty(self): + rows = [self._row(1, "AWS.s3.list_buckets.Buckets", "S3")] + mock_poc, _ = self._run(rows) + + args, _ = mock_poc.call_args + self.assertEqual(args[1], "list_buckets") + self.assertEqual(args[2], ["Buckets"]) + + def test_two_part_placeholder_is_skipped_at_debug(self): + rows = [self._row(1, "AWS.iam", "IAM")] + mock_poc, records = self._run(rows) + + mock_poc.assert_not_called() + self.assertFalse(any(r.levelno >= logging.WARNING for r in records)) + self.assertTrue( + any( + "AWS.iam" in r.getMessage() + for r in records + if r.levelno == logging.DEBUG + ) + ) + + def test_malformed_code_is_skipped_with_warning(self): + rows = [self._row(1, "AWS.ec2.describe_instances", "Broken")] + mock_poc, records = self._run(rows) + + mock_poc.assert_not_called() + self.assertTrue( + any( + r.levelno == logging.WARNING + and "AWS.ec2.describe_instances" in r.getMessage() + for r in records + ) + ) + + def test_item_cap_hit_logs_warning(self): + rows = [self._row(1, "AWS.ec2.describe_images.Images", "AMI")] + with patch("core.utils_aws.MAX_ITEMS_PER_RESOURCE_TYPE", 3): + _, records = self._run(rows, poc_return=[{"ImageId": "ami"}] * 3) + + self.assertTrue( + any( + r.levelno == logging.WARNING + and "describe_images" in r.getMessage() + and "ec2" in r.getMessage() + for r in records + ) + ) + + def test_per_row_debug_line_names_the_spec_source(self): + rows = [ + self._row( + 1, + "AWS.cloudfront.list_distributions.DistributionList.Items", + "CloudFront", + '{"service": "cloudfront", "operation": "list_distributions", ' + '"result_path": ["DistributionList", "Items"]}', + ), + self._row(2, "AWS.s3.list_buckets.Buckets", "S3"), + ] + _, records = self._run(rows) + messages = [r.getMessage() for r in records if r.levelno == logging.DEBUG] + + self.assertTrue( + any("AWS.cloudfront" in m and "from params" in m for m in messages), + messages, + ) + self.assertTrue( + any("AWS.s3.list_buckets" in m and "from code" in m for m in messages), + messages, + ) + + def test_summary_line_counts_every_resolution_outcome(self): + rows = [ + self._row( + 1, + "AWS.cloudfront.list_distributions.DistributionList.Items", + "CloudFront", + '{"service": "cloudfront", "operation": "list_distributions", ' + '"result_path": ["DistributionList", "Items"]}', + ), + self._row(2, "AWS.s3.list_buckets.Buckets", "S3"), + self._row(3, "AWS.iam", "IAM"), + self._row(4, "AWS.ec2.describe_instances", "Broken"), + ] + _, records = self._run(rows) + summary = [ + r.getMessage() + for r in records + if r.levelno == logging.INFO + and "Resource type resolution" in r.getMessage() + ] + + self.assertEqual(len(summary), 1, records) + self.assertEqual( + summary[0], + "Resource type resolution: 2 of 4 rows scanned (1 from params, " + "1 from code), 1 placeholders skipped, 1 invalid.", + ) + + def test_summary_is_info_so_it_stays_off_the_default_console(self): + rows = [self._row(1, "AWS.s3.list_buckets.Buckets", "S3")] + _, records = self._run(rows) + + summary = [r for r in records if "Resource type resolution" in r.getMessage()] + self.assertEqual([r.levelno for r in summary], [logging.INFO]) + + def test_item_cap_is_passed_to_paginate_or_call(self): + from core.utils_aws import MAX_ITEMS_PER_RESOURCE_TYPE + + rows = [self._row(1, "AWS.s3.list_buckets.Buckets", "S3")] + mock_poc, _ = self._run(rows) + + args, _ = mock_poc.call_args + self.assertEqual(args[3], MAX_ITEMS_PER_RESOURCE_TYPE) + + class PaginateTests(unittest.TestCase): def _fake_client(self, pages): """Build a stub client whose paginator yields the given pages.""" @@ -283,15 +649,49 @@ def test_collects_items_across_pages(self): client = self._fake_client( [{"Items": [1, 2, 3]}, {"Items": [4, 5]}, {"Items": [6]}] ) - self.assertEqual(paginate(client, "any_op", "Items"), [1, 2, 3, 4, 5, 6]) + self.assertEqual(paginate(client, "any_op", ["Items"]), [1, 2, 3, 4, 5, 6]) def test_missing_result_key_treated_as_empty(self): client = self._fake_client([{"Items": [1]}, {}]) - self.assertEqual(paginate(client, "any_op", "Items"), [1]) + self.assertEqual(paginate(client, "any_op", ["Items"]), [1]) + + def test_walks_nested_result_path_per_page(self): + client = self._fake_client( + [ + {"DistributionList": {"Items": ["d1", "d2"]}}, + {"DistributionList": {"Items": ["d3"]}}, + ] + ) + self.assertEqual( + paginate(client, "list_distributions", ["DistributionList", "Items"]), + ["d1", "d2", "d3"], + ) + + def test_page_missing_intermediate_key_yields_nothing_for_that_page(self): + client = self._fake_client( + [{"DistributionList": {"Items": ["d1"]}}, {}, {"DistributionList": {}}] + ) + self.assertEqual( + paginate(client, "list_distributions", ["DistributionList", "Items"]), + ["d1"], + ) + + def test_bare_string_result_path_is_treated_as_single_key(self): + client = self._fake_client([{"Volumes": ["v1", "v2"]}]) + self.assertEqual(paginate(client, "describe_volumes", "Volumes"), ["v1", "v2"]) + + def test_stops_and_truncates_at_max_items(self): + pages = [{"Items": list(range(4))} for _ in range(10)] + client = self._fake_client(pages) + + result = paginate(client, "any_op", ["Items"], 10) + + self.assertEqual(len(result), 10) + self.assertEqual(result, list(range(4)) * 2 + [0, 1]) def test_forwards_kwargs_to_paginator(self): client = self._fake_client([{"Items": []}]) - paginate(client, "any_op", "Items", MaxResults=50) + paginate(client, "any_op", ["Items"], MaxResults=50) client.get_paginator.return_value.paginate.assert_called_once_with( MaxResults=50 ) @@ -322,7 +722,7 @@ def test_uses_paginator_when_available(self): client.can_paginate.return_value = True client.get_paginator.return_value = paginator - self.assertEqual(paginate_or_call(client, "list_things", "Items"), [1, 2, 3]) + self.assertEqual(paginate_or_call(client, "list_things", ["Items"]), [1, 2, 3]) client.can_paginate.assert_called_once_with("list_things") def test_falls_back_to_single_call_when_not_paginable(self): @@ -330,7 +730,7 @@ def test_falls_back_to_single_call_when_not_paginable(self): client.can_paginate.return_value = False client.list_things.return_value = {"Items": ["a", "b"]} - self.assertEqual(paginate_or_call(client, "list_things", "Items"), ["a", "b"]) + self.assertEqual(paginate_or_call(client, "list_things", ["Items"]), ["a", "b"]) client.list_things.assert_called_once_with() def test_non_dict_response_returns_empty_list(self): @@ -338,7 +738,38 @@ def test_non_dict_response_returns_empty_list(self): client.can_paginate.return_value = False client.list_things.return_value = None - self.assertEqual(paginate_or_call(client, "list_things", "Items"), []) + self.assertEqual(paginate_or_call(client, "list_things", ["Items"]), []) + + def test_walks_nested_result_path_on_single_call(self): + client = MagicMock() + client.can_paginate.return_value = False + client.get_apps.return_value = { + "ApplicationsResponse": {"Applications": ["a1", "a2"]} + } + + self.assertEqual( + paginate_or_call( + client, "get_apps", ["ApplicationsResponse", "Applications"] + ), + ["a1", "a2"], + ) + + def test_forwards_kwargs_to_single_call(self): + client = MagicMock() + client.can_paginate.return_value = False + client.describe_images.return_value = {"Images": []} + + paginate_or_call(client, "describe_images", ["Images"], Owners=["self"]) + client.describe_images.assert_called_once_with(Owners=["self"]) + + def test_truncates_single_call_at_max_items(self): + client = MagicMock() + client.can_paginate.return_value = False + client.list_things.return_value = {"Items": list(range(20))} + + self.assertEqual( + paginate_or_call(client, "list_things", ["Items"], 5), [0, 1, 2, 3, 4] + ) if __name__ == "__main__":