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

---
 .../inf-org-crossaccount/test-cross-org.py    | 74 +++++++++++--------
 1 file changed, 43 insertions(+), 31 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 ec9f088f..22927268 100755
--- a/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py
+++ b/infrastructure/global/stacksets/inf-org-crossaccount/test-cross-org.py
@@ -1,29 +1,46 @@
 import boto3
 import argparse
+import os
 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")
+    parser = argparse.ArgumentParser(description="Test AWS Org Connectivity by assuming roles.")
+    
+    # Defaults are pulled from environment variables if the flags are omitted
+    parser.add_argument("--profile", 
+                        default=os.environ.get("AWS_PROFILE"),
+                        help="AWS CLI profile (default: AWS_PROFILE env var)")
+    
+    parser.add_argument("--region", 
+                        default=os.environ.get("AWS_DEFAULT_REGION", "us-east-1"),
+                        help="AWS Region (default: AWS_DEFAULT_REGION or 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
+        # Initialize session. Boto3 naturally handles None for profile_name 
+        # by falling back to 'default' or environment credentials.
         session = boto3.Session(profile_name=args.profile, region_name=args.region)
-        org_client = session.client('organizations')
+        
         sts_client = session.client('sts')
+        org_client = session.client('organizations')
         
-        # Verify who you are in the Management Account first
+        # Verify Management Account Identity
         caller = sts_client.get_caller_identity()
-        print(f"Authenticated as: {caller['Arn']}")
-        print(f"Targeting role: {args.role_name}\n")
+        print(f"--- Session Initiated ---")
+        print(f"Identity: {caller['Arn']}")
+        print(f"Profile:  {args.profile or 'Environment/Default'}")
+        print(f"Region:   {session.region_name}")
+        print(f"Role:     {args.role_name}\n")
         
         print(f"{'Account Name':<25} | {'Account ID':<15} | {'Status'}")
-        print("-" * 60)
+        print("-" * 65)
 
         paginator = org_client.get_paginator('list_accounts')
         for page in paginator.paginate():
@@ -36,36 +53,31 @@ def test_connectivity(args):
                 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(
+                    # Assume role using the Management Account's credentials
+                    assumed = sts_client.assume_role(
                         RoleArn=role_arn,
-                        RoleSessionName="OrgTestSession"
+                        RoleSessionName="OrgTest"
                     )
                     
-                    # 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
+                    # Verify by calling STS inside the member account
+                    temp = assumed['Credentials']
+                    member_sts = boto3.client(
+                        'sts',
+                        aws_access_key_id=temp['AccessKeyId'],
+                        aws_secret_access_key=temp['SecretAccessKey'],
+                        aws_session_token=temp['SessionToken']
                     )
+                    member_id = member_sts.get_caller_identity()['Account']
                     
-                    # 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})")
+                    print(f"{acc_name[:25]:<25} | {acc_id:<15} | ✅ Success")
 
                 except ClientError as e:
-                    print(f"{acc_name[:25]:<25} | {acc_id:<15} | ❌ Failed ({e.response['Error']['Code']})")
+                    print(f"{acc_name[:25]:<25} | {acc_id:<15} | ❌ {e.response['Error']['Code']}")
 
     except ProfileNotFound:
-        print(f"Error: The profile '{args.profile}' was not found in your AWS config.")
-        sys.exit(1)
+        print(f"Error: Profile '{args.profile}' not found.")
     except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        sys.exit(1)
+        print(f"Error: {e}")
 
 if __name__ == "__main__":
-    cli_args = get_args()
-    test_connectivity(cli_args)
-
+    test_connectivity(get_args())