regscale-cli 6.27.2.0__py3-none-any.whl → 6.28.0.0__py3-none-any.whl

regscale/integrations/commercial/amazon/common.py

@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 """RegScale AWS Integrations"""
 import re
+import logging
 from datetime import datetime, timedelta
 from typing import Any, Optional, Tuple
 
@@ -11,7 +12,7 @@ from dateutil import parser
 
 from regscale.core.app.utils.app_utils import create_logger
 
-logger = create_logger()
+logger = logging.getLogger("regscale")
 
 
 def check_finding_severity(comment: Optional[str]) -> str:
@@ -22,17 +23,17 @@ def check_finding_severity(comment: Optional[str]) -> str:
     :rtype: str
     """
     result = ""
-    match = re.search(r"(?<=Finding Severity: ).*", comment)
-    if match:
-        severity = match.group()
-        result = severity  # Output: "High"
+    if comment:
+        match = re.search(r"(?<=Finding Severity: ).*", comment)
+        if match:
+            result = match.group()
     return result
 
 
-def get_due_date(earliest_date_performed: datetime, days: int) -> datetime:
+def get_due_date(earliest_date_performed: str, days: int) -> datetime:
     """Returns the due date for an issue
 
-    :param datetime earliest_date_performed: Earliest date performed
+    :param str earliest_date_performed: Earliest date performed (string format)
     :param int days: Days to add to the earliest date performed
     :return: Due date
     :rtype: datetime
@@ -42,7 +43,7 @@ def get_due_date(earliest_date_performed: datetime, days: int) -> datetime:
         due_date = datetime.strptime(earliest_date_performed, fmt) + timedelta(days=days)
     except ValueError:
         # Try to determine the date format from a string
-        due_date = parser.parse(earliest_date_performed) + timedelta(days)
+        due_date = parser.parse(earliest_date_performed) + timedelta(days=days)
     return due_date
 
 
@@ -94,6 +95,31 @@ def get_comments(finding: dict) -> str:
         return "No remediation recommendation available"
 
 
+def _fetch_with_error_handling(
+    aws_client: BaseClient, method_name: str, result_key: str, resource_type: str = "items"
+) -> list:
+    """Generic fetch method with error handling for AWS Security Hub
+
+    :param BaseClient aws_client: AWS Security Hub Client
+    :param str method_name: Name of the method to call on the client
+    :param str result_key: Key to extract from the response
+    :param str resource_type: Type of resource being fetched (for logging)
+    :return: List of items from AWS
+    :rtype: list
+    """
+    items = []
+    try:
+        method = getattr(aws_client, method_name)
+        response = method()
+        items = response.get(result_key, [])
+        logger.info("Fetched %d %s from Security Hub", len(items), resource_type)
+    except ClientError as cex:
+        logger.error("Unexpected error when fetching %s from AWS: %s", resource_type, cex)
+    except AttributeError as aex:
+        logger.error("Method %s not found on client: %s", method_name, aex)
+    return items
+
+
 def fetch_aws_findings(aws_client: BaseClient) -> list:
     """Fetch AWS Findings with optimized rate limiting and pagination
 
@@ -101,70 +127,41 @@ def fetch_aws_findings(aws_client: BaseClient) -> list:
     :return: AWS Findings
     :rtype: list
     """
-    findings = []
     try:
         # Use optimized SecurityHubPuller for better performance
         from regscale.integrations.commercial.aws.security_hub import SecurityHubPuller
 
-        # Extract credentials from the client to create SecurityHubPuller
-        session = aws_client._client_config.__dict__.get("_user_provided_options", {})
-        region = session.meta.region_name
+        # Extract region from the client's meta information
+        region = aws_client.meta.region_name
 
-        # Create SecurityHubPuller with same credentials as the client
+        # Create SecurityHubPuller with same region and credentials
         puller = SecurityHubPuller(region_name=region)
-
-        # Use existing client instead of creating new one to maintain credentials
         puller.client = aws_client
 
         # Fetch all findings with optimized pagination and rate limiting
         logger.info("Using optimized SecurityHubPuller for findings retrieval...")
         findings = puller.get_all_findings_with_retries()
 
-        logger.info(f"Successfully fetched {len(findings)} findings with rate limiting")
+        logger.info("Successfully fetched %d findings with rate limiting", len(findings))
+        return findings
 
-    except ImportError:
-        # Fallback to original method if SecurityHubPuller not available
-        logger.warning("SecurityHubPuller not available, falling back to basic client")
-        findings = fallback_fetch_aws_findings(aws_client)
-    except ClientError as cex:
-        logger.error("Unexpected error: %s", cex)
+    except (ImportError, AttributeError) as e:
+        # Fallback to original method if SecurityHubPuller not available or region not accessible
+        logger.warning("SecurityHubPuller not available (%s), falling back to basic client", type(e).__name__)
+        return _fetch_with_error_handling(aws_client, "get_findings", "Findings", "findings")
     except Exception as e:
-        logger.error("Error using SecurityHubPuller, falling back to basic client: %s", e)
-        findings = fallback_fetch_aws_findings(aws_client)
-
-    return findings
-
-
-def fallback_fetch_aws_findings(aws_client: BaseClient) -> list:
-    """Fallback method to fetch AWS Findings without pagination
-
-    :param BaseClient aws_client: AWS Security Hub Client
-    :return: AWS Findings
-    :rtype: list
-    """
-    findings = []
-    try:
-        response = aws_client.get_findings()
-        findings = response.get("Findings", [])
-    except ClientError as cex:
-        create_logger().error("Unexpected error when fetching resources from AWS: %s", cex)
-    return findings
+        logger.error("Error using SecurityHubPuller (%s), falling back to basic client", e)
+        return _fetch_with_error_handling(aws_client, "get_findings", "Findings", "findings")
 
 
 def fetch_aws_findings_v2(aws_client: BaseClient) -> list:
-    """Fetch AWS Findings
+    """Fetch AWS Findings using v2 API
 
     :param BaseClient aws_client: AWS Security Hub Client
     :return: AWS Findings
     :rtype: list
     """
-    findings = []
-    try:
-        response = aws_client.get_findings_v2()
-        findings = response.get("Findings", [])
-    except ClientError as cex:
-        create_logger().error("Unexpected error when fetching resources from AWS: %s", cex)
-    return findings
+    return _fetch_with_error_handling(aws_client, "get_findings_v2", "Findings", "findings")
 
 
 def fetch_aws_resources(aws_client: BaseClient) -> list:
@@ -174,11 +171,4 @@ def fetch_aws_resources(aws_client: BaseClient) -> list:
     :return: AWS Resources
     :rtype: list
     """
-    resources = []
-    try:
-        response = aws_client.get_resources_v2()
-        resources = response.get("Resources", [])
-        logger.info(f"Fetched {len(resources)} resources from Security Hub")
-    except ClientError as cex:
-        create_logger().error("Unexpected error when fetching resources from AWS: %s", cex)
-    return resources
+    return _fetch_with_error_handling(aws_client, "get_resources_v2", "Resources", "resources")