Helper script to quickly restart all of the instances of an app running in AWS under a load balancer

aws-restarter.py

'''
Helper script to quickly restart all of the instances of a specific app
Assumes lots of things about the setup of the app (load balancer, ssh access, supervisor, etc.)
'''
import boto3
import subprocess
import argparse


def load_balancer_instances(name):
    current_elb = elb.describe_load_balancers(LoadBalancerNames=[name])
    return [i['InstanceId'] for i in current_elb['LoadBalancerDescriptions'][0]['Instances']]

def private_ip(instance_id):
    return ec2.Instance(instance_id).private_ip_address

# Expects that the service is managed by supervisor
def restart_service(service_name, ip_address):
    subprocess.run(["ssh", ip_address, "supervisorctl", "restart", servicename])

def main(service_name, load_balancer_name):
    print("Restarting {} service via {} load balancer".format(service_name, load_balancer_name))
    instance_ids = load_balancer_instances(load_balancer_name)
    print("Found the following associated instances: {}".format(instance_ids))
    ips = [private_ip(i) for i in instance_ids]
    for ip in ips:
        print("Restarting {} box with IP {}".format(service_name, ip))
        restart_service(service_name, ip)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Restarts services by sshing into all of their instances and performing a supervisor restart')
    parser.add_argument('--aws-profile', help='AWS profile to use to perform the request', type=str, default='default')
    parser.add_argument('--service', help='Service name (used by supervisor)', type=str, required=True)
    parser.add_argument('--load-balancer', help="Which load balancer's instances should be restarted", type=str, required=True)
    args = parser.parse_args()

    session = boto3.Session(profile_name=args.aws_profile)
    elb = session.client('elb')
    ec2 = session.resource('ec2')

    main(args.service, args.load_balancer)