diff --git a/local-app/python-tools/cross-organization/org_runner.py b/local-app/python-tools/cross-organization/org_runner.py index 6bd73ea9..3f95b757 100755 --- a/local-app/python-tools/cross-organization/org_runner.py +++ b/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,64 +58,59 @@ 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: @@ -115,7 +118,12 @@ 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") + # 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() + +