fix partition

local-app/python-tools/aws-service-linked-roles/aws-slr-generate.py

@@ -5,19 +5,31 @@
 import sys
 from botocore.exceptions import ClientError, ProfileNotFound
 
-VERSION = "1.0.2"
+VERSION = "1.0.4"
 
-def get_account_id(session):
+def get_session_info(session):
+    """
+    Retrieves the Account ID and Partition from the current session.
+    """
     sts = session.client('sts')
-    return sts.get_caller_identity()['Account']
+    identity = sts.get_caller_identity()
+    arn = identity['Arn']
+
+    # ARN format: arn:partition:service:region:account-id:resource-type/resource-id
+    # We split by ':' and take the second element for the partition.
+    arn_parts = arn.split(':')
+    partition = arn_parts[1]
+    account_id = identity['Account']
+
+    return account_id, partition
 
 def generate_slr_terraform():
-    parser = argparse.ArgumentParser(description="Audit AWS Service Linked Roles and generate Terraform import/vars files.")
+    parser = argparse.ArgumentParser(description="Audit AWS SLRs and generate TF import/vars files.")
     parser.add_argument("--profile", help="AWS CLI profile to use")
     parser.add_argument("--region", help="AWS region (e.g., us-gov-west-1)")
     parser.add_argument("--import-file", default="import.tf", help="Output filename for import blocks")
     parser.add_argument("--variables-file", default="variables.service-linked-roles.auto.tfvars", help="Output filename for tfvars")
-    parser.add_argument("--dry-run", action="store_true", help="Print the generated content to stdout without creating files")
+    parser.add_argument("--dry-run", action="store_true", help="Print content to stdout instead of files")
     parser.add_argument("--version", action="version", version=f"%(prog)s {VERSION}")
 
     args = parser.parse_args()
@@ -26,7 +38,9 @@ def generate_slr_terraform():
         # Initialize Session
         session = boto3.Session(profile_name=args.profile, region_name=args.region)
         iam = session.client('iam')
-        account_id = get_account_id(session)
+
+        # Detect Partition and Account ID dynamically
+        account_id, partition = get_session_info(session)
 
         service_names = []
         slr_count = 0
@@ -40,6 +54,7 @@ def generate_slr_terraform():
                     principal = statement.get('Principal', {}).get('Service')
                     if principal:
                         slr_count += 1
+                        # Service can be a string or a list
                         if isinstance(principal, list):
                             service_names.extend(principal)
                         else:
@@ -48,20 +63,21 @@ def generate_slr_terraform():
         # De-duplicate and sort
         unique_services = sorted(list(set(service_names)))
 
-        # Prepare Content
-        import_content = [f"# Generated by SLR-to-TF v{VERSION}\n"]
+        # Prepare File Content
+        import_content = [f"# Generated by aws-slr-generate v{VERSION}\n"]
         for service in unique_services:
             import_content.append(f'import {{\n')
             import_content.append(f'  to = aws_iam_service_linked_role.roles["{service}"]\n')
-            import_content.append(f'  id = "arn:aws:iam::{account_id}:role/aws-service-role/{service}"\n')
+            # Use the detected partition here
+            import_content.append(f'  id = "arn:{partition}:iam::{account_id}:role/aws-service-role/{service}"\n')
             import_content.append(f'}}\n\n')
 
-        vars_content = [f"# Generated by SLR-to-TF v{VERSION}\n", "service_linked_roles = [\n"]
+        vars_content = [f"# Generated by aws-slr-generate v{VERSION}\n", "service_linked_roles = [\n"]
         for service in unique_services:
             vars_content.append(f'  "{service}",\n')
         vars_content.append("]\n")
 
-        # Output Logic
+        # Write or Print
         if args.dry_run:
             print("\n" + "="*20 + " DRY RUN: IMPORT FILE " + "="*20)
             print("".join(import_content))
@@ -78,27 +94,26 @@ def generate_slr_terraform():
         print("-" * 45)
         print(f"SLR-to-TF Generator v{VERSION}")
         print("-" * 45)
+        print(f"AWS Partition:    {partition}")
         print(f"AWS Account:      {account_id}")
-        print(f"Region:           {session.region_name or 'Default'}")
         print(f"Roles Found:      {slr_count}")
         print(f"Unique Services:  {len(unique_services)}")
 
         if args.dry_run:
-            print(f"Mode:             DRY RUN (No files created)")
+            print(f"Mode:             DRY RUN (Files not created)")
         else:
             print(f"Import File:      {args.import_file}")
             print(f"Variables File:   {args.variables_file}")
-
         print("-" * 45)
 
     except ProfileNotFound:
-        print(f"Error: The profile '{args.profile}' was not found.")
+        print(f"Error: Profile '{args.profile}' not found.")
         sys.exit(1)
     except ClientError as e:
         print(f"AWS Error: {e}")
         sys.exit(1)
     except Exception as e:
-        print(f"An unexpected error occurred: {e}")
+        print(f"Unexpected error: {e}")
         sys.exit(1)
 
 if __name__ == "__main__":