diff --git a/local-app/python-tools/test-cross-organization/test-cross-org.py b/local-app/python-tools/test-cross-organization/test-cross-org.py index 361398e5..d12733b1 100755 --- a/local-app/python-tools/test-cross-organization/test-cross-org.py +++ b/local-app/python-tools/test-cross-organization/test-cross-org.py @@ -7,19 +7,7 @@ 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__}") @@ -27,7 +15,7 @@ def get_args(): 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())