simplify

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

@@ -7,27 +7,15 @@
 from botocore.exceptions import ClientError, ProfileNotFound
 
 # --- VERSIONING ---
-__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()
+__version__ = "1.0.19"
 
 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("--output", help="Optional filename to save the results to at the end")
     parser.add_argument("--debug", action="store_true", help="Enable verbose wire-trace logging")
     return parser.parse_args()
 
@@ -53,16 +41,14 @@ def get_full_path(org_client, entity_id, cache):
 
 def test_connectivity(args):
     start_time = time.perf_counter()
+
+    # List to capture all lines for final file dump
+    output_log = []
 
-    # 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}")
+    def log_print(message=""):
+        """Prints to console and appends to the output_log list."""
+        print(message)
+        output_log.append(message)
 
     if args.debug:
         boto3.set_stream_logger(name='botocore', level=logging.DEBUG)
@@ -75,21 +61,20 @@ def test_connectivity(args):
         caller = sts_client.get_caller_identity()
         partition = caller['Arn'].split(':')[1]
 
-        # 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}")
+        log_print("-" * 100)
+        log_print(f"AWS ORG CONNECTIVITY TEST - Version {__version__}")
+        log_print("-" * 100)
+        log_print(f"CONFIGURATION:")
+        log_print(f"  Partition:       {partition}")
+        log_print(f"  Profile:         {args.profile or 'Environment/Default'}")
+        log_print(f"  Source Identity: {caller['Arn']}")
+        log_print(f"  Base Region:     {session.region_name}")
+        log_print(f"  Target Role:     {args.role_name}")
         if args.output:
-            print(f"  Output File:   {args.output}")
-        print("-" * 100)
+            log_print(f"  Output File:     {args.output}")
+        log_print("-" * 100)
 
-        print("Gathering Organization accounts...")
+        log_print("Gathering Organization accounts...")
         all_accounts = []
         paginator = org_client.get_paginator('list_accounts')
         for page in paginator.paginate():
@@ -99,9 +84,9 @@ def test_connectivity(args):
         all_accounts.sort(key=lambda x: x[sort_key].lower())
 
         table_width = 215
-        print("-" * table_width)
-        print(f"{'#':<4} | {'Account Name':<25} | {'Alias':<64} | {'Account ID':<15} | {'Status':<12} | {'Full OU Path'}")
-        print("-" * table_width)
+        log_print("-" * table_width)
+        log_print(f"{'#':<4} | {'Account Name':<25} | {'Alias':<64} | {'Account ID':<15} | {'Status':<12} | {'Full OU Path'}")
+        log_print("-" * table_width)
 
         hierarchy_cache = {} 
         ou_summary = {} 
@@ -135,27 +120,32 @@ def test_connectivity(args):
                 status_text = f"FAIL({e.response['Error']['Code']})"
                 ou_summary[summary_key]["fail"] += 1
 
-            print(f"{i:<4} | {acc_name[:25]:<25} | {alias:<64} | {acc_id:<15} | {status_text:<12} | {ou_path}")
+            log_print(f"{i:<4} | {acc_name[:25]:<25} | {alias:<64} | {acc_id:<15} | {status_text:<12} | {ou_path}")
 
         end_time = time.perf_counter()
         elapsed_seconds = end_time - start_time
         minutes, seconds = divmod(int(elapsed_seconds), 60)
 
-        print("\n" + "=" * 85)
-        print(f"{'OU PATH SUMMARY':<45} | {'OU ID':<18} | {'SUCC':<6} | {'FAIL':<6}")
-        print("=" * 85)
+        log_print("\n" + "=" * 85)
+        log_print(f"{'OU PATH SUMMARY':<45} | {'OU ID':<18} | {'SUCC':<6} | {'FAIL':<6}")
+        log_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 | Total Accounts: {len(all_accounts)}")
-        print("=" * 85)
+            log_print(f"{path[:45]:<45} | {ou_id:<18} | {counts['success']:<6} | {counts['fail']:<6}")
+        log_print("=" * 85)
+        log_print(f"Total Elapsed Time: {minutes}m {seconds}s | Total Accounts: {len(all_accounts)}")
+        log_print("=" * 85)
+
+        # FINAL FILE DUMP
+        if args.output:
+            try:
+                with open(args.output, 'w') as f:
+                    f.write("\n".join(output_log))
+                print(f"\nDone. Results saved to {args.output}")
+            except Exception as e:
+                print(f"Warning: Could not save output file: {e}")
 
     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())