Skip to content

Commit

Permalink
interrupt handling and stats based summary
Browse files Browse the repository at this point in the history
  • Loading branch information
@morga471
morga471 committed Mar 14, 2025
1 parent 8e797d9 commit 0ccd2eb
Showing 1 changed file with 196 additions and 56 deletions.
252 changes: 196 additions & 56 deletions local-app/python-tools/gfl-resource-actions/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,166 @@ def show_detailed_help():
"""
print(help_text)

def display_summary(stats, action, was_interrupted=False):
"""Display a summary of the actions taken"""
print("\n" + "="*50)
title = f"SUMMARY OF {action.upper()} OPERATION"
if was_interrupted:
title += " (INTERRUPTED)"
print(title)
print("="*50)

print(f"\nAccounts processed: {stats['accounts_processed']}")
print(f"Accounts skipped: {stats['accounts_skipped']}")

# EC2 summary
print("\nEC2 Instances:")
print(f" Found: {stats['ec2_found']}")
print(f" Skipped: {stats['ec2_skipped']}")
print(f" {action.capitalize()}ed: {stats['ec2_' + action + 'ed']}")
print(f" Errors: {stats['ec2_errors']}")

# RDS summary
print("\nRDS Instances:")
print(f" Found: {stats['rds_found']}")
print(f" Skipped: {stats['rds_skipped']}")
print(f" {action.capitalize()}ed: {stats['rds_' + action + 'ed']}")
print(f" Errors: {stats['rds_errors']}")

# EKS summary
print("\nEKS Clusters:")
print(f" Found: {stats['eks_found']}")
print(f" Skipped: {stats['eks_skipped']}")
print(f" {action.capitalize()}ed: {stats['eks_' + action + 'ed']}")
print(f" Errors: {stats['eks_errors']}")

# EMR summary
print("\nEMR Clusters:")
print(f" Found: {stats['emr_found']}")
print(f" Skipped: {stats['emr_skipped']}")
print(f" {action.capitalize()}ed: {stats['emr_' + action + 'ed']}")
print(f" Errors: {stats['emr_errors']}")

print("\n" + "="*50)

def process_account(account, credentials, regions, action, dry_run, exclude_resources, stats):
"""Process a single account and update stats"""
logger.info(f"Processing account: {account['name']} ({account['id']})")

# Get resources
resources = get_account_resources(credentials, regions, exclude_resources)

# Count service-managed EC2 instances
eks_managed = sum(1 for i in resources['ec2_instances'] if config.EKS_TAG in i.get('tags', {}))
emr_managed = sum(1 for i in resources['ec2_instances'] if config.EMR_TAG in i.get('tags', {}))

# Update resource counts
total_ec2 = len(resources['ec2_instances'])
excluded_ec2 = sum(1 for i in resources['ec2_instances'] if config.EXCLUSION_TAG in i.get('tags', {}))
stats['ec2_found'] += total_ec2
stats['ec2_skipped'] += excluded_ec2

stats['rds_found'] += len(resources['rds_instances'])
stats['eks_found'] += len(resources['eks_clusters'])
stats['emr_found'] += len(resources['emr_clusters'])

# Log discovered resources with service-managed counts
logger.info(f"Account {account['id']} - Found {total_ec2} EC2 instances ({excluded_ec2} explicitly excluded, {eks_managed} EKS-managed, {emr_managed} EMR-managed)")
logger.info(f"Account {account['id']} - Found {len(resources['rds_instances'])} RDS instances")
logger.info(f"Account {account['id']} - Found {len(resources['eks_clusters'])} EKS clusters")
logger.info(f"Account {account['id']} - Found {len(resources['emr_clusters'])} EMR clusters")

# Initialize resource managers
ec2_manager = EC2Manager(credentials, account['id'])
rds_manager = RDSManager(credentials, account['id'])
eks_manager = EKSManager(credentials, account['id'])
emr_manager = EMRManager(credentials, account['id'])

# Perform actions based on selected mode
if action == "stop":
if "ec2" not in exclude_resources:
result = ec2_manager.stop(resources['ec2_instances'], dry_run)
stats['ec2_stopped'] += result.get('success', 0)
stats['ec2_errors'] += result.get('errors', 0)
else:
stats['ec2_skipped'] += total_ec2 - excluded_ec2

if "rds" not in exclude_resources:
result = rds_manager.stop(resources['rds_instances'], dry_run)
stats['rds_stopped'] += result.get('success', 0)
stats['rds_errors'] += result.get('errors', 0)
else:
stats['rds_skipped'] += len(resources['rds_instances'])

if "eks" not in exclude_resources:
result = eks_manager.stop(resources['eks_clusters'], dry_run)
stats['eks_stopped'] += result.get('success', 0)
stats['eks_errors'] += result.get('errors', 0)
else:
stats['eks_skipped'] += len(resources['eks_clusters'])

if "emr" not in exclude_resources:
result = emr_manager.stop(resources['emr_clusters'], dry_run)
stats['emr_stopped'] += result.get('success', 0)
stats['emr_errors'] += result.get('errors', 0)
else:
stats['emr_skipped'] += len(resources['emr_clusters'])
else: # action == "start"
if "ec2" not in exclude_resources:
result = ec2_manager.start(resources['ec2_instances'], dry_run)
stats['ec2_started'] += result.get('success', 0)
stats['ec2_errors'] += result.get('errors', 0)
else:
stats['ec2_skipped'] += total_ec2 - excluded_ec2

if "rds" not in exclude_resources:
result = rds_manager.start(resources['rds_instances'], dry_run)
stats['rds_started'] += result.get('success', 0)
stats['rds_errors'] += result.get('errors', 0)
else:
stats['rds_skipped'] += len(resources['rds_instances'])

if "eks" not in exclude_resources:
result = eks_manager.start(resources['eks_clusters'], dry_run)
stats['eks_started'] += result.get('success', 0)
stats['eks_errors'] += result.get('errors', 0)
else:
stats['eks_skipped'] += len(resources['eks_clusters'])

if "emr" not in exclude_resources:
result = emr_manager.start(resources['emr_clusters'], dry_run)
stats['emr_started'] += result.get('success', 0)
stats['emr_errors'] += result.get('errors', 0)
else:
stats['emr_skipped'] += len(resources['emr_clusters'])

def initialize_stats():
"""Initialize the statistics tracking dictionary"""
return {
'accounts_processed': 0,
'accounts_skipped': 0,
'ec2_found': 0,
'ec2_skipped': 0,
'ec2_stopped': 0,
'ec2_started': 0,
'ec2_errors': 0,
'rds_found': 0,
'rds_skipped': 0,
'rds_stopped': 0,
'rds_started': 0,
'rds_errors': 0,
'eks_found': 0,
'eks_skipped': 0,
'eks_stopped': 0,
'eks_started': 0,
'eks_errors': 0,
'emr_found': 0,
'emr_skipped': 0,
'emr_stopped': 0,
'emr_started': 0,
'emr_errors': 0,
}

def main():
"""Main execution function."""
args = parse_arguments()
Expand All @@ -83,70 +243,50 @@ def main():
exclude_accounts = args.exclude_accounts.split(',') if args.exclude_accounts else []
exclude_resources = args.exclude_resources or []

# Initialize statistics dictionary
stats = initialize_stats()

logger.info(f"Starting AWS organization resource {action} {'(DRY RUN)' if dry_run else ''}")

# Get organization accounts
accounts = get_organization_accounts(exclude_accounts)

# Process each account
for account in accounts:
logger.info(f"Processing account: {account['name']} ({account['id']})")

# Assume role in the account
credentials = assume_role(account['id'])
if not credentials:
logger.warning(f"Skipping account {account['id']} due to role assumption failure")
continue

# Get resources
resources = get_account_resources(credentials, regions, exclude_resources)

# Count service-managed EC2 instances
eks_managed = sum(1 for i in resources['ec2_instances'] if config.EKS_TAG in i.get('tags', {}))
emr_managed = sum(1 for i in resources['ec2_instances'] if config.EMR_TAG in i.get('tags', {}))

# Log discovered resources with service-managed counts
total_ec2 = len(resources['ec2_instances'])
excluded_ec2 = sum(1 for i in resources['ec2_instances'] if config.EXCLUSION_TAG in i.get('tags', {}))

logger.info(f"Account {account['id']} - Found {total_ec2} EC2 instances ({excluded_ec2} explicitly excluded, {eks_managed} EKS-managed, {emr_managed} EMR-managed)")
logger.info(f"Account {account['id']} - Found {len(resources['rds_instances'])} RDS instances")
logger.info(f"Account {account['id']} - Found {len(resources['eks_clusters'])} EKS clusters")
logger.info(f"Account {account['id']} - Found {len(resources['emr_clusters'])} EMR clusters")

# Initialize resource managers
ec2_manager = EC2Manager(credentials, account['id'])
rds_manager = RDSManager(credentials, account['id'])
eks_manager = EKSManager(credentials, account['id'])
emr_manager = EMRManager(credentials, account['id'])

# Perform actions based on selected mode
if action == "stop":
if "ec2" not in exclude_resources:
ec2_manager.stop(resources['ec2_instances'], dry_run)

if "rds" not in exclude_resources:
rds_manager.stop(resources['rds_instances'], dry_run)

if "eks" not in exclude_resources:
eks_manager.stop(resources['eks_clusters'], dry_run)
# Flag to track interruptions
was_interrupted = False

try:
# Process each account
for account in accounts:
try:
# Assume role in the account
credentials = assume_role(account['id'])
if not credentials:
logger.warning(f"Skipping account {account['id']} due to role assumption failure")
stats['accounts_skipped'] += 1
continue

if "emr" not in exclude_resources:
emr_manager.stop(resources['emr_clusters'], dry_run)
else: # action == "start"
if "ec2" not in exclude_resources:
ec2_manager.start(resources['ec2_instances'], dry_run)

if "rds" not in exclude_resources:
rds_manager.start(resources['rds_instances'], dry_run)

if "eks" not in exclude_resources:
eks_manager.start(resources['eks_clusters'], dry_run)
stats['accounts_processed'] += 1

if "emr" not in exclude_resources:
emr_manager.start(resources['emr_clusters'], dry_run)
# Process this account
process_account(account, credentials, regions, action, dry_run, exclude_resources, stats)

except Exception as e:
logger.error(f"Error processing account {account['id']}: {str(e)}")
stats['accounts_skipped'] += 1

logger.info(f"Resource {action} completed {'(DRY RUN)' if dry_run else ''}")

except KeyboardInterrupt:
was_interrupted = True
print("\n\nProcess interrupted by user (Ctrl+C)")
logger.warning("Process was interrupted by user (Ctrl+C)")

# Display summary of actions (will run even if interrupted)
display_summary(stats, action, was_interrupted)

logger.info(f"Resource {action} completed {'(DRY RUN)' if dry_run else ''}")
# Exit with appropriate code
if was_interrupted:
sys.exit(130) # Standard exit code for Ctrl+C/SIGINT

if __name__ == "__main__":
main()

0 comments on commit 0ccd2eb

Please sign in to comment.