#!/usr/bin/env python3

import boto3
import argparse

# ==========================================================================
# SCRIPT: privatelink_setup.py
# DESCRIPTION: Sets up VPC Gateway Endpoints for Amazon S3 and Amazon DynamoDB.
#              VPC Gateway Endpoints allow instances in your VPC to access
#              these AWS services privately, without traversing the public internet.
#              This enhances security and can reduce data transfer costs.
#
# USE CASE SCENARIO:
# An application running in a private subnet within a VPC needs to securely
# access Amazon S3 and DynamoDB without exposing traffic to the public internet.
# This script automates the creation of the necessary Gateway Endpoints and
# associates them with specified route tables.
#
# PREREQUISITES:
# 1.  **AWS Credentials:** The AWS CLI or SDK must be configured with credentials
#     that have the necessary permissions. This can be done via `~/.aws/credentials`,
#     environment variables, or an instance profile/IAM role for EC2 instances.
# 2.  **boto3 Library:** The AWS SDK for Python (`boto3`) must be installed.
#     You can install it using `pip install boto3`.
# 3.  **IAM Permissions:** The principal executing this script must have:
#     - `ec2:DescribeVpcEndpoints`
#     - `ec2:CreateVpcEndpoint`
#     - `ec2:DescribeRouteTables` (implicitly needed if you need to find route tables)
# 4.  **Existing Resources:**
#     - A VPC where the endpoints will be created.
#     - One or more Route Tables within that VPC that need to be updated to use the endpoint.
#
# HOW TO USE:
# 1.  **Save the script:** Save this content as `privatelink_setup.py`.
# 2.  **Make it executable (Linux/macOS):** `chmod +x privatelink_setup.py`
# 3.  **Run from your terminal:**
#     python privatelink_setup.py \
#       --vpc-id "vpc-0abcdef1234567890" \
#       --route-table-ids "rtb-0abcdef1234567890" "rtb-0fedcba9876543210" \
#       --region "us-east-1" \
#       --service "all"
#
#     **Arguments:**
#     - `--vpc-id`: The ID of the VPC where the endpoints will be created.
#     - `--route-table-ids`: Space-separated list of Route Table IDs to associate with the endpoint.
#                            Traffic destined for S3/DynamoDB from subnets associated with these
#                            route tables will be routed through the endpoint.
#     - `--region`: (Optional) The AWS region. Defaults to `us-east-1`.
#     - `--service`: (Optional) Specify `s3`, `dynamodb`, or `all` to create endpoints for specific services.
#                    Defaults to `all`.
#
# IMPORTANT CONSIDERATIONS:
# - This script creates Gateway Endpoints, which are free. Interface Endpoints (for other services)
#   incur charges.
# - Ensure your security groups and NACLs allow traffic to/from the endpoint.
# - The script checks for existing endpoints to prevent recreation errors.
# ==========================================================================

def setup_s3_gateway_endpoint(
    vpc_id: str,
    route_table_ids: list,
    region_name: str = 'us-east-1'
) -> str or None:
    """
    Sets up a VPC Gateway Endpoint for Amazon S3 in the specified VPC.

    Args:
        vpc_id (str): The ID of the VPC where the endpoint will be created.
        route_table_ids (list): A list of Route Table IDs to associate with the endpoint.
                                This updates the route tables to include a route to the S3 endpoint.
        region_name (str): The AWS region where the VPC and S3 service are located.
                           Defaults to 'us-east-1'.

    Returns:
        str or None: The ID of the created or existing S3 Gateway Endpoint, or None if an error occurred.
    """
    # Initialize the boto3 EC2 client for programmatic access to AWS EC2 service.
    ec2_client = boto3.client('ec2', region_name=region_name)

    print(f"\n>>> Setting up S3 Gateway Endpoint for VPC '{vpc_id}' in region {region_name}...")

    try:
        # Before creating, check if an S3 Gateway Endpoint already exists in this VPC.
        # This prevents errors if the script is run multiple times.
        existing_endpoints = ec2_client.describe_vpc_endpoints(
            Filters=[
                {'Name': 'vpc-id', 'Values': [vpc_id]},
                {'Name': 'service-name', 'Values': [f'com.amazonaws.{region_name}.s3']}
            ]
        )['VpcEndpoints']

        if existing_endpoints:
            # If an endpoint is found, print its ID and return it.
            print(f"   S3 Gateway Endpoint already exists in VPC '{vpc_id}'. Endpoint ID: {existing_endpoints[0]['VpcEndpointId']}")
            return existing_endpoints[0]['VpcEndpointId']

        # If no existing endpoint, proceed to create a new one.
        # `VpcEndpointType='Gateway'` specifies it's a Gateway Endpoint.
        # `ServiceName` is the specific AWS service endpoint to connect to.
        # `RouteTableIds` are updated to route S3 traffic through this endpoint.
        response = ec2_client.create_vpc_endpoint(
            VpcEndpointType='Gateway',
            VpcId=vpc_id,
            ServiceName=f'com.amazonaws.{region_name}.s3',
            RouteTableIds=route_table_ids,
            TagSpecifications=[
                {
                    'ResourceType': 'vpc-endpoint',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{vpc_id}-s3-gateway-endpoint'}
                    ]
                },
            ]
        )
        endpoint_id = response['VpcEndpoint']['VpcEndpointId']
        print(f"   S3 Gateway Endpoint '{endpoint_id}' created successfully.")
        return endpoint_id

    except Exception as e:
        # Handle any exceptions that occur during the API call.
        print(f"Error setting up S3 Gateway Endpoint: {e}")
        return None # Return None to indicate failure.

def setup_dynamodb_gateway_endpoint(
    vpc_id: str,
    route_table_ids: list,
    region_name: str = 'us-east-1'
) -> str or None:
    """
    Sets up a VPC Gateway Endpoint for Amazon DynamoDB in the specified VPC.

    Args:
        vpc_id (str): The ID of the VPC where the endpoint will be created.
        route_table_ids (list): A list of Route Table IDs to associate with the endpoint.
                                This updates the route tables to include a route to the DynamoDB endpoint.
        region_name (str): The AWS region where the VPC and DynamoDB service are located.
                           Defaults to 'us-east-1'.

    Returns:
        str or None: The ID of the created or existing DynamoDB Gateway Endpoint, or None if an error occurred.
    """
    # Initialize the boto3 EC2 client for programmatic access to AWS EC2 service.
    ec2_client = boto3.client('ec2', region_name=region_name)

    print(f"\n>>> Setting up DynamoDB Gateway Endpoint for VPC '{vpc_id}' in region {region_name}...")

    try:
        # Before creating, check if a DynamoDB Gateway Endpoint already exists in this VPC.
        existing_endpoints = ec2_client.describe_vpc_endpoints(
            Filters=[
                {'Name': 'vpc-id', 'Values': [vpc_id]},
                {'Name': 'service-name', 'Values': [f'com.amazonaws.{region_name}.dynamodb']}
            ]
        )['VpcEndpoints']

        if existing_endpoints:
            # If an endpoint is found, print its ID and return it.
            print(f"   DynamoDB Gateway Endpoint already exists in VPC '{vpc_id}'. Endpoint ID: {existing_endpoints[0]['VpcEndpointId']}")
            return existing_endpoints[0]['VpcEndpointId']

        # If no existing endpoint, proceed to create a new one.
        response = ec2_client.create_vpc_endpoint(
            VpcEndpointType='Gateway',
            VpcId=vpc_id,
            ServiceName=f'com.amazonaws.{region_name}.dynamodb',
            RouteTableIds=route_table_ids,
            TagSpecifications=[
                {
                    'ResourceType': 'vpc-endpoint',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{vpc_id}-dynamodb-gateway-endpoint'}
                    ]
                },
            ]
        )
        endpoint_id = response['VpcEndpoint']['VpcEndpointId']
        print(f"   DynamoDB Gateway Endpoint '{endpoint_id}' created successfully.")
        return endpoint_id

    except Exception as e:
        # Handle any exceptions that occur during the API call.
        print(f"Error setting up DynamoDB Gateway Endpoint: {e}")
        return None # Return None to indicate failure.

# ==========================================================================
# Main execution block to parse command-line arguments and call the functions.
# This allows the script to be run from the command line with various options.
# ==========================================================================
if __name__ == "__main__":
    # Create an argument parser object.
    parser = argparse.ArgumentParser(
        description="""
        Automates the setup of VPC Gateway Endpoints for Amazon S3 and DynamoDB.
        These endpoints enable private and secure access to these AWS services
        from within your VPC, bypassing the public internet.
        """
    )
    # Define the command-line arguments the script expects.
    parser.add_argument(
        "--vpc-id", 
        required=True, 
        help="The ID of the VPC where the endpoints will be created (e.g., 'vpc-xxxxxxxxxxxxxxxxx')."
    )
    parser.add_argument(
        "--route-table-ids", 
        nargs='+', # '+' means one or more arguments
        required=True, 
        help="A list of Route Table IDs to associate with the endpoints. Provide as space-separated values (e.g., rtb-xxxxxxxxxxxx rtb-yyyyyyyyyyyy)."
    )
    parser.add_argument(
        "--region", 
        default="us-east-1", 
        help="AWS region where the VPC and services are located (default: us-east-1)."
    )
    parser.add_argument(
        "--service", 
        choices=["s3", "dynamodb", "all"], 
        default="all", 
        help="Specify 's3', 'dynamodb', or 'all' to create endpoints for specific services (default: all)."
    )

    # Parse the arguments provided by the user when running the script.
    args = parser.parse_args()

    # Call the appropriate setup functions based on the --service argument.
    if args.service == "s3" or args.service == "all":
        s3_endpoint_id = setup_s3_gateway_endpoint(
            vpc_id=args.vpc_id,
            route_table_ids=args.route_table_ids,
            region_name=args.region
        )
        if s3_endpoint_id:
            print(f"S3 Gateway Endpoint setup complete. Endpoint ID: {s3_endpoint_id}")

    if args.service == "dynamodb" or args.service == "all":
        dynamodb_endpoint_id = setup_dynamodb_gateway_endpoint(
            vpc_id=args.vpc_id,
            route_table_ids=args.route_table_ids,
            region_name=args.region
        )
        if dynamodb_endpoint_id:
            print(f"DynamoDB Gateway Endpoint setup complete. Endpoint ID: {dynamodb_endpoint_id}")

    print("\n=== VPC Gateway Endpoints setup script finished. ===")
