[aws][feat] Reimplement SSM Compliance resource collection (#2280)

Author: 1101-1

plugins/aws/fix_plugin_aws/resource/s3.py

@@ -1,3 +1,4 @@
+from functools import partial
 import logging
 from collections import defaultdict
 from datetime import timedelta
@@ -363,9 +364,22 @@ def _get_tags(self, client: AwsClient) -> Dict[str, str]:
         return tags_as_dict(tag_list)  # type: ignore
 
     def collect_usage_metrics(self, builder: GraphBuilder) -> List[AwsCloudwatchQuery]:
+        def _calculate_total_size(bucket_instance: AwsS3Bucket) -> None:
+            # Calculate the total bucket size for each bucket by summing up the sizes of all storage types
+            bucket_size: Dict[str, float] = defaultdict(float)
+            for metric_name, metric_values in bucket_instance._resource_usage.items():
+                if metric_name.endswith("_bucket_size_bytes"):
+                    for name, value in metric_values.items():
+                        bucket_size[name] += value
+            if bucket_size:
+                bucket_instance._resource_usage["bucket_size_bytes"] = dict(bucket_size)
+
         # Filter out metrics with the 'aws-controltower' dimension value
         if "aws-controltower" in self.safe_name:
             return []
+
+        # calculate all bucket sizes after usage metrics collection
+        builder.after_collect_actions.append(partial(_calculate_total_size, self))
         storage_types = {
             "StandardStorage": "standard_storage",
             "IntelligentTieringStorage": "intelligent_tiering_storage",
@@ -415,16 +429,6 @@ def collect_usage_metrics(self, builder: GraphBuilder) -> List[AwsCloudwatchQuer
                 )
         return queries
 
-    def complete_graph(self, builder: GraphBuilder, source: Json) -> None:
-        # Calculate the total bucket size for each bucket by summing up the sizes of all storage types
-        bucket_size: Dict[str, float] = defaultdict(float)
-        for metric_name, metric_values in self._resource_usage.items():
-            if metric_name.endswith("_bucket_size_bytes"):
-                for name, value in metric_values.items():
-                    bucket_size[name] += value
-        if bucket_size:
-            self._resource_usage["bucket_size_bytes"] = dict(bucket_size)
-
     def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:
         tags = self._get_tags(client)
         tags[key] = value

plugins/aws/fix_plugin_aws/resource/ssm.py

@@ -1,4 +1,5 @@
-import json
+from functools import partial
+from json import loads as json_loads
 import logging
 from datetime import datetime
 from typing import ClassVar, Dict, Optional, List, Type, Any
@@ -12,9 +13,10 @@
 from fix_plugin_aws.resource.ec2 import AwsEc2Instance
 from fix_plugin_aws.resource.s3 import AwsS3Bucket
 from fix_plugin_aws.utils import ToDict
-from fixlib.baseresources import ModelReference
-from fixlib.json_bender import Bender, S, Bend, AsDateString, ForallBend, K
+from fixlib.baseresources import SEVERITY_MAPPING, Finding, ModelReference, PhantomBaseResource, Severity
+from fixlib.json_bender import Bender, S, Bend, AsDateString, ForallBend
 from fixlib.types import Json
+from fixlib.utils import chunks
 
 log = logging.getLogger("fix.plugins.aws")
 service_name = "ssm"
@@ -239,7 +241,7 @@ def collect_document(name: str) -> None:
                     and (instance := cls.from_api(js, builder))
                 ):
                     if content_format == "JSON":
-                        instance.content = json.loads(content)
+                        instance.content = json_loads(content)
                     elif content_format == "YAML":
                         instance.content = yaml.safe_load(content)
                     else:
@@ -341,7 +343,7 @@ class AwsSSMNonCompliantSummary:
     severity_summary: Optional[AwsSSMSeveritySummary] = field(default=None, metadata={"description": "A summary of the non-compliance severity by compliance type"})  # fmt: skip
 
 
-ResourceTypeLookup = {
+ResourceTypeLookup: Dict[str, Type[AwsResource]] = {
     "ManagedInstance": AwsEc2Instance,
     "AWS::EC2::Instance": AwsEc2Instance,
     "AWS::DynamoDB::Table": AwsDynamoDbTable,
@@ -351,45 +353,88 @@ class AwsSSMNonCompliantSummary:
 
 
 @define(eq=False, slots=False)
-class AwsSSMResourceCompliance(AwsResource):
+class AwsSSMResourceCompliance(AwsResource, PhantomBaseResource):
     kind: ClassVar[str] = "aws_ssm_resource_compliance"
-    _kind_display: ClassVar[str] = "AWS SSM Resource Compliance"
-    _kind_description: ClassVar[str] = "AWS SSM Resource Compliance is a feature within AWS Systems Manager that evaluates and reports on the compliance status of AWS resources. It checks resources against predefined or custom rules, identifying non-compliant configurations and security issues. Users can view compliance data, generate reports, and take corrective actions to maintain resource adherence to organizational standards and best practices."  # fmt: skip
-    _docs_url: ClassVar[str] = (
-        "https://docs.aws.amazon.com/systems-manager/latest/userguide/sysman-compliance-about.html"
-    )
-    _kind_service: ClassVar[Optional[str]] = service_name
-    _metadata: ClassVar[Dict[str, Any]] = {"icon": "resource", "group": "management"}
-    _aws_metadata: ClassVar[Dict[str, Any]] = {"arn_tpl": "arn:{partition}:ssm:{region}:{account}:resource-compliance/{id}"}  # fmt: skip
+    _model_export: ClassVar[bool] = False  # do not export this class, since there will be no instances of it
     api_spec: ClassVar[AwsApiSpec] = AwsApiSpec(
-        "ssm", "list-resource-compliance-summaries", "ResourceComplianceSummaryItems"
+        "ssm",
+        "list-resource-compliance-summaries",
+        "ResourceComplianceSummaryItems",
+        {"Filters": [{"Key": "Status", "Values": ["COMPLIANT"], "Type": "EQUAL"}]},
     )
-    _reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["aws_ec2_instance", "aws_dynamodb_table", "aws_s3_bucket", "aws_ssm_document"]}
-    }
     mapping: ClassVar[Dict[str, Bender]] = {
-        "id": S("ComplianceType") + K("_") + S("ResourceType") + K("_") + S("ResourceId"),
+        "id": S("Id"),
+        "name": S("title"),
         "compliance_type": S("ComplianceType"),
         "resource_type": S("ResourceType"),
         "resource_id": S("ResourceId"),
+        "title": S("Title"),
         "status": S("Status"),
-        "overall_severity": S("OverallSeverity"),
+        "severity": S("Severity"),
         "execution_summary": S("ExecutionSummary") >> Bend(AwsSSMComplianceExecutionSummary.mapping),
-        "compliant_summary": S("CompliantSummary") >> Bend(AwsSSMCompliantSummary.mapping),
-        "non_compliant_summary": S("NonCompliantSummary") >> Bend(AwsSSMNonCompliantSummary.mapping),
+        "compliance_details": S("Details"),
     }
-    compliance_type: Optional[str] = field(default=None, metadata={"description": "The compliance type."})  # fmt: skip
-    resource_type: Optional[str] = field(default=None, metadata={"description": "The resource type."})  # fmt: skip
-    resource_id: Optional[str] = field(default=None, metadata={"description": "The resource ID."})  # fmt: skip
-    status: Optional[str] = field(default=None, metadata={"description": "The compliance status for the resource."})  # fmt: skip
-    overall_severity: Optional[str] = field(default=None, metadata={"description": "The highest severity item found for the resource. The resource is compliant for this item."})  # fmt: skip
-    execution_summary: Optional[AwsSSMComplianceExecutionSummary] = field(default=None, metadata={"ignore_history": True, "description": "Information about the execution."})  # fmt: skip
-    compliant_summary: Optional[AwsSSMCompliantSummary] = field(default=None, metadata={"description": "A list of items that are compliant for the resource."})  # fmt: skip
-    non_compliant_summary: Optional[AwsSSMNonCompliantSummary] = field(default=None, metadata={"description": "A list of items that aren't compliant for the resource."})  # fmt: skip
+    compliance_type: Optional[str] = field(default=None, metadata={"description": "The compliance type. For example, Association (for a State Manager association), Patch, or Custom:string are all valid compliance types."})  # fmt: skip
+    resource_type: Optional[str] = field(default=None, metadata={"description": "The type of resource. ManagedInstance is currently the only supported resource type."})  # fmt: skip
+    resource_id: Optional[str] = field(default=None, metadata={"description": "An ID for the resource. For a managed node, this is the node ID."})  # fmt: skip
+    title: Optional[str] = field(default=None, metadata={"description": "A title for the compliance item. For example, if the compliance item is a Windows patch, the title could be the title of the KB article for the patch; for example: Security Update for Active Directory Federation Services."})  # fmt: skip
+    status: Optional[str] = field(default=None, metadata={"description": "The status of the compliance item. An item is either COMPLIANT, NON_COMPLIANT, or an empty string (for Windows patches that aren't applicable)."})  # fmt: skip
+    severity: Optional[str] = field(default=None, metadata={"description": "The severity of the compliance status. Severity can be one of the following: Critical, High, Medium, Low, Informational, Unspecified."})  # fmt: skip
+    execution_summary: Optional[AwsSSMComplianceExecutionSummary] = field(default=None, metadata={"description": "A summary for the compliance item. The summary includes an execution ID, the execution type (for example, command), and the execution time."})  # fmt: skip
+    compliance_details: Optional[Dict[str, str]] = field(default=None, metadata={"description": "A Key:Value tag combination for the compliance item."})  # fmt: skip
+
+    def parse_finding(self) -> Finding:
+        title = self.title or ""
+        severity = SEVERITY_MAPPING.get(self.severity or "", Severity.medium)
+        details = self.compliance_details
+        if self.execution_summary:
+            updated_at = self.execution_summary.execution_time
+        else:
+            updated_at = None
+        return Finding(title, severity, None, None, updated_at, details)
 
-    def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        if (rt := self.resource_type) and (rid := self.resource_id) and (clazz := ResourceTypeLookup.get(rt)):
-            builder.add_edge(self, clazz=clazz, id=rid)
+    @classmethod
+    def collect(cls, json: List[Json], builder: GraphBuilder) -> None:
+        def add_finding(
+            provider: str, finding: Finding, clazz: Optional[Type[AwsResource]] = None, **node: Any
+        ) -> None:
+            if resource := builder.node(clazz=clazz, **node):
+                resource.add_finding(provider, finding)
+
+        def collect_compliance_items(jsons: List[Json]) -> None:
+            spec = AwsApiSpec("ssm", "list-compliance-items", "ComplianceItems")
+            compliance_ids = [item["ResourceId"] for item in jsons]
+            for result in builder.client.list(
+                aws_service="ssm",
+                action=spec.api_action,
+                result_name=spec.result_property,
+                expected_errors=spec.expected_errors,
+                ResourceIds=compliance_ids,
+            ):
+                if finding := AwsSSMResourceCompliance.from_api(result, builder):
+                    if (
+                        (rt := finding.resource_type)
+                        and (rid := finding.resource_id)
+                        and (clazz := ResourceTypeLookup.get(rt))
+                    ):
+                        # append the finding when all resources have been collected
+                        builder.after_collect_actions.append(
+                            partial(
+                                add_finding,
+                                "amazon_ssm_compliance",
+                                finding.parse_finding(),
+                                clazz,
+                                id=rid,
+                            )
+                        )
+
+        # we can request only 40 items per request
+        for jsons in chunks(json, 39):
+            builder.submit_work("ssm", collect_compliance_items, jsons)
+
+    @classmethod
+    def called_collect_apis(cls) -> List[AwsApiSpec]:
+        return [cls.api_spec, AwsApiSpec("ssm", "list-compliance-items", "ComplianceItems")]
 
 
 resources: List[Type[AwsResource]] = [AwsSSMInstance, AwsSSMDocument, AwsSSMResourceCompliance]

plugins/aws/test/collector_test.py

@@ -37,8 +37,8 @@ def count_kind(clazz: Type[AwsResource]) -> int:
     # make sure all threads have been joined
     assert len(threading.enumerate()) == 1
     # ensure the correct number of nodes and edges
-    assert count_kind(AwsResource) == 260
-    assert len(account_collector.graph.edges) == 574
+    assert count_kind(AwsResource) == 257
+    assert len(account_collector.graph.edges) == 571
     assert len(account_collector.graph.deferred_edges) == 2
     for node in account_collector.graph.nodes:
         if isinstance(node, AwsRegion):

plugins/aws/test/resources/files/ssm/list-compliance-items__foo_foo_foo.json

@@ -0,0 +1,23 @@
+{
+    "ComplianceItems": [
+        {
+            "ComplianceType": "Association",
+            "ResourceType": "ManagedInstance",
+            "ResourceId": "i-1",
+            "Id": "SSM-Association-1",
+            "Title": "State Manager Association Compliance",
+            "Status": "NON_COMPLIANT",
+            "Severity": "HIGH",
+            "ExecutionSummary": {
+                "ExecutionTime": "2024-10-01T12:34:56Z",
+                "ExecutionId": "xyz5678-execution-id",
+                "ExecutionType": "Association"
+            },
+            "Details": {
+                "LastExecutionStatus": "Failed",
+                "ErrorDetails": "Failed to apply association"
+            }
+        }
+    ],
+    "NextToken": "next-token-value"
+}

plugins/aws/test/resources/files/ssm/list-resource-compliance-summaries.json renamed to plugins/aws/test/resources/files/ssm/list-resource-compliance-summaries__Status_COMPLIANT_EQUAL.json

@@ -104,4 +104,4 @@
         }
     ],
     "NextToken": "foo"
-}
+}

plugins/aws/test/resources/ssm_test.py

@@ -1,3 +1,4 @@
+from fix_plugin_aws.resource.ec2 import AwsEc2Instance
 from fix_plugin_aws.resource.ssm import (
     AwsSSMInstance,
     AwsSSMDocument,
@@ -6,14 +7,19 @@
 )
 from test.resources import round_trip_for
 
+from fixlib.baseresources import Severity
+
 
 def test_instances() -> None:
     first, builder = round_trip_for(AwsSSMInstance)
     assert len(builder.resources_of(AwsSSMInstance)) == 2
 
 
 def test_resource_compliance() -> None:
-    round_trip_for(AwsSSMResourceCompliance)
+    collected, _ = round_trip_for(AwsEc2Instance, region_name="global", collect_also=[AwsSSMResourceCompliance])
+    asseessments = collected._assessments
+    assert asseessments[0].findings[0].title == "State Manager Association Compliance"
+    assert asseessments[0].findings[0].severity == Severity.high
 
 
 def test_documents() -> None: