ecr_image_update

ecr_image_update.py

#!/usr/bin/env python3
"""
ECR Image Updater Script
Pulls Docker images from AWS ECR, adds pip install requests, and pushes with new tag
"""

import subprocess
import sys
import re
import tempfile
import os
from pathlib import Path


# List of ECR images to pull
IMAGE_LIST = [
    "905418109013.dkr.ecr.ap-south-1.amazonaws.com/local/hello-world:python-test"
]

# New tag suffix to append to modified images
NEW_TAG_SUFFIX = "-with-requests"


def get_ecr_login_command(region: str, registry: str) -> list:
    """Generate AWS ECR login command."""
    return [
        "aws", "ecr", "get-login-password",
        "--region", region
    ]


def docker_login(registry: str, password: str) -> bool:
    """Login to Docker registry."""
    try:
        result = subprocess.run(
            f"docker login --username AWS --password-stdin {registry}",
            input=password,
            capture_output=True,
            text=True,
            shell=True
        )
        if result.returncode == 0:
            print(f"✓ Successfully logged in to {registry}")
            return True
        else:
            print(f"✗ Failed to login to {registry}: {result.stderr}")
            return False
    except Exception as e:
        print(f"✗ Error during docker login: {e}")
        return False


def extract_registry_info(image_uri: str) -> tuple:
    """Extract registry URL and region from ECR image URI."""
    # Pattern: <account_id>.dkr.ecr.<region>.amazonaws.com/<repo>:<tag>
    pattern = r"(\d+\.dkr\.ecr\.([a-z0-9-]+)\.amazonaws\.com)"
    match = re.search(pattern, image_uri)
    if match:
        return match.group(1), match.group(2)
    return None, None


def get_ecr_password(region: str) -> str:
    """Get ECR login password using AWS CLI."""
    try:
        result = subprocess.run(
            f"aws ecr get-login-password --region {region}",
            capture_output=True,
            text=True,
            shell=True
        )
        if result.returncode == 0:
            return result.stdout.strip()
        else:
            print(f"✗ Failed to get ECR password: {result.stderr}")
            return None
    except Exception as e:
        print(f"✗ Error getting ECR password: {e}")
        return None


def pull_image(image_uri: str) -> bool:
    """Pull a Docker image."""
    try:
        print(f"Pulling image: {image_uri}")
        result = subprocess.run(
            f"docker pull {image_uri}",
            capture_output=True,
            text=True,
            shell=True
        )
        if result.returncode == 0:
            print(f"✓ Successfully pulled {image_uri}")
            return True
        else:
            print(f"✗ Failed to pull {image_uri}: {result.stderr}")
            return False
    except Exception as e:
        print(f"✗ Error pulling image: {e}")
        return False


def read_image_list(file_path: str) -> list:
    """Read image URIs from a file (optional, for loading from external file)."""
    images = []
    path = Path(file_path)
    
    if not path.exists():
        print(f"✗ Image list file not found: {file_path}")
        return images
    
    with open(path, 'r') as f:
        for line in f:
            line = line.strip()
            # Skip empty lines and comments
            if line and not line.startswith('#'):
                images.append(line)
    
    return images


def get_new_tag(image_uri: str) -> str:
    """Generate new image URI with modified tag."""
    if ':' in image_uri:
        base, tag = image_uri.rsplit(':', 1)
        return f"{base}:{tag}{NEW_TAG_SUFFIX}"
    else:
        return f"{image_uri}:latest{NEW_TAG_SUFFIX}"


def build_modified_image(base_image: str, new_image: str) -> bool:
    """Build a new image with pip install requests added."""
    try:
        # Create a temporary directory for Dockerfile
        with tempfile.TemporaryDirectory() as tmpdir:
            dockerfile_path = os.path.join(tmpdir, "Dockerfile")
            
            # Create Dockerfile
            dockerfile_content = f"""FROM {base_image}
RUN pip install requests
"""
            with open(dockerfile_path, 'w') as f:
                f.write(dockerfile_content)
            
            print(f"Building modified image: {new_image}")
            print(f"  Adding: RUN pip install requests")
            
            # Build the image
            result = subprocess.run(
                f'docker build -t "{new_image}" "{tmpdir}"',
                capture_output=True,
                text=True,
                shell=True
            )
            
            if result.returncode == 0:
                print(f"✓ Successfully built {new_image}")
                return True
            else:
                print(f"✗ Failed to build image: {result.stderr}")
                return False
    except Exception as e:
        print(f"✗ Error building image: {e}")
        return False


def push_image(image_uri: str) -> bool:
    """Push a Docker image to ECR."""
    try:
        print(f"Pushing image: {image_uri}")
        result = subprocess.run(
            f"docker push {image_uri}",
            capture_output=True,
            text=True,
            shell=True
        )
        if result.returncode == 0:
            print(f"✓ Successfully pushed {image_uri}")
            return True
        else:
            print(f"✗ Failed to push {image_uri}: {result.stderr}")
            return False
    except Exception as e:
        print(f"✗ Error pushing image: {e}")
        return False


def remove_local_image(image_uri: str) -> bool:
    """Remove a Docker image from local machine."""
    try:
        print(f"Removing local image: {image_uri}")
        result = subprocess.run(
            f"docker rmi {image_uri} -f",
            capture_output=True,
            text=True,
            shell=True
        )
        if result.returncode == 0:
            print(f"✓ Successfully removed {image_uri}")
            return True
        else:
            print(f"✗ Failed to remove {image_uri}: {result.stderr}")
            return False
    except Exception as e:
        print(f"✗ Error removing image: {e}")
        return False


def main():
    """Main function to pull, modify, and push ECR images."""
    print("=" * 60)
    print("ECR Image Updater")
    print("=" * 60)
    
    # Use the IMAGE_LIST constant
    images = IMAGE_LIST
    
    if not images:
        print("No images found in IMAGE_LIST")
        sys.exit(1)
    
    print(f"\nFound {len(images)} image(s) to process:\n")
    for img in images:
        new_tag = get_new_tag(img)
        print(f"  - {img}")
        print(f"    → {new_tag}")
    print()
    
    # Track unique registries for login
    logged_in_registries = set()
    
    # Process each image
    pull_success = 0
    pull_failed = 0
    build_success = 0
    build_failed = 0
    push_success = 0
    push_failed = 0
    
    for image_uri in images:
        print("-" * 60)
        
        # Extract registry and region
        registry, region = extract_registry_info(image_uri)
        
        if not registry or not region:
            print(f"✗ Could not parse ECR URI: {image_uri}")
            pull_failed += 1
            continue
        
        # Login to registry if not already logged in
        if registry not in logged_in_registries:
            print(f"Authenticating with ECR registry: {registry}")
            password = get_ecr_password(region)
            
            if not password:
                print(f"✗ Could not authenticate with {registry}")
                pull_failed += 1
                continue
            
            if docker_login(registry, password):
                logged_in_registries.add(registry)
            else:
                pull_failed += 1
                continue
        
        # Step 1: Pull the image
        if not pull_image(image_uri):
            pull_failed += 1
            continue
        pull_success += 1
        
        # Step 2: Build modified image with new tag
        new_image_uri = get_new_tag(image_uri)
        if not build_modified_image(image_uri, new_image_uri):
            build_failed += 1
            continue
        build_success += 1
        
        # Step 3: Push the new image to ECR
        if not push_image(new_image_uri):
            push_failed += 1
            continue
        push_success += 1
        
        # Step 4: Remove local images to save space
        print("Cleaning up local images...")
        remove_local_image(image_uri)
        remove_local_image(new_image_uri)
    
    # Summary
    print("=" * 60)
    print("Summary")
    print("=" * 60)
    print(f"Total images:    {len(images)}")
    print(f"Pull success:    {pull_success}")
    print(f"Pull failed:     {pull_failed}")
    print(f"Build success:   {build_success}")
    print(f"Build failed:    {build_failed}")
    print(f"Push success:    {push_success}")
    print(f"Push failed:     {push_failed}")
    print("Local cleanup:   Done")
    
    if pull_failed > 0 or build_failed > 0 or push_failed > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()