regscale-cli 6.21.2.0__py3-none-any.whl → 6.28.2.1__py3-none-any.whl

regscale/integrations/commercial/aws/inventory/resources/integration.py

@@ -1,12 +1,32 @@
 """AWS application integration resource collectors."""
 
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Optional
 
 from ..base import BaseCollector
 
 
 class IntegrationCollector(BaseCollector):
-    """Collector for AWS application integration resources."""
+    """Collector for AWS application integration resources with filtering support."""
+
+    def __init__(
+        self,
+        session: Any,
+        region: str,
+        account_id: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+        enabled_services: Optional[Dict[str, bool]] = None,
+    ):
+        """
+        Initialize integration collector with filtering support.
+
+        :param session: AWS session to use for API calls
+        :param str region: AWS region to collect from
+        :param str account_id: Optional AWS account ID to filter resources
+        :param dict tags: Optional tag filters (AND logic)
+        :param dict enabled_services: Optional dict of service names to boolean flags for enabling/disabling collection
+        """
+        super().__init__(session, region, account_id, tags)
+        self.enabled_services = enabled_services or {}
 
     def get_api_gateways(self) -> Dict[str, List[Dict[str, Any]]]:
         """
@@ -84,9 +104,75 @@ class IntegrationCollector(BaseCollector):
             self._handle_error(e, "API Gateway APIs")
         return apis
 
+    def _should_include_topic(self, sns, topic_arn: str) -> bool:
+        """
+        Check if topic should be included based on account and tag filters.
+
+        :param sns: SNS client
+        :param str topic_arn: Topic ARN
+        :return: True if topic should be included, False otherwise
+        :rtype: bool
+        """
+        if not self._matches_account(topic_arn):
+            return False
+
+        if self.tags:
+            try:
+                tags_response = sns.list_tags_for_resource(ResourceArn=topic_arn)
+                topic_tags = tags_response.get("Tags", [])
+                return self._matches_tags(topic_tags)
+            except Exception:
+                return False
+
+        return True
+
+    def _get_topic_subscriptions(self, sns, topic_arn: str) -> List[Dict[str, Any]]:
+        """
+        Get all subscriptions for a topic.
+
+        :param sns: SNS client
+        :param str topic_arn: Topic ARN
+        :return: List of subscriptions
+        :rtype: List[Dict[str, Any]]
+        """
+        subs = []
+        sub_paginator = sns.get_paginator("list_subscriptions_by_topic")
+        for sub_page in sub_paginator.paginate(TopicArn=topic_arn):
+            subs.extend(sub_page.get("Subscriptions", []))
+        return subs
+
+    def _build_topic_data(self, topic_arn: str, attrs: Dict[str, Any], subs: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Build topic data dictionary.
+
+        :param str topic_arn: Topic ARN
+        :param attrs: Topic attributes
+        :param subs: List of subscriptions
+        :return: Processed topic data
+        :rtype: Dict[str, Any]
+        """
+        return {
+            "Region": self.region,
+            "TopicArn": topic_arn,
+            "Owner": attrs.get("Owner"),
+            "Policy": attrs.get("Policy"),
+            "DisplayName": attrs.get("DisplayName"),
+            "SubscriptionsConfirmed": attrs.get("SubscriptionsConfirmed"),
+            "SubscriptionsPending": attrs.get("SubscriptionsPending"),
+            "SubscriptionsDeleted": attrs.get("SubscriptionsDeleted"),
+            "Subscriptions": [
+                {
+                    "SubscriptionArn": sub.get("SubscriptionArn"),
+                    "Protocol": sub.get("Protocol"),
+                    "Endpoint": sub.get("Endpoint"),
+                }
+                for sub in subs
+            ],
+        }
+
     def get_sns_topics(self) -> List[Dict[str, Any]]:
         """
-        Get information about SNS topics.
+        Get information about SNS topics with filtering.
 
         :return: List of SNS topic information
         :rtype: List[Dict[str, Any]]
@@ -99,43 +185,73 @@ class IntegrationCollector(BaseCollector):
             for page in paginator.paginate():
                 for topic in page.get("Topics", []):
                     try:
-                        # Get topic attributes
-                        attrs = sns.get_topic_attributes(TopicArn=topic["TopicArn"])["Attributes"]
-                        # Get subscriptions
-                        subs = []
-                        sub_paginator = sns.get_paginator("list_subscriptions_by_topic")
-                        for sub_page in sub_paginator.paginate(TopicArn=topic["TopicArn"]):
-                            subs.extend(sub_page.get("Subscriptions", []))
-
-                        topics.append(
-                            {
-                                "Region": self.region,
-                                "TopicArn": topic.get("TopicArn"),
-                                "Owner": attrs.get("Owner"),
-                                "Policy": attrs.get("Policy"),
-                                "DisplayName": attrs.get("DisplayName"),
-                                "SubscriptionsConfirmed": attrs.get("SubscriptionsConfirmed"),
-                                "SubscriptionsPending": attrs.get("SubscriptionsPending"),
-                                "SubscriptionsDeleted": attrs.get("SubscriptionsDeleted"),
-                                "Subscriptions": [
-                                    {
-                                        "SubscriptionArn": sub.get("SubscriptionArn"),
-                                        "Protocol": sub.get("Protocol"),
-                                        "Endpoint": sub.get("Endpoint"),
-                                    }
-                                    for sub in subs
-                                ],
-                            }
-                        )
+                        topic_arn = topic["TopicArn"]
+
+                        if not self._should_include_topic(sns, topic_arn):
+                            continue
+
+                        attrs = sns.get_topic_attributes(TopicArn=topic_arn)["Attributes"]
+                        subs = self._get_topic_subscriptions(sns, topic_arn)
+                        topic_data = self._build_topic_data(topic_arn, attrs, subs)
+                        topics.append(topic_data)
                     except Exception as e:
                         self._handle_error(e, f"SNS topic {topic['TopicArn']}")
         except Exception as e:
             self._handle_error(e, "SNS topics")
         return topics
 
+    def _should_include_queue(self, sqs, queue_arn: str, queue_url: str) -> bool:
+        """
+        Check if queue should be included based on account and tag filters.
+
+        :param sqs: SQS client
+        :param str queue_arn: Queue ARN
+        :param str queue_url: Queue URL
+        :return: True if queue should be included, False otherwise
+        :rtype: bool
+        """
+        if not self._matches_account(queue_arn):
+            return False
+
+        if self.tags:
+            try:
+                tags_response = sqs.list_queue_tags(QueueUrl=queue_url)
+                queue_tags = tags_response.get("Tags", {})
+                return self._matches_tags(queue_tags)
+            except Exception:
+                return False
+
+        return True
+
+    def _build_queue_data(self, queue_url: str, attrs: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Build queue data dictionary.
+
+        :param str queue_url: Queue URL
+        :param attrs: Queue attributes
+        :return: Processed queue data
+        :rtype: Dict[str, Any]
+        """
+        return {
+            "Region": self.region,
+            "QueueUrl": queue_url,
+            "QueueArn": attrs.get("QueueArn"),
+            "ApproximateNumberOfMessages": attrs.get("ApproximateNumberOfMessages"),
+            "ApproximateNumberOfMessagesNotVisible": attrs.get("ApproximateNumberOfMessagesNotVisible"),
+            "ApproximateNumberOfMessagesDelayed": attrs.get("ApproximateNumberOfMessagesDelayed"),
+            "CreatedTimestamp": attrs.get("CreatedTimestamp"),
+            "LastModifiedTimestamp": attrs.get("LastModifiedTimestamp"),
+            "VisibilityTimeout": attrs.get("VisibilityTimeout"),
+            "MaximumMessageSize": attrs.get("MaximumMessageSize"),
+            "MessageRetentionPeriod": attrs.get("MessageRetentionPeriod"),
+            "DelaySeconds": attrs.get("DelaySeconds"),
+            "Policy": attrs.get("Policy"),
+            "RedrivePolicy": attrs.get("RedrivePolicy"),
+        }
+
     def get_sqs_queues(self) -> List[Dict[str, Any]]:
         """
-        Get information about SQS queues.
+        Get information about SQS queues with filtering.
 
         :return: List of SQS queue information
         :rtype: List[Dict[str, Any]]
@@ -148,29 +264,14 @@ class IntegrationCollector(BaseCollector):
             for page in paginator.paginate():
                 for queue_url in page.get("QueueUrls", []):
                     try:
-                        # Get queue attributes
                         attrs = sqs.get_queue_attributes(QueueUrl=queue_url, AttributeNames=["All"])["Attributes"]
+                        queue_arn = attrs.get("QueueArn", "")
 
-                        queues.append(
-                            {
-                                "Region": self.region,
-                                "QueueUrl": queue_url,
-                                "QueueArn": attrs.get("QueueArn"),
-                                "ApproximateNumberOfMessages": attrs.get("ApproximateNumberOfMessages"),
-                                "ApproximateNumberOfMessagesNotVisible": attrs.get(
-                                    "ApproximateNumberOfMessagesNotVisible"
-                                ),
-                                "ApproximateNumberOfMessagesDelayed": attrs.get("ApproximateNumberOfMessagesDelayed"),
-                                "CreatedTimestamp": attrs.get("CreatedTimestamp"),
-                                "LastModifiedTimestamp": attrs.get("LastModifiedTimestamp"),
-                                "VisibilityTimeout": attrs.get("VisibilityTimeout"),
-                                "MaximumMessageSize": attrs.get("MaximumMessageSize"),
-                                "MessageRetentionPeriod": attrs.get("MessageRetentionPeriod"),
-                                "DelaySeconds": attrs.get("DelaySeconds"),
-                                "Policy": attrs.get("Policy"),
-                                "RedrivePolicy": attrs.get("RedrivePolicy"),
-                            }
-                        )
+                        if not self._should_include_queue(sqs, queue_arn, queue_url):
+                            continue
+
+                        queue_data = self._build_queue_data(queue_url, attrs)
+                        queues.append(queue_data)
                     except Exception as e:
                         self._handle_error(e, f"SQS queue {queue_url}")
         except Exception as e:
@@ -224,14 +325,27 @@ class IntegrationCollector(BaseCollector):
 
     def collect(self) -> Dict[str, Any]:
         """
-        Collect all application integration resources.
+        Collect application integration resources based on enabled_services configuration.
 
-        :return: Dictionary containing all application integration resource information
+        :return: Dictionary containing enabled application integration resource information
         :rtype: Dict[str, Any]
         """
-        return {
-            "APIGateway": self.get_api_gateways(),
-            "SNSTopics": self.get_sns_topics(),
-            "SQSQueues": self.get_sqs_queues(),
-            "EventBridgeRules": self.get_eventbridge_rules(),
-        }
+        result = {}
+
+        # API Gateway
+        if self.enabled_services.get("api_gateway", True):
+            result["APIGateway"] = self.get_api_gateways()
+
+        # SNS Topics
+        if self.enabled_services.get("sns", True):
+            result["SNSTopics"] = self.get_sns_topics()
+
+        # SQS Queues
+        if self.enabled_services.get("sqs", True):
+            result["SQSQueues"] = self.get_sqs_queues()
+
+        # EventBridge Rules
+        if self.enabled_services.get("eventbridge", True):
+            result["EventBridgeRules"] = self.get_eventbridge_rules()
+
+        return result