add output to file

local-app/python-tools/test-cross-organization/test-cross-org.py

@@ -7,48 +7,63 @@
 from botocore.exceptions import ClientError, ProfileNotFound
 
 # --- VERSIONING ---
-__version__ = "1.0.17"
+__version__ = "1.0.18"
+
+class Tee(object):
+    """Custom stream to write to multiple destinations (terminal and file)."""
+    def __init__(self, *files):
+        self.files = files
+    def write(self, obj):
+        for f in self.files:
+            f.write(obj)
+            f.flush() # Ensure real-time updates
+    def flush(self):
+        for f in self.files:
+            f.flush()
 
 def get_args():
     parser = argparse.ArgumentParser(description=f"AWS Org Connectivity Test v{__version__}")
     parser.add_argument("--profile", default=os.environ.get("AWS_PROFILE"), help="AWS CLI profile")
     parser.add_argument("--region", default=os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), help="AWS Region")
     parser.add_argument("--role-name", required=True, help="Role name to assume in member accounts")
     parser.add_argument("--sort", choices=['id', 'name'], default='name', help="Field to sort results by")
+    parser.add_argument("--output", help="Optional filename to save the results to")
     parser.add_argument("--debug", action="store_true", help="Enable verbose wire-trace logging")
     return parser.parse_args()
 
 def get_full_path(org_client, entity_id, cache):
     """Resolves hierarchy path and caches results. Strips 'Root'."""
     if entity_id in cache:
         return cache[entity_id]
-
     if entity_id.startswith('r-'):
         cache[entity_id] = (None, entity_id)
         return None, entity_id
-
     try:
         parent_resp = org_client.list_parents(ChildId=entity_id)
         parents = parent_resp.get('Parents', [])
-        if not parents:
-            return "Unknown", entity_id
-
-        parent_id = parents[0]['Id']
+        parent_id = parents[0]['Id'] if parents else None
         ou_desc = org_client.describe_organizational_unit(OrganizationalUnitId=entity_id)
         ou_name = ou_desc['OrganizationalUnit']['Name']
-
-        parent_path, _ = get_full_path(org_client, parent_id, cache)
+        parent_path, _ = get_full_path(org_client, parent_id, cache) if parent_id else (None, None)
         full_path = f"{parent_path}:{ou_name}" if parent_path else ou_name
-
         cache[entity_id] = (full_path, entity_id)
         return full_path, entity_id
     except Exception:
         return f"Error({entity_id})", entity_id
 
 def test_connectivity(args):
-    # Start high-precision timer
     start_time = time.perf_counter()
 
+    # Setup redirection if --output is provided
+    f = None
+    original_stdout = sys.stdout
+    if args.output:
+        try:
+            f = open(args.output, 'w')
+            sys.stdout = Tee(sys.stdout, f)
+        except Exception as e:
+            print(f"Warning: Could not open output file: {e}")
+
     if args.debug:
         boto3.set_stream_logger(name='botocore', level=logging.DEBUG)
 
@@ -60,7 +75,21 @@ def test_connectivity(args):
         caller = sts_client.get_caller_identity()
         partition = caller['Arn'].split(':')[1]
 
-        print("Gathering Organization accounts and hierarchy...")
+        # RESTORED HEADER
+        print("-" * 100)
+        print(f"AWS ORG CONNECTIVITY TEST - Version {__version__}")
+        print("-" * 100)
+        print(f"CONFIGURATION:")
+        print(f"  Partition:     {partition}")
+        print(f"  Profile:       {args.profile or 'Environment/Default'}")
+        print(f"  Source Identity: {caller['Arn']}")
+        print(f"  Base Region:   {session.region_name}")
+        print(f"  Target Role:   {args.role_name}")
+        if args.output:
+            print(f"  Output File:   {args.output}")
+        print("-" * 100)
+
+        print("Gathering Organization accounts...")
         all_accounts = []
         paginator = org_client.get_paginator('list_accounts')
         for page in paginator.paginate():
@@ -71,49 +100,35 @@ def test_connectivity(args):
 
         table_width = 215
         print("-" * table_width)
-        print(f"AWS ORG CONNECTIVITY TEST - v{__version__} | Total Accounts: {len(all_accounts)}")
-        print("-" * table_width)
         print(f"{'#':<4} | {'Account Name':<25} | {'Alias':<64} | {'Account ID':<15} | {'Status':<12} | {'Full OU Path'}")
         print("-" * table_width)
 
         hierarchy_cache = {} 
         ou_summary = {} 
 
         for i, acc in enumerate(all_accounts, 1):
-            acc_id = acc['Id']
-            acc_name = acc['Name']
+            acc_id, acc_name = acc['Id'], acc['Name']
             role_arn = f"arn:{partition}:iam::{acc_id}:role/{args.role_name}"
 
             parent_resp = org_client.list_parents(ChildId=acc_id)
             parents = parent_resp.get('Parents', [])
-
-            if parents:
-                parent_id = parents[0]['Id']
-                raw_path, ou_id = get_full_path(org_client, parent_id, hierarchy_cache)
-                ou_path = raw_path if raw_path else "N/A (Root)"
-            else:
-                ou_path = "Orphaned"
-                ou_id = "N/A"
+            raw_path, ou_id = get_full_path(org_client, parents[0]['Id'], hierarchy_cache) if parents else ("Orphaned", "N/A")
+            ou_path = raw_path if raw_path else "N/A (Root)"
 
             summary_key = (ou_path, ou_id)
-            if summary_key not in ou_summary:
-                ou_summary[summary_key] = {"success": 0, "fail": 0}
-
-            alias = "None"
-            status_text = "PENDING"
+            if summary_key not in ou_summary: ou_summary[summary_key] = {"success": 0, "fail": 0}
 
+            alias, status_text = "None", "PENDING"
             try:
                 assumed = sts_client.assume_role(RoleArn=role_arn, RoleSessionName="ConnTest")
-                creds = assumed['Credentials']
                 m_sess = boto3.Session(
-                    aws_access_key_id=creds['AccessKeyId'],
-                    aws_secret_access_key=creds['SecretAccessKey'],
-                    aws_session_token=creds['SessionToken'],
+                    aws_access_key_id=assumed['Credentials']['AccessKeyId'],
+                    aws_secret_access_key=assumed['Credentials']['SecretAccessKey'],
+                    aws_session_token=assumed['Credentials']['SessionToken'],
                     region_name=session.region_name
                 )
                 aliases = m_sess.client('iam').list_account_aliases().get('AccountAliases', [])
-                if aliases:
-                    alias = aliases[0]
+                alias = aliases[0] if aliases else "None"
                 status_text = "SUCCESS"
                 ou_summary[summary_key]["success"] += 1
             except ClientError as e:
@@ -122,24 +137,25 @@ def test_connectivity(args):
 
             print(f"{i:<4} | {acc_name[:25]:<25} | {alias:<64} | {acc_id:<15} | {status_text:<12} | {ou_path}")
 
-        # Stop timer
         end_time = time.perf_counter()
         elapsed_seconds = end_time - start_time
         minutes, seconds = divmod(int(elapsed_seconds), 60)
 
-        # Output Summary with Elapsed Time
         print("\n" + "=" * 85)
         print(f"{'OU PATH SUMMARY':<45} | {'OU ID':<18} | {'SUCC':<6} | {'FAIL':<6}")
         print("=" * 85)
         for (path, ou_id), counts in sorted(ou_summary.items(), key=lambda x: x[0][0] if x[0][0] else ""):
             print(f"{path[:45]:<45} | {ou_id:<18} | {counts['success']:<6} | {counts['fail']:<6}")
         print("=" * 85)
-        print(f"Total Elapsed Time: {minutes}m {seconds}s")
+        print(f"Total Elapsed Time: {minutes}m {seconds}s | Total Accounts: {len(all_accounts)}")
         print("=" * 85)
 
     except Exception as e:
         print(f"CRITICAL ERROR: {e}")
+    finally:
+        if f:
+            sys.stdout = original_stdout
+            f.close()
 
 if __name__ == "__main__":
     test_connectivity(get_args())
-