Automating AWS EC2 Instance Management with Python

Managing AWS EC2 instances can often require repetitive tasks such as listing available instances, starting stopped ones, and securely logging into them. Automating these tasks not only saves time but also reduces the possibility of human error. In this blog post, we'll explore how to use Python and Boto3, AWS's SDK for Python, to create a simple yet powerful script that lists your EC2 instances, powers on any that are off, and then securely logs into them.

Getting Started

Before diving into the code, ensure you have Python and Boto3 installed in your environment. Boto3 allows you to directly interact with AWS services in Python. If you haven't installed Boto3 yet, you can do so by running pip install boto3.

The Script Explained

Our script performs several key functions to manage EC2 instances efficiently. Let's break down each part of the script:

Initialization

First, we initialize a Boto3 EC2 resource and a client. These objects allow us to interact with the EC2 instances and perform actions like starting or stopping them.

import boto3

ec2 = boto3.resource('ec2')
ec2_client = boto3.client('ec2')

Listing Instances

The list_all_instances function lists all EC2 instances based on specific filters. It can filter instances by name and whether they are running or stopped. This makes it easy to identify which instances are available for connection.

Selecting an Instance

The select_instance function prompts the user to select an instance from the list of available ones. This interaction ensures that users have control over which instance they're managing.

Handling SSH Keys

To securely log into an instance, the correct SSH key is required. The find_key_for_instance function locates the SSH key associated with the instance, ensuring a secure connection.

Starting Stopped Instances

One of the script's key features is its ability to start instances that are found to be stopped. This is handled in the ssh_into_instance function, which checks the instance's state before attempting to log in. If the instance is stopped, the script starts it and waits until it's in the 'running' state.

Securely Logging Into Instances

Finally, once the instance is ready, the script logs into it using SSH. It constructs the SSH command with the correct parameters, including the path to the SSH key, and executes it.

Running the Script

To run the script, simply execute it in your terminal. You'll be prompted to decide whether to include stopped instances in the list and to optionally filter instances by name. Then, select the instance you wish to log into from the provided list. If the selected instance is stopped, the script will start it and then automatically log you in.

Conclusion

This script is a starting point for automating EC2 instance management. By leveraging Python and Boto3, you can customize and extend this script to fit your specific needs, such as automating backups, deploying applications, or managing instance lifecycles. Automating these tasks not only saves time but also ensures that your operations are more secure and efficient.

Remember, automation is key to effective cloud management. By incorporating scripts like this into your workflow, you can focus on more important tasks while letting your code handle the routine work.

import boto3
import subprocess
import os
import time

# Initialize a boto3 EC2 resource
ec2 = boto3.resource('ec2')
ec2_client = boto3.client('ec2')  # Added for starting instances

def list_all_instances(include_stopped=False, search_term=None):
    """List all EC2 instances, optionally excluding stopped instances and filtering by search term."""
    filters = [{'Name': 'tag:Name', 'Values': ['*'+search_term+'*']} if search_term else {'Name': 'instance-state-name', 'Values': ['running', 'stopped']}]

    if not include_stopped:
        filters.append({'Name': 'instance-state-name', 'Values': ['running']})

    instances = ec2.instances.filter(Filters=filters)
    return instances

def get_instance_name(instance):
    """Extract the name of the instance from its tags."""
    for tag in instance.tags or []:
        if tag['Key'] == 'Name':
            return tag['Value']
    return "No Name"

def select_instance(instances):
    """Allow the user to select an instance to log into."""
    print("Available instances:")
    if not instances:
        print("No matching instances found.")
        return None

    for i, instance in enumerate(instances, start=1):
        name = get_instance_name(instance)
        print(f"{i}) Name: {name}, Instance ID: {instance.id}, State: {instance.state['Name']}")

    selection = input("Enter the number of the instance you want to log into (or 'exit' to quit): ")
    if selection.lower() == 'exit':
        return None
    try:
        selection = int(selection) - 1
        return list(instances)[selection]
    except (ValueError, IndexError):
        print("Invalid selection.")
        return None

def find_key_for_instance(instance):
    """Find the SSH key for the instance based on its KeyName."""
    key_name = instance.key_name
    keys_directory = os.path.expanduser("~/.ssh")
    for key_file in os.listdir(keys_directory):
        if key_file.startswith(key_name) and key_file.endswith(".pem"):
            return os.path.join(keys_directory, key_file)
    return None

def ssh_into_instance(instance, remote_user="ec2-user"):
    """SSH into the selected instance, if any."""
    if instance is None:
        return

    # Check instance state and start if stopped
    if instance.state['Name'] == 'stopped':
        print(f"Instance {instance.id} is stopped. Starting instance...")
        ec2_client.start_instances(InstanceIds=[instance.id])
        print("Waiting for instance to be in 'running' state...")
        instance.wait_until_running()
        instance.reload()  # Refresh instance state and details

    ssh_key_path = find_key_for_instance(instance)
    if not ssh_key_path:
        print(f"No matching SSH key found for instance {instance.id} with KeyName {instance.key_name}")
        return

    print(f"Logging into {get_instance_name(instance)} ({instance.id})...")
    private_ip = instance.private_ip_address
    ssh_cmd = f'ssh -o StrictHostKeyChecking=no -i {ssh_key_path} {remote_user}@{private_ip}'
    subprocess.run(ssh_cmd, shell=True)

def main():
    """Main function to list instances and allow user selection for SSH login."""
    include_stopped = input("Include stopped instances? (yes/no): ").lower().startswith('y')
    search_term = input("Enter a search term to filter by instance name (leave empty for no filter): ").strip() or None
    instances = list(list_all_instances(include_stopped, search_term))

    selected_instance = select_instance(instances)
    ssh_into_instance(selected_instance)

if __name__ == "__main__":
    main()