From 789c26065ed9c8ba51c3fa234efab74776c04007 Mon Sep 17 00:00:00 2001 From: badra001 Date: Wed, 31 Dec 2025 10:22:27 -0500 Subject: [PATCH] refactor some, make more generic --- .../test-cross-organization/test-cross-org.py | 120 +++++++++++------- 1 file changed, 76 insertions(+), 44 deletions(-) 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 741d2df2..be7f4da4 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 @@ -1,4 +1,4 @@ -#!/bin/env python +#!/usr/bin/env python import boto3 import argparse @@ -11,21 +11,47 @@ from botocore.exceptions import ClientError, ProfileNotFound # --- VERSIONING --- -__version__ = "1.0.21" - -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") +__version__ = "1.1.0" + +# ============================================================================== +# GENERAL PURPOSE WRAPPER LOGIC +# ============================================================================== +def account_task(account_session, account_id, account_name, region): + """ + General purpose wrapper for per-account logic. + Replace the logic inside this function to run different tasks. - # Updated output argument to accept optional values (nargs='?') - parser.add_argument("--output", nargs='?', const='DEFAULT', - help="Save results to file. If no name provided, uses results_ORGID.ISO-DATE.txt") + :param account_session: A boto3.Session authenticated in the member account. + :param account_id: The ID of the account being processed. + :param account_name: The name of the account in the Organization. + :param region: The region to operate in. + :return: A dictionary containing the results of the task. + """ + results = {"alias": "None", "status": "SUCCESS", "error": None} - parser.add_argument("--debug", action="store_true", help="Enable verbose wire-trace logging") - return parser.parse_args() + try: + # EXAMPLE TASK: Get IAM Alias + iam = account_session.client('iam') + aliases = iam.list_account_aliases().get('AccountAliases', []) + if aliases: + results["alias"] = aliases[0] + + # ADD YOUR CUSTOM LOGIC HERE: + # e.g., config = account_session.client('config') + # e.g., s3 = account_session.client('s3') + + except ClientError as e: + results["status"] = f"FAIL({e.response['Error']['Code']})" + results["error"] = str(e) + except Exception as e: + results["status"] = "FAIL(Unexpected)" + results["error"] = str(e) + + return results + +# ============================================================================== +# HIERARCHY & UTILITY LOGIC +# ============================================================================== def get_full_path(org_client, entity_id, cache): """Resolves hierarchy path and caches results. Strips 'Root'.""" @@ -47,7 +73,18 @@ def get_full_path(org_client, entity_id, cache): except Exception: return f"Error({entity_id})", entity_id -def test_connectivity(args): +def get_args(): + parser = argparse.ArgumentParser(description=f"AWS Org Task Runner 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", nargs='?', const='DEFAULT', help="Save results to file.") + parser.add_argument("--debug", action="store_true", help="Enable verbose wire-trace logging") + return parser.parse_args() + +def main(): + args = get_args() start_time = time.perf_counter() output_log = [] @@ -71,27 +108,21 @@ def log_print(message=""): try: org_info = org_client.describe_organization() org_id = org_info['Organization']['Id'] - except Exception: - pass + except Exception: pass - # Handle Default Naming Logic final_filename = None if args.output: - # ISO timestamp: YYYY-MM-DDTHHMMSS datestamp = datetime.now().strftime("%Y-%m-%dT%H%M%S") - if args.output == 'DEFAULT': final_filename = f"results_{org_id}.{datestamp}.txt" else: p = Path(args.output) - # If user provided a name, we still append the timestamp for safety final_filename = f"{p.stem}_{datestamp}{p.suffix}" log_print("-" * 100) - log_print(f"AWS ORG CONNECTIVITY TEST - Version {__version__}") + log_print(f"AWS ORG TASK RUNNER - 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}") @@ -100,7 +131,6 @@ def log_print(message=""): log_print(f" Output File: {final_filename}") log_print("-" * 100) - log_print("Gathering Organization accounts...") all_accounts = [] paginator = org_client.get_paginator('list_accounts') for page in paginator.paginate(): @@ -109,13 +139,10 @@ def log_print(message=""): sort_key = 'Name' if args.sort == 'name' else 'Id' all_accounts.sort(key=lambda x: x[sort_key].lower()) - table_width = 215 - 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) + log_print(f"{'#':<4} | {'Account Name':<25} | {'Alias':<40} | {'Account ID':<15} | {'Status':<12} | {'OU Path'}") + log_print("-" * 160) - hierarchy_cache = {} - ou_summary = {} + hierarchy_cache, ou_summary = {}, {} for i, acc in enumerate(all_accounts, 1): acc_id, acc_name = acc['Id'], acc['Name'] @@ -129,24 +156,32 @@ def log_print(message=""): summary_key = (ou_path, ou_id) 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") + assumed = sts_client.assume_role(RoleArn=role_arn, RoleSessionName="OrgTaskRunner") m_sess = boto3.Session( 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', []) - alias = aliases[0] if aliases else "None" - status_text = "SUCCESS" - ou_summary[summary_key]["success"] += 1 + + # EXECUTE THE MODULAR TASK + task_result = account_task(m_sess, acc_id, acc_name, session.region_name) + + status_text = task_result['status'] + alias = task_result['alias'] + + if "SUCCESS" in status_text: + ou_summary[summary_key]["success"] += 1 + else: + ou_summary[summary_key]["fail"] += 1 + except ClientError as e: status_text = f"FAIL({e.response['Error']['Code']})" + alias = "ERROR" ou_summary[summary_key]["fail"] += 1 - log_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[:40]:<40} | {acc_id:<15} | {status_text:<12} | {ou_path}") end_time = time.perf_counter() elapsed_seconds = end_time - start_time @@ -162,15 +197,12 @@ def log_print(message=""): log_print("=" * 85) if final_filename: - try: - with open(final_filename, 'w') as f: - f.write("\n".join(output_log)) - print(f"\nDone. Results saved to {final_filename}") - except Exception as e: - print(f"Warning: Could not save output file: {e}") + with open(final_filename, 'w') as f: + f.write("\n".join(output_log)) + print(f"\nResults saved to {final_filename}") except Exception as e: print(f"CRITICAL ERROR: {e}") if __name__ == "__main__": - test_connectivity(get_args()) + main()