diff --git a/local-app/python-tools/cross-organization/purge_check_s3_buckets.py b/local-app/python-tools/cross-organization/purge_check_s3_buckets.py index b7176d04..8776aa70 100755 --- a/local-app/python-tools/cross-organization/purge_check_s3_buckets.py +++ b/local-app/python-tools/cross-organization/purge_check_s3_buckets.py @@ -1,136 +1,287 @@ #!/usr/bin/env python3 -import csv -import argparse +import os import sys import re +import csv +import queue import time +import argparse import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor + import boto3 +from rich.live import Live +from rich.console import Console +from rich.table import Table +from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn, SpinnerColumn -__version__ = "1.0.0" +# --- VERSIONING --- +__version__ = "2.0.0" -def purge_bucket_worker(account_id, bucket_name, role_name, dry_run): - """Worker task designed to clear all object versions and markers from a target bucket.""" - start_time = time.time() - deleted_count = 0 - - # 1. Assume cross-account role context natively - sts = boto3.client("sts") - role_arn = f"arn:aws:iam::{account_id}:role/{role_name}" - +console = Console() + +def build_assumed_role_session(target_account_id, role_name, region_name, profile=None): + """Assumes the cross-account role in the target member account and returns an authorized session.""" + base_session = boto3.Session(profile_name=profile) if profile else boto3.Session() + if not role_name: + return base_session + try: - assumed_role = sts.assume_role( + sts_client = base_session.client("sts", region_name=region_name) + partition = "aws" + try: + caller_identity = sts_client.get_caller_identity() + partition = caller_identity["Arn"].split(":")[1] + except Exception: + if region_name.startswith("us-gov-"): + partition = "aws-us-gov" + + role_arn = f"arn:{partition}:iam::{target_account_id}:role/{role_name}" + role_session_name = f"PurgeLane-{target_account_id}" + + assumed_role_object = sts_client.assume_role( RoleArn=role_arn, - RoleSessionName="S3PurgeWorkerSession" + RoleSessionName=role_session_name, + DurationSeconds=3600 ) - creds = assumed_role["Credentials"] - session = boto3.Session( - aws_access_key_id=creds["AccessKeyId"], - aws_secret_access_key=creds["SecretAccessKey"], - aws_session_token=creds["SessionToken"] + credentials = assumed_role_object['Credentials'] + return boto3.Session( + aws_access_key_id=credentials['AccessKeyId'], + aws_secret_access_key=credentials['SecretAccessKey'], + aws_session_token=credentials['SessionToken'], + region_name=region_name ) except Exception as e: - return {"bucket": bucket_name, "status": f"FAILED_ROLE_ASSUME: {e}", "time": 0.0, "count": 0} + raise RuntimeError(f"Cross-Account Delegation Failure for role '{role_name}' on account {target_account_id}: {str(e)}") - s3_resource = session.resource("s3") - bucket = s3_resource.Bucket(bucket_name) +def purge_bucket_worker_task(target_block, role_name, primary_region, dry_run, aws_profile, progress_queue, task_index): + """Worker lifecycle executing object purgation over live queue updates.""" + start_time = time.time() + account_id = target_block["account_id"] + bucket_name = target_block["bucket_name"] + deleted_count = 0 - if dry_run: - # Dry run counts targets safely via low friction collection queries - try: + def update_lane_ui(current_msg, total_weight): + if progress_queue: + progress_queue.put({ + "bucket_name": bucket_name, + "task_index": task_index, + "current_test": current_msg, + "total_tests": total_weight + }) + + try: + update_lane_ui("Assuming_Role", "100") + session = build_assumed_role_session(account_id, role_name, primary_region, aws_profile) + s3_resource = session.resource("s3") + bucket = s3_resource.Bucket(bucket_name) + + if dry_run: + update_lane_ui("Evaluating_Size", "100") + # Safe object counting simulation loop total_targets = sum(1 for _ in bucket.object_versions.all()) + update_lane_ui(f"Dry_Run_Count_{total_targets}", str(total_targets)) elapsed = time.time() - start_time return {"bucket": bucket_name, "status": "DRY_RUN_APPROVED", "time": elapsed, "count": total_targets} - except Exception as e: - return {"bucket": bucket_name, "status": f"DRY_RUN_FAILED: {e}", "time": 0.0, "count": 0} - # 2. Complete Execution Path: Delete both object versions and delete markers - try: - # Standard object batch deletion loop + # Armed Destructive Execution Logic Block + update_lane_ui("Listing_Objects", "Evaluating") + + # Pull down versions collection references version_paginator = bucket.object_versions batch = [] + # First phase evaluation count + update_lane_ui("Purging_Objects_0", "100") + for version in version_paginator.all(): batch.append({"Key": version.object_key, "VersionId": version.id}) deleted_count += 1 - # Flush delete queue in standard 1,000 object batches if len(batch) == 1000: bucket.delete_objects(Delete={"Objects": batch, "Quiet": True}) + update_lane_ui(f"Purging_Objects_{deleted_count}", f"{deleted_count + 1000}") batch = [] if batch: bucket.delete_objects(Delete={"Objects": batch, "Quiet": True}) - + elapsed = time.time() - start_time + update_lane_ui(f"Purged_{deleted_count}", str(deleted_count)) return {"bucket": bucket_name, "status": "PURGED_SUCCESSFULLY", "time": elapsed, "count": deleted_count} - - except Exception as e: + + except Exception as err: elapsed = time.time() - start_time - return {"bucket": bucket_name, "status": f"PURGE_FAILED: {e}", "time": elapsed, "count": deleted_count} + update_lane_ui("PURGE_FAILED", "0") + return {"bucket": bucket_name, "status": f"FAILED: {str(err)}", "time": elapsed, "count": deleted_count} def main(): - parser = argparse.ArgumentParser(description="Multi-Threaded Cross-Account S3 Purge Engine v1.0.0") + global_start = time.time() + parser = argparse.ArgumentParser(description=f"Enterprise Multi-Threaded S3 Purge Engine Framework v{__version__}") + parser.add_argument("--csv", required=True, help="Path to generated s3_global_inventory_*.csv asset file") - parser.add_argument("--filter", required=True, help="Python Regex string matching target bucket names to clear") - parser.add_argument("--role-name", default="SecOpsAuditRole", help="Cross-account assume role identifier map") - parser.add_argument("--workers", type=int, default=10, help="Max concurrent worker threads") - parser.add_argument("--execute", action="store_true", help="Arm execution logic block (default is Dry Run safe mode)") - args = parser.parse_args() + parser.add_argument("--filter-bucket", required=True, help="Regex criteria pattern matching target bucket names") + parser.add_argument("--filter-account", help="Optional regex criteria matching sub-account IDs or environment alias descriptors") + parser.add_argument("--role-name", help="Target IAM cross-account role naming assignment configuration") + parser.add_argument("--workers", type=int, default=4, help="Maximum concurrent worker slots channels") + parser.add_argument("--region", default="us-east-1", help="Execution regional fallback context") + parser.add_argument("--profile", help="AWS configuration context configuration handle target") + parser.add_argument("--execute", action="store_true", help="Arm structural execution destructive deletions") + args = parser.parse_args() dry_run = not args.execute + + # Compile Regular Expression filters try: - compiled_regex = re.compile(args.filter) + bucket_pattern = re.compile(args.filter_bucket, re.IGNORECASE) + account_pattern = re.compile(args.filter_account, re.IGNORECASE) if args.filter_account else None except Exception as re_err: - print(f"[-] Invalid Regex filter pattern specified: {re_err}"); sys.exit(1) + console.print(f"[bold red]Error: Invalid regular expression pattern parameter syntax: {re_err}[/bold red]") + sys.exit(1) targets = [] + + # Process dataset file + if not os.path.exists(args.csv): + console.print(f"[bold red]Error: Target inventory asset file path '{args.csv}' does not exist.[/bold red]") + sys.exit(1) + with open(args.csv, "r", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: b_name = row["Bucket Name"] - if compiled_regex.search(b_name): - targets.append({ - "account_id": row["Account ID"], - "alias": row["Account Alias"], - "bucket_name": b_name - }) - - print(f"[*] Discovery Phase: Located {len(targets)} buckets matching filter template string: '{args.filter}'") - if dry_run: - print("[!] SAFE MODE ACTIVE: Evaluating tracking weights and counts. No deletion operations will occur.") - else: - print("[WARN] ARMED EXECUTION MODE: Destructive data purge sequence initializing across fleet pools.") + acct_id = row["Account ID"] + acct_alias = row["Account Alias"] + # Filter out resource bucket profiles + if not bucket_pattern.search(b_name): + continue + + # Account subset layer pattern evaluations + if account_pattern: + if not (account_pattern.search(acct_id) or account_pattern.search(acct_alias)): + continue + + targets.append({ + "account_id": acct_id, + "alias": acct_alias, + "bucket_name": b_name + }) + + console.print(f"[bold blue]======================================================================[/bold blue]") + console.print(f"Info: Purge Worker Engine Lifecycle v{__version__} Active") + console.print(f"[bold blue]======================================================================[/bold blue]") + console.print(f"Target Bucket Filter Regex : [cyan]{args.filter_bucket}[/cyan]") + console.print(f"Target Account Filter Regex: [cyan]{args.filter_account if args.filter_account else 'None (All Accounts Eligible)'}[/cyan]") + console.print(f"Discovered Filter Targets : [green]{len(targets)} Buckets[/green]") + console.print(f"Safe Execution Dry-Run Mode: [white]{dry_run}[/white]") + console.print(f"[bold blue]----------------------------------------------------------------------[/bold blue]") + + if not targets: + console.print("[bold yellow]Warning: Zero buckets matched filtering configurations criteria rules. Safely exiting.[/bold yellow]") + sys.exit(0) + + progress_queue = queue.Queue() account_timers = defaultdict(float) - global_start = time.time() + summary_results = [] - # 3. Execution Thread Pool Lifecycle - with ThreadPoolExecutor(max_workers=args.workers) as executor: - future_to_bucket = { - executor.submit( - purge_bucket_worker, t["account_id"], t["bucket_name"], args.role_name, dry_run - ): t for t in targets - } - - for future in as_completed(future_to_bucket): - meta = future_to_bucket[future] - try: - res = future.result() - account_timers[meta["alias"]] += res["time"] - - print(f" • [BUCKET LOG] Account: {meta['alias']} | Bucket: {res['bucket']:<40} | Status: {res['status']:<20} | Objects Affected: {res['count']:<6} | Duration: {res['time']:.2f}s") - except Exception as exc: - print(f" • [FATAL THREAD EXCEPTION] Target {meta['bucket_name']} threw unhandled loop block error: {exc}") + # Dynamic Layout Rendering Interface + progress_table = Table.grid(expand=True) + + overall_progress = Progress( + SpinnerColumn(), + TextColumn("[bold yellow]Overall Purge Progress:[/bold yellow]"), + BarColumn(bar_width=40), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + TextColumn("[bold green]Bucket {task.completed} of {task.total}[/bold green]"), + TimeElapsedColumn() + ) + macro_task_id = overall_progress.add_task("Purge Fleet", total=len(targets)) + + lane_progress = Progress( + TextColumn("[bold blue]Lane #{task.fields[lane_idx]}[/bold blue]"), + TextColumn("([bold white]{task.fields[bucket_label]}[/bold white])"), + BarColumn(bar_width=25), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + TextColumn("[bold yellow]{task.fields[current_test]}/{task.fields[total_tests]}[/bold yellow]"), + TimeElapsedColumn() + ) + + render_lanes_count = min(args.workers, len(targets)) + lane_task_ids = {} + for i in range(render_lanes_count): + tid = lane_progress.add_task( + "Awaiting batch assignment", total=100, lane_idx=i, bucket_label="IDLE", current_test="START", total_tests="0" + ) + lane_task_ids[i] = tid + + progress_table.add_row(overall_progress) + progress_table.add_row("[bold blue]----------------------------------------------------------------------[/bold blue]") + progress_table.add_row(lane_progress) + + with Live(progress_table, refresh_per_second=10, console=console): + with ThreadPoolExecutor(max_workers=args.workers) as executor: + futures_map = {} + + for idx, item in enumerate(targets): + lane_idx = idx % args.workers + fut = executor.submit( + purge_bucket_worker_task, + item, args.role_name, args.region, dry_run, args.profile, progress_queue, lane_idx + ) + futures_map[fut] = (item, lane_idx) + + while len(futures_map) > 0: + while not progress_queue.empty(): + try: + signal = progress_queue.get_nowait() + lane_id = signal["task_index"] + t_id = lane_task_ids[lane_id] + + found_digits = re.findall(r'\d+', str(signal["current_test"])) + current_val = int(found_digits[0]) if found_digits else 1 + total_w = int(signal["total_tests"]) if str(signal["total_tests"]).isdigit() else (current_val + 1000) + + pct = int((current_val / total_w) * 100) if total_w > 0 else 0 + if pct > 100: pct = 100 + + # Squeeze long bucket designations to make sure text columns do not clip terminal rows + b_display = signal["bucket_name"] + if len(b_display) > 30: + b_display = f"...{b_display[-27:]}" + + lane_progress.update( + t_id, bucket_label=b_display, completed=pct, current_test=signal["current_test"], total_tests=signal["total_tests"] + ) + except queue.Empty: + break + + completed_futures = [f for f in futures_map if f.done()] + for fut in completed_futures: + meta, assigned_lane = futures_map.pop(fut) + try: + res = fut.result() + account_timers[meta["alias"]] += res["time"] + summary_results.append(res) + except Exception as fatal_fut_err: + summary_results.append({"bucket": meta["bucket_name"], "status": f"CRITICAL_EXCEPTION: {fatal_fut_err}", "time": 0.0, "count": 0}) + + overall_progress.update(macro_task_id, advance=1) + lane_progress.update( + lane_task_ids[assigned_lane], bucket_label="IDLE", completed=0, current_test="IDLE", total_tests="0" + ) + time.sleep(0.05) global_duration = time.time() - global_start - print("\n" + "=" * 80) - print(f"[*] PURGE SEQUENCE SUMMARY COMPLETED | Total Runtime Duration: {global_duration:.2f}s") - print("=" * 80) - for acct, total_time in account_timers.items(): - print(f" • Account Module Alias Focus Group: {acct:<25} | Cumulative Active Worker Duration: {total_time:.2f}s") - print("=" * 80 + "\n") + console.print(f"\n[bold blue]======================================================================[/bold blue]") + console.print(f"[*] PURGE SEQUENCE LIFECYCLE CONCLUDED | Total Run Time: {global_duration:.2f}s") + console.print(f"[bold blue]======================================================================[/bold blue]") + + # Calculate performance times grouped per target profile sector + for alias_name, duration_seconds in sorted(account_timers.items()): + console.print(f" • Account Sector Group Focus: [white]{alias_name:<25}[/white] | Cumulative Active Operations Duration: [green]{duration_seconds:.2f}s[/green]") + console.print(f"[bold blue]======================================================================[/bold blue]\n") if __name__ == "__main__": main() \ No newline at end of file