add parallel operation

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

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