Merge pull request #1 from meese-enterprises/fix/regions-issue

Improved error handling

Author: ajmeese7

commands/collect.py

@@ -11,7 +11,7 @@
 import pyjq
 import urllib.parse
 from botocore.exceptions import ClientError, EndpointConnectionError, NoCredentialsError
-from shared.common import get_account, custom_serializer
+from shared.common import custom_serializer, get_account, get_default_region
 from botocore.config import Config
 
 __description__ = "Run AWS API calls to collect data from the account"
@@ -223,13 +223,7 @@ def collect(arguments):
     make_directory("account-data/{}".format(account_dir))
 
     # Identify the default region used by global services such as IAM
-    default_region = os.environ.get("AWS_REGION", "us-east-1")
-    if "gov-" in default_region:
-        default_region = "us-gov-west-1"
-    elif "cn-" in default_region:
-        default_region = "cn-north-1"
-    else:
-        default_region = "us-east-1"
+    default_region = get_default_region()
 
     regions_filter = None
     if len(arguments.regions_filter) > 0:

commands/weboftrust.py

@@ -198,6 +198,7 @@ def get_iam_trusts(account, nodes, connections, connections_to_get):
                         # Validate that the federated principal and the SAML provider is coming from known accounts.
                         # WoT will show us the direction of that trust for further inspection.
                         # this enables cross_account_admin_sts (STS between accounts)
+                        saml_provider_arn = ""
                         for saml in saml_providers:
                             if saml["Arn"] == federated_principal:
                                 saml_provider_arn = saml["Arn"]
@@ -296,7 +297,7 @@ def get_iam_trusts(account, nodes, connections, connections_to_get):
                                 }
                             )
                             continue
-                        else:
+                        elif saml_provider_arn != "":
                             raise Exception(
                                 "Unknown federation provider: {}".format(
                                     saml_provider_arn.lower()

shared/common.py

@@ -1,11 +1,12 @@
 from __future__ import print_function
+from netaddr import IPNetwork
 import argparse
-import json
 import datetime
+import json
+import os
 import pyjq
-import yaml
 import sys
-from netaddr import IPNetwork
+import yaml
 
 from shared.nodes import Account, Region
 from shared.query import query_aws, get_parameter_file
@@ -152,11 +153,25 @@ def is_unblockable_cidr(cidr):
     return False
 
 
+def get_default_region():
+    default_region = os.environ.get("AWS_REGION", "us-east-1")
+    if "gov-" in default_region:
+        return "us-gov-west-1"
+    elif "cn-" in default_region:
+        return "cn-north-1"
+    else:
+        return "us-east-1"
+
+
 def get_regions(account, outputfilter={}):
     # aws ec2 describe-regions
     region_data = query_aws(account, "describe-regions")
+
+    # Fallback to default region if no region data is found
     if not region_data:
-        raise InvalidAccountData("region data not found for {}".format(account.name))
+        default_region = get_default_region()
+        log_warning(f"No region data found for account {account.name}. Falling back to default region: {default_region}")
+        return [{"RegionName": default_region}]
 
     region_filter = ""
     if "regions" in outputfilter: