rework

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