From c77a4f38536194c1839224febe94f456e47dc060 Mon Sep 17 00:00:00 2001 From: badra001 Date: Fri, 2 Jan 2026 08:19:16 -0500 Subject: [PATCH] add parallel operation --- .../cross-organization/assess_check_config.py | 4 + .../cross-organization/org_runner.py | 89 ++++++++++++------- 2 files changed, 60 insertions(+), 33 deletions(-) diff --git a/local-app/python-tools/cross-organization/assess_check_config.py b/local-app/python-tools/cross-organization/assess_check_config.py index 327e864e..548a07c4 100644 --- a/local-app/python-tools/cross-organization/assess_check_config.py +++ b/local-app/python-tools/cross-organization/assess_check_config.py @@ -1,7 +1,11 @@ +#!/usr/bin/env python + import json import argparse import re +__version__ = "1.0.0" + def main(): parser = argparse.ArgumentParser() parser.add_argument("--input", required=True, help="The JSON audit file") diff --git a/local-app/python-tools/cross-organization/org_runner.py b/local-app/python-tools/cross-organization/org_runner.py index 4d75a51f..6bd73ea9 100755 --- a/local-app/python-tools/cross-organization/org_runner.py +++ b/local-app/python-tools/cross-organization/org_runner.py @@ -10,9 +10,10 @@ import importlib from datetime import datetime from pathlib import Path +from concurrent.futures import ThreadPoolExecutor, as_completed # --- VERSIONING --- -__version__ = "1.4.0" +__version__ = "1.5.0" class OrgTaskRunner: def __init__(self, args): @@ -22,14 +23,45 @@ def __init__(self, args): self.start_time = 0 def log_print(self, message=""): + # Note: In multi-threading, print() is mostly thread-safe in Python 3. print(message) self.output_log.append(message) + def process_account(self, acc_id, acc_name, partition, tasks, total_count, index): + """Worker function for each thread/account.""" + # Thread-safe: Each worker gets its OWN session + thread_session = boto3.Session(profile_name=self.args.profile, region_name=self.args.region) + sts = thread_session.client('sts') + + role_arn = f"arn:{partition}:iam::{acc_id}:role/{self.args.role_name}" + account_data = {"account_id": acc_id, "account_name": acc_name, "checks": {}, "alias": "N/A"} + + try: + assumed = sts.assume_role(RoleArn=role_arn, RoleSessionName="OrgRunnerTask") + 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=self.args.region + ) + + for t_func in tasks: + res = t_func(m_sess, acc_id, acc_name, self.args.region) + account_data["alias"] = res.get("alias") + account_data["checks"].update(res.get("data", {})) + + return account_data, f"[{index}/{total_count}] Processed {acc_name}" + except Exception as e: + return None, f"[{index}/{total_count}] FAILED {acc_name}: {str(e)}" + def run(self): self.start_time = time.perf_counter() - session = boto3.Session(profile_name=self.args.profile, region_name=self.args.region) - sts = session.client('sts') - org = session.client('organizations') + + # Setup main session to list accounts + main_session = boto3.Session(profile_name=self.args.profile, region_name=self.args.region) + org = main_session.client('organizations') + sts = main_session.client('sts') + partition = sts.get_caller_identity()['Arn'].split(':')[1] # Load tasks and capture versions tasks = [] @@ -43,50 +75,38 @@ def run(self): tasks.append(getattr(module, 'account_task')) task_versions.append(f"{module_name} (v{v})") - paginator = org.get_paginator('list_accounts') - all_accounts = [acc for page in paginator.paginate() for acc in page['Accounts'] if acc['Status'] == 'ACTIVE'] + all_accounts = [acc for page in org.get_paginator('list_accounts').paginate() + for acc in page['Accounts'] if acc['Status'] == 'ACTIVE'] all_accounts.sort(key=lambda x: x['Name' if self.args.sort == 'name' else 'Id'].lower()) - # HEADER WITH VERSIONS self.log_print("-" * 100) - self.log_print(f"AWS ORG TASK RUNNER - v{__version__}") - self.log_print(f"Enabled Checks: {', '.join(task_versions) if task_versions else 'None'}") + self.log_print(f"AWS ORG TASK RUNNER - v{__version__} (Parallel Mode)") + self.log_print(f"Workers: {self.args.max_workers} | Checks: {', '.join(task_versions)}") self.log_print("-" * 100) - for i, acc in enumerate(all_accounts, 1): - acc_id, acc_name = acc['Id'], acc['Name'] - role_arn = f"arn:aws:iam::{acc_id}:role/{self.args.role_name}" + # PARALLEL EXECUTION + with ThreadPoolExecutor(max_workers=self.args.max_workers) as executor: + futures = [ + executor.submit(self.process_account, acc['Id'], acc['Name'], partition, tasks, len(all_accounts), i) + for i, acc in enumerate(all_accounts, 1) + ] - try: - assumed = sts.assume_role(RoleArn=role_arn, RoleSessionName="OrgRunnerTask") - 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=self.args.region - ) - - account_data = {"account_id": acc_id, "account_name": acc_name, "checks": {}} - - # EXECUTION WITH TIMING - for t_func in tasks: - res = t_func(m_sess, acc_id, acc_name, self.args.region) - account_data["alias"] = res.get("alias") - account_data["checks"].update(res.get("data", {})) - - self.full_results.append(account_data) - self.log_print(f"[{i}/{len(all_accounts)}] Processed {acc_name}") - except Exception as e: - self.log_print(f"[{i}/{len(all_accounts)}] FAILED {acc_name}: {e}") + for future in as_completed(futures): + data, log_msg = future.result() + if data: + self.full_results.append(data) + self.log_print(log_msg) # EXPORTS (JSON & CSV) if self.args.output: ds = datetime.now().strftime("%Y-%m-%dT%H%M%S") base = self.args.output if self.args.output != 'DEFAULT' else "audit_results" + # Save JSON with open(f"{base}.{ds}.json", 'w') as f: json.dump(self.full_results, f, indent=2) + # Save CSV with open(f"{base}.{ds}.csv", 'w', newline='') as f: writer = csv.writer(f) writer.writerow(["account_id", "account_alias", "region", "field_name", "field_value"]) @@ -95,6 +115,8 @@ def run(self): for k, v in fields.items(): writer.writerow([acc["account_id"], acc.get("alias"), reg, k, v]) + self.log_print(f"\nTotal Time: {round(time.perf_counter() - self.start_time, 2)}s") + if __name__ == "__main__": p = argparse.ArgumentParser() p.add_argument("--profile", default=os.environ.get("AWS_PROFILE")) @@ -103,4 +125,5 @@ def run(self): p.add_argument("--sort", choices=['id', 'name'], default='name') p.add_argument("--output", nargs='?', const='DEFAULT') p.add_argument("--enable-checks", nargs='+') + p.add_argument("--max-workers", type=int, default=4, help="Number of concurrent accounts to process") OrgTaskRunner(p.parse_args()).run()