refactor some, make more generic

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()