vpc

#! /usr/bin/env python

import argparse
import logging
import sys
import time

from itertools import islice

import boto3
import netaddr


logging.basicConfig(level=logging.INFO)

MAX_AZS = 3

def create_vpc(name, cidr, tenancy, ec2):
    resp = ec2.create_vpc(CidrBlock=cidr, InstanceTenancy=tenancy)

    vpc_id =  resp['Vpc']['VpcId']

    logging.info("Created vpc %s", vpc_id)
    
    ec2.get_waiter('vpc_available').wait(VpcIds=[vpc_id])
    
    ec2.modify_vpc_attribute(VpcId=vpc_id, EnableDnsSupport={ 'Value': True })
    ec2.modify_vpc_attribute(VpcId=vpc_id, EnableDnsHostnames={ 'Value': True })

    set_name(vpc_id, name, ec2)

    return vpc_id


def is_public(subnet):
    words = subnet.ip.words
    return words[2] > 10 and words[2] % 10 == 2


def is_private(subnet):
    words = subnet.ip.words
    return words[2] > 10 and words[2] % 10 == 1


def availability_zones(ec2):
    return islice((az['ZoneName'] for az in ec2.describe_availability_zones()['AvailabilityZones']), MAX_AZS)
                                                                                             

def _create_subnet(vpc_id, cidr, az, description, ec2):
    resp = ec2.create_subnet(VpcId=vpc_id, CidrBlock=cidr, AvailabilityZone=az)
    subnet_id = resp['Subnet']['SubnetId']
    
    ec2.get_waiter('subnet_available').wait(SubnetIds=[subnet_id])

    set_name(subnet_id, "%s-%s" % (description, az), ec2)

    logging.info("Created subnet %s", subnet_id)
    
    return subnet_id


def create_public_subnets(vpc_id, cidr, description, ec2):
    """Create a public subnet for each AZ available to vpc_id."""
    logging.info("Creating public subnets")
    network = netaddr.IPNetwork(cidr)
    mapping = zip(availability_zones(ec2),
                  (subnet for subnet in network.subnet(24) if is_public(subnet)))

    desc = "%s-public" % description
    subnets = [ (_create_subnet(vpc_id, str(net[1]), net[0], desc, ec2), net[0]) for net in mapping ]

    return subnets


def create_private_subnets(vpc_id, cidr, description, ec2):
    """Create a private subnet for each AZ available to vpc_id."""
    logging.info("Creating private subnets")
    network = netaddr.IPNetwork(cidr)
    mapping = zip(availability_zones(ec2),
                  (subnet for subnet in network.subnet(24) if is_private(subnet)))

    desc = "%s-private" % description
    subnets = [ (_create_subnet(vpc_id, str(net[1]), net[0], desc, ec2), net[0]) for net in mapping ]

    return subnets


def create_gateway(vpc_id, description, ec2):
    """Create an internet gateway and attach it to vpc_id."""
    resp = ec2.create_internet_gateway()
    gateway_id = resp['InternetGateway']['InternetGatewayId']

    logging.info("Created gateway %s", gateway_id)
    
    time.sleep(15)
    
    set_name(gateway_id, "%s-igw" % description, ec2)
        
    ec2.attach_internet_gateway(InternetGatewayId=gateway_id, VpcId=vpc_id)

    return gateway_id


def create_private_routing_tables(vpc_id, description, subnets, ec2):
    for (subnet_id, az) in subnets:
        resp = ec2.create_route_table(VpcId=vpc_id)
        route_table_id = resp["RouteTable"]["RoutingTableId"]

        set_name(route_table_id, "%s-private-%s" % (description, az))
        ec2.associate_route_table(SubnetId=subnet, RouteTableId=route_table_id)
                 
                
def create_command(args, ec2):
    tenancy = 'dedicated' if args.dedicated else 'default'
    vpc_id = create_vpc(args.name, args.cidr, tenancy, ec2)

    public_subnets = create_public_subnets(vpc_id, args.cidr, args.name, ec2)
    private_subnets = create_private_subnets(vpc_id, args.cidr, args.name, ec2)

    gateway_id = create_gateway(vpc_id, args.name, ec2)

    # Temporary output since I'm not done yet but need to get
    # auth-prod up and running
    print("VpcId: %s" % vpc_id)
    print("GatewayId: %s" % gateway_id)
    print(net[0] for net in public_subnets)
    print(net[0] for net in private_subnets)

    
def set_name(resource_id, name, ec2):
    """Set the Name tag on resource_id to name."""
    ec2.create_tags(
        Resources=[resource_id],
        Tags=[
            {
                "Key": "Name",
                "Value": name,
            }
        ]
    )


def name(resource):
    if "Tags" not in resource:
        return "N/A"
    return next((t["Value"] for t in resource["Tags"] if t["Key"] == "Name"), "N/A")


def list_command(args, ec2):
    resp = ec2.describe_vpcs()
    output = []
    for vpc in resp["Vpcs"]:
        vpc["Name"] = name(vpc)
        output.append("{Name}\t{VpcId}\t{CidrBlock}".format(**vpc))

    print "\n".join(sorted(output))


def subnets_command(args, ec2):
    resp = ec2.describe_subnets(Filters=[
        {
            "Name": "vpc-id",
            "Values": [ args.vpc_id ]
        }
    ])
    output = []
    for subnet in resp["Subnets"]:
        subnet["Name"] = name(subnet)
        output.append("{Name}\t{SubnetId}\t{CidrBlock}\t{AvailabilityZone}".format(**subnet))

    print "\n".join(sorted(output))
            
        
def parse_args(argv):
    parser = argparse.ArgumentParser(
        add_help=True,
        description='VPC management'
    )
    subparsers = parser.add_subparsers()

    create_parser = subparsers.add_parser("create", help="Create a VPC")
    create_parser.add_argument("name", help="VPC name")
    create_parser.add_argument("cidr", help="CIDR block")
    create_parser.add_argument("--dedicated", action='store_true', help="Force dedicated instances")
    create_parser.set_defaults(func=create_command)

    list_parser = subparsers.add_parser("list", help="List VPCs")
    list_parser.set_defaults(func=list_command)

    subnets_parser = subparsers.add_parser('subnets', help="List VPC subnets")
    subnets_parser.add_argument("vpc_id")
    subnets_parser.set_defaults(func=subnets_command)
    
    return parser.parse_args(argv)


def main(args):
    ec2 = boto3.client('ec2')    
    args.func(args, ec2)

    
if __name__ == "__main__":
    main(parse_args(sys.argv[1:]))