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 3c58e3da..361398e5 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,7 +7,19 @@ 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__}") @@ -15,6 +27,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("--debug", action="store_true", help="Enable verbose wire-trace logging") return parser.parse_args() @@ -22,33 +35,35 @@ 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,8 +100,6 @@ 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) @@ -80,40 +107,28 @@ def test_connectivity(args): 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()) -