From 8939ec2d9257cad1674ea5ea4ffabf0e05b4d158 Mon Sep 17 00:00:00 2001
From: badra001 <donald.e.badrak.ii@census.gov>
Date: Tue, 30 Dec 2025 16:17:39 -0500
Subject: [PATCH] v1.0.2

---
 .../inf-org-crossaccount/test-cross-org.py    | 121 ++++++++++--------
 1 file changed, 68 insertions(+), 53 deletions(-)

diff --git a/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py b/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py
index f9586017..ec9f088f 100755
--- a/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py
+++ b/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py
@@ -1,56 +1,71 @@
 import boto3
-from botocore.exceptions import ClientError
-
-# --- SET THESE TO MATCH YOUR TERRAFORM ---
-CROSS_ACCOUNT_ROLE_NAME = "OrganizationAdminAccessRole" 
-TARGET_REGION = "us-east-1"
-
-def test_connectivity():
-    # 1. Initialize the Organizations client from Management Account
-    org_client = boto3.client('organizations')
-    sts_client = boto3.client('sts')
-    
-    print(f"Starting connectivity test using role: {CROSS_ACCOUNT_ROLE_NAME}\n")
-    print(f"{'Account Name':<25} | {'Account ID':<15} | {'Status'}")
-    print("-" * 55)
-
-    # 2. Page through all accounts in the Org
-    paginator = org_client.get_paginator('list_accounts')
-    for page in paginator.paginate():
-        for account in page['Accounts']:
-            # Skip accounts that aren't active (suspended/pending)
-            if account['State'] != 'ACTIVE':
-                continue
-
-            acc_id = account['Id']
-            acc_name = account['Name']
-            role_arn = f"arn:aws:iam::{acc_id}:role/{CROSS_ACCOUNT_ROLE_NAME}"
-
-            try:
-                # 3. Attempt to assume the role
-                response = sts_client.assume_role(
-                    RoleArn=role_arn,
-                    RoleSessionName="ConnectivityTestSession"
-                )
-
-                # 4. Use temporary credentials to create a target session
-                creds = response['Credentials']
-                target_session = boto3.Session(
-                    aws_access_key_id=creds['AccessKeyId'],
-                    aws_secret_access_key=creds['SecretAccessKey'],
-                    aws_session_token=creds['SessionToken']
-                )
-
-                # 5. Perform a "whoami" call in the target account
-                target_sts = target_session.client('sts')
-                identity = target_sts.get_caller_identity()
-                
-                print(f"{acc_name[:25]:<25} | {acc_id:<15} | ✅ Success")
-
-            except ClientError as e:
-                # Catch failures (e.g., role hasn't finished deploying or trust policy issue)
-                error_code = e.response['Error']['Code']
-                print(f"{acc_name[:25]:<25} | {acc_id:<15} | ❌ Failed ({error_code})")
+import argparse
+import sys
+from botocore.exceptions import ClientError, ProfileNotFound
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Test AWS Org Connectivity with assuming roles.")
+    parser.add_argument("--profile", required=True, help="AWS CLI profile to use (e.g., your-sso-profile)")
+    parser.add_argument("--region", default="us-east-1", help="Default region for the session (default: us-east-1)")
+    parser.add_argument("--role-name", required=True, help="The name of the role to assume in member accounts")
+    return parser.parse_args()
+
+def test_connectivity(args):
+    try:
+        # Create the base session using your local profile
+        session = boto3.Session(profile_name=args.profile, region_name=args.region)
+        org_client = session.client('organizations')
+        sts_client = session.client('sts')
+        
+        # Verify who you are in the Management Account first
+        caller = sts_client.get_caller_identity()
+        print(f"Authenticated as: {caller['Arn']}")
+        print(f"Targeting role: {args.role_name}\n")
+        
+        print(f"{'Account Name':<25} | {'Account ID':<15} | {'Status'}")
+        print("-" * 60)
+
+        paginator = org_client.get_paginator('list_accounts')
+        for page in paginator.paginate():
+            for account in page['Accounts']:
+                if account['Status'] != 'ACTIVE':
+                    continue
+
+                acc_id = account['Id']
+                acc_name = account['Name']
+                role_arn = f"arn:aws:iam::{acc_id}:role/{args.role_name}"
+
+                try:
+                    # Assume role in the member account
+                    assumed_role = sts_client.assume_role(
+                        RoleArn=role_arn,
+                        RoleSessionName="OrgTestSession"
+                    )
+                    
+                    # Prove it worked by creating a client with the temp credentials
+                    creds = assumed_role['Credentials']
+                    member_session = boto3.Session(
+                        aws_access_key_id=creds['AccessKeyId'],
+                        aws_secret_access_key=creds['SecretAccessKey'],
+                        aws_session_token=creds['SessionToken'],
+                        region_name=args.region
+                    )
+                    
+                    # Call target account identity
+                    member_id = member_session.client('sts').get_caller_identity()['Account']
+                    print(f"{acc_name[:25]:<25} | {acc_id:<15} | ✅ Success (Verified: {member_id})")
+
+                except ClientError as e:
+                    print(f"{acc_name[:25]:<25} | {acc_id:<15} | ❌ Failed ({e.response['Error']['Code']})")
+
+    except ProfileNotFound:
+        print(f"Error: The profile '{args.profile}' was not found in your AWS config.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        sys.exit(1)
 
 if __name__ == "__main__":
-    test_connectivity()
+    cli_args = get_args()
+    test_connectivity(cli_args)
+