Skip to content

Commit

Permalink
update queue iterator (as .new so as not to break running thing)
Browse files Browse the repository at this point in the history
  • Loading branch information
badra001 committed Jul 17, 2026
1 parent 4d8230b commit 7a1bf3c
Showing 1 changed file with 384 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,384 @@
#!/usr/bin/env python3
import os
import sys
import re
import csv
import json
import queue
import time
import argparse
from datetime import datetime, timezone
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

# --- VERSIONING ---
__version__ = "2.6.0"

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:
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=role_session_name,
DurationSeconds=3600
)
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:
raise RuntimeError(f"Cross-Account Delegation Failure for role '{role_name}' on account {target_account_id}: {str(e)}")

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 with fine-grained tracking metrics."""
time_format = "%Y-%m-%dT%H:%M:%S.%f"
start_epoch = time.time()
start_timestamp_str = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"

account_id = target_block["account_id"]
account_name = target_block["alias"]
bucket_name = target_block["bucket_name"]
source_size = target_block["source_size"]
source_objects = target_block["source_objects"]

deleted_count = 0
error_count = 0
status_message = "INITIALIZED"

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")
total_targets = sum(1 for _ in bucket.object_versions.all())
update_lane_ui(f"Dry_Run_Count_{total_targets}", str(total_targets))

end_epoch = time.time()
end_timestamp_str = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"
return {
"account_id": account_id, "account_name": account_name, "region": primary_region,
"bucket_name": bucket_name, "source_size": source_size, "object_count": source_objects,
"start_time": start_timestamp_str, "end_time": end_timestamp_str,
"elapsed_time_seconds": round(end_epoch - start_epoch, 3), "error_count": 0, "status": "DRY_RUN_APPROVED"
}

update_lane_ui("Listing_Objects", "Evaluating")
version_paginator = bucket.object_versions
batch = []

for version in version_paginator.all():
try:
batch.append({"Key": version.object_key, "VersionId": version.id})
deleted_count += 1

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 = []
except Exception:
error_count += 1

if batch:
try:
bucket.delete_objects(Delete={"Objects": batch, "Quiet": True})
except Exception:
error_count += 1

status_message = "PURGED_SUCCESSFULLY" if error_count == 0 else "PURGED_WITH_ERRORS"

except Exception as err:
error_count += 1
status_message = f"FAILED: {str(err)}"
update_lane_ui("PURGE_FAILED", "0")

end_epoch = time.time()
end_timestamp_str = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"

return {
"account_id": account_id, "account_name": account_name, "region": primary_region,
"bucket_name": bucket_name, "source_size": source_size, "object_count": source_objects,
"start_time": start_timestamp_str, "end_time": end_timestamp_str,
"elapsed_time_seconds": round(end_epoch - start_epoch, 3), "error_count": error_count, "status": status_message
}

def main():
global_start_epoch = time.time()
time_format = "%Y-%m-%dT%H:%M:%S.%f"
global_start_timestamp = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"

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-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("--output", default=".", help="Folder destination directory path to save logs (Default: current directory)")
parser.add_argument("--execute", action="store_true", help="Arm structural execution destructive deletions")

args = parser.parse_args()
dry_run = not args.execute

try:
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:
console.print(f"[bold red]Error: Invalid regular expression pattern parameter syntax: {re_err}[/bold red]")
sys.exit(1)

targets = []
detected_size_unit_label = "Size"

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:
clean_row = {str(k).strip(): str(v).strip() for k, v in row.items() if k is not None}

b_name = clean_row.get("Bucket Name")
acct_id = clean_row.get("Account ID")
acct_alias = clean_row.get("Account Alias")
b_region = clean_row.get("Region", args.region)

if not b_name: continue

raw_size = 0.0
raw_objects = 0
for key_str, val_str in clean_row.items():
normalized_key = key_str.lower()
clean_val = val_str.replace(",", "")

if "size" in normalized_key:
detected_size_unit_label = key_str
try: raw_size = float(clean_val)
except Exception: pass
if "object" in normalized_key:
try: raw_objects = int(float(clean_val))
except Exception: pass

if not bucket_pattern.search(b_name):
continue

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, "region": b_region,
"source_size": raw_size, "source_objects": raw_objects
})

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"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)
operation_tracking_results = []

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_name]}[/bold white])"),
BarColumn(bar_width=30),
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_name="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)

# Convert the static targets dataset list into a fluid, stateful iterator pool
targets_iterator = iter(targets)
active_lane_allocations = {i: None for i in range(args.workers)}
futures_map = {}

with Live(progress_table, refresh_per_second=10, console=console):
with ThreadPoolExecutor(max_workers=args.workers) as executor:

# Keep looping as long as we have jobs left in the pool OR threads actively executing work
while len(futures_map) > 0 or any(v is not None for v in active_lane_allocations.values()):

# --- FIXED: Dynamically allocate upcoming buckets to open lanes immediately ---
for lane_idx, active_fut in active_lane_allocations.items():
if active_fut is None:
try:
next_item = next(targets_iterator)
fut = executor.submit(
purge_bucket_worker_task,
next_item, args.role_name, next_item["region"], dry_run, args.profile, progress_queue, lane_idx
)
futures_map[fut] = next_item
active_lane_allocations[lane_idx] = fut
except StopIteration:
break # Inventory array completely exhausted

# Consume real-time progress events
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

lane_progress.update(t_id, bucket_name=signal["bucket_name"], completed=pct, current_test=signal["current_test"], total_tests=signal["total_tests"])
except queue.Empty:
break

# --- FIXED: Process completed threads and release their lease slots immediately ---
completed_futures = [f for f in futures_map if f.done()]
for fut in completed_futures:
meta = futures_map.pop(fut)

# Identify and wipe the lease mapping to open the lane back up
assigned_lane = next(k for k, v in active_lane_allocations.items() if v == fut)
active_lane_allocations[assigned_lane] = None

try:
res = fut.result()
account_timers[meta["alias"]] += res["elapsed_time_seconds"]
operation_tracking_results.append(res)
except Exception as fatal_fut_err:
time_now_str = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"
operation_tracking_results.append({
"account_id": meta["account_id"], "account_name": meta["alias"], "region": meta["region"],
"bucket_name": meta["bucket_name"], "source_size": meta["source_size"], "object_count": meta["source_objects"],
"start_time": time_now_str, "end_time": time_now_str, "elapsed_time_seconds": 0.0,
"error_count": 1, "status": f"CRITICAL_EXCEPTION: {str(fatal_fut_err)}"
})

overall_progress.update(macro_task_id, advance=1)
lane_progress.update(lane_task_ids[assigned_lane], bucket_name="IDLE", completed=0, current_test="IDLE", total_tests="0")

time.sleep(0.05)

global_end_epoch = time.time()
global_end_timestamp = datetime.now(timezone.utc).strftime(time_format)[:-3] + "Z"
global_total_elapsed = round(global_end_epoch - global_start_epoch, 3)

if args.output != "." and not os.path.exists(args.output):
os.makedirs(args.output, exist_ok=True)

timestamp_suffix = datetime.now().strftime("%Y%m%dT%H%M%S")
csv_out_path = os.path.join(args.output, f"s3_purge_summary_{timestamp_suffix}.csv")
json_out_path = os.path.join(args.output, f"s3_purge_summary_{timestamp_suffix}.json")

csv_size_header = f"Source {detected_size_unit_label}"

# Export CSV Log
csv_fields = ["Account ID", "Account Name", "Region", "Bucket Name", csv_size_header, "Source Object Count", "Operation Start", "Operation End", "Elapsed Time (s)", "Error Count", "Status"]
with open(csv_out_path, "w", newline="", encoding="utf-8") as cf:
writer = csv.writer(cf)
writer.writerow(csv_fields)
for r in operation_tracking_results:
writer.writerow([
r["account_id"], r["account_name"], r["region"], r["bucket_name"],
r["source_size"], r["object_count"], r["start_time"], r["end_time"],
r["elapsed_time_seconds"], r["error_count"], r["status"]
])

# Export Master JSON Document with Global Engine Telemetry
json_telemetry_payload = {
"orchestrator_metadata": {
"engine_version": __version__,
"command_args": sys.argv,
"global_start_time": global_start_timestamp,
"global_end_time": global_end_timestamp,
"global_elapsed_seconds": global_total_elapsed,
"dry_run_active": dry_run,
"detected_size_unit": detected_size_unit_label,
"filters_applied": {
"bucket_regex": args.filter_bucket,
"account_regex": args.filter_account if args.filter_account else "all"
}
},
"operation_results": operation_tracking_results
}
with open(json_out_path, "w", encoding="utf-8") as jf:
json.dump(json_telemetry_payload, jf, indent=2)

console.print(f"\n[bold blue]======================================================================[/bold blue]")
console.print(f"[*] PURGE SEQUENCE LIFECYCLE CONCLUDED | Total Run Time: {global_total_elapsed:.2f}s")
console.print(f"[bold blue]======================================================================[/bold blue]")
console.print(f"Artifact CSV Export : [bold white]{csv_out_path}[/bold white]")
console.print(f"Artifact JSON Export : [bold white]{json_out_path}[/bold white]")
console.print(f"[bold blue]======================================================================[/bold blue]\n")

if __name__ == "__main__":
main()

0 comments on commit 7a1bf3c

Please sign in to comment.