Monday, 1 January 2024

Transfer file from EC2 to local using SSM

 import boto3

import base64

import os

import time


def get_instance_id_by_private_ip(private_ip, region="us-east-1", aws_profile=None):

    # Set the AWS profile, if specified

    if aws_profile:

        boto3.setup_default_session(profile_name=aws_profile)


    # Initialize EC2 client

    ec2_client = boto3.client("ec2", region_name=region)

    

    # Describe instances filtered by private IP

    response = ec2_client.describe_instances(

        Filters=[{"Name": "private-ip-address", "Values": [private_ip]}]

    )

    

    # Extract the instance ID from the response

    for reservation in response["Reservations"]:

        for instance in reservation["Instances"]:

            return instance["InstanceId"]

    

    # If no instance found, return None

    return None


def copy_files_from_ec2(file_paths_on_ec2, local_save_dir, instance_id, region="us-east-1", aws_profile=None):

    # Set the AWS profile, if specified

    if aws_profile:

        boto3.setup_default_session(profile_name=aws_profile)


    # Initialize SSM client

    ssm_client = boto3.client("ssm", region_name=region)


    for file_path_on_ec2 in file_paths_on_ec2:

        # Extract filename to save locally

        filename = os.path.basename(file_path_on_ec2)

        local_save_path = os.path.join(local_save_dir, filename)


        # Base64 encode the file on the EC2 instance using an SSM command

        commands = [f'base64 {file_path_on_ec2}']

        

        # Send the command to EC2 via SSM

        response = ssm_client.send_command(

            DocumentName="AWS-RunShellScript",

            Parameters={"commands": commands},

            InstanceIds=[instance_id],

        )

        

        # Fetch the command ID to track the output

        command_id = response["Command"]["CommandId"]


        # Wait for the command to complete

        time.sleep(2)  # Initial wait before polling for command status

        while True:

            result = ssm_client.get_command_invocation(

                CommandId=command_id,

                InstanceId=instance_id,

            )

            if result["Status"] == "Success":

                # Retrieve the Base64 output

                encoded_content = result["StandardOutputContent"]

                break

            elif result["Status"] in ["Failed", "Cancelled", "TimedOut"]:

                raise Exception(f"SSM command failed for '{file_path_on_ec2}' with status: {result['Status']}")

            time.sleep(1)  # Poll every second


        # Decode the Base64 content and write it to a local file

        decoded_content = base64.b64decode(encoded_content)

        with open(local_save_path, "wb") as file:

            file.write(decoded_content)


        print(f"File '{file_path_on_ec2}' copied from instance '{instance_id}' to local path '{local_save_path}'.")


# Example usage


# Retrieve the instance ID using a private IP

private_ip = "10.0.0.1"  # Replace with the private IP address you're looking for

region = "us-east-1"  # Replace with your AWS region if different

aws_profile = "your-profile-name"  # Replace with your AWS CLI profile name


instance_id = get_instance_id_by_private_ip(private_ip, region, aws_profile)

if instance_id:

    print(f"The instance ID for private IP {private_ip} is {instance_id}.")


    # Specify files to copy from EC2 and the local save directory

    file_paths_on_ec2 = ["/home/ec2-user/test1.py", "/home/ec2-user/test2.py"]  # Paths to files on EC2 instance

    local_save_dir = "./"  # Local directory to save the files


    # Copy files from EC2 to local machine

    copy_files_from_ec2(file_paths_on_ec2, local_save_dir, instance_id, region, aws_profile)

else:

    print(f"No instance found with private IP {private_ip}.")


No comments:

Post a Comment