add progress bar

local-app/python-tools/cross-organization/org_runner.py

@@ -12,8 +12,16 @@
 from pathlib import Path
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
+# Try to import tqdm, fallback to a dummy if not installed
+try:
+    from tqdm import tqdm
+except ImportError:
+    print("Warning: 'tqdm' not found. Install it with 'pip install tqdm' for a progress bar.")
+    # Simple fallback to avoid crashing
+    def tqdm(iterable, **kwargs): return iterable
+
 # --- VERSIONING ---
-__version__ = "1.5.0"
+__version__ = "1.5.1"
 
 class OrgTaskRunner:
     def __init__(self, args):
@@ -23,13 +31,13 @@ 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)
+        # We only log to the internal buffer for the final file write
+        # Standard printing is handled by tqdm to avoid screen flickering
         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
+    def process_account(self, acc_id, acc_name, partition, tasks):
+        """Worker function for each thread."""
+        # Create a fresh session for thread safety
         thread_session = boto3.Session(profile_name=self.args.profile, region_name=self.args.region)
         sts = thread_session.client('sts')
 
@@ -50,72 +58,72 @@ def process_account(self, acc_id, acc_name, partition, tasks, total_count, index
                 account_data["alias"] = res.get("alias")
                 account_data["checks"].update(res.get("data", {}))
 
-            return account_data, f"[{index}/{total_count}] Processed {acc_name}"
+            return account_data, None # Success
         except Exception as e:
-            return None, f"[{index}/{total_count}] FAILED {acc_name}: {str(e)}"
+            return None, f"FAILED {acc_name} ({acc_id}): {str(e)}"
 
     def run(self):
         self.start_time = time.perf_counter()
 
-        # 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 = []
-        task_versions = []
         if self.args.enable_checks:
             sys.path.append(os.getcwd())
             for m in self.args.enable_checks:
-                module_name = m.replace('.py', '')
-                module = importlib.import_module(module_name)
-                v = getattr(module, '__version__', 'Unknown')
+                module = importlib.import_module(m.replace('.py', ''))
                 tasks.append(getattr(module, 'account_task'))
-                task_versions.append(f"{module_name} (v{v})")
+
+        print("-" * 100)
+        print(f"AWS ORG TASK RUNNER - v{__version__} | Workers: {self.args.max_workers}")
+        print("-" * 100)
 
         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())
 
-        self.log_print("-" * 100)
-        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)
-
-        # PARALLEL EXECUTION
+        # PARALLEL EXECUTION WITH TQDM
         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)
-            ]
+            futures = {
+                executor.submit(self.process_account, acc['Id'], acc['Name'], partition, tasks): acc 
+                for acc in all_accounts
+            }
 
-            for future in as_completed(futures):
-                data, log_msg = future.result()
-                if data:
-                    self.full_results.append(data)
-                self.log_print(log_msg)
+            # tqdm context manager wraps the as_completed generator
+            with tqdm(total=len(all_accounts), desc="Processing Accounts", unit="acc", colour="green") as pbar:
+                for future in as_completed(futures):
+                    data, error = future.result()
+                    if data:
+                        self.full_results.append(data)
+                        self.log_print(f"SUCCESS: {data['account_name']}")
+                    else:
+                        self.log_print(error)
+                    pbar.update(1)
 
-        # EXPORTS (JSON & CSV)
+        # FINAL EXPORTS
         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:
+            csv_file = f"{base}.{ds}.csv"
+            with open(csv_file, 'w', newline='') as f:
                 writer = csv.writer(f)
                 writer.writerow(["account_id", "account_alias", "region", "field_name", "field_value"])
                 for acc in self.full_results:
                     for reg, fields in acc["checks"].items():
                         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")
+            # Save JSON
+            json_file = f"{base}.{ds}.json"
+            with open(json_file, 'w') as f:
+                json.dump(self.full_results, f, indent=2)
+
+        print(f"\nTotal Time: {round(time.perf_counter() - self.start_time, 2)}s")
 
 if __name__ == "__main__":
     p = argparse.ArgumentParser()
@@ -125,5 +133,7 @@ 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")
+    p.add_argument("--max-workers", type=int, default=8)
     OrgTaskRunner(p.parse_args()).run()
+
+