#!/usr/bin/python

import sys
import os
import argparse
import boto
import boto.vpc
import boto.ec2
from boto.utils import get_instance_metadata
import requests
import json
import time

parser = argparse.ArgumentParser(description='HA NAT with consul')
parser.add_argument('command', choices=['start','takeover'])
args = parser.parse_args()

#region = get_instance_metadata()['placement']['availability-zone'][:-1]
region = 'eu-west-1'
vpc_conn = boto.vpc.connect_to_region(region)
ec2_conn = boto.ec2.connect_to_region(region)

def become_gateway(vpc_conn):
    nat_rt = vpc_conn.get_all_route_tables(filters={'tag:Name':'NAT'})[0]
    has_default_route = False
    for route in nat_rt.routes:
        if route.destination_cidr_block == '0.0.0.0/0':
            has_default_route = True
    if has_default_route:
        vpc_conn.replace_route(
                route_table_id=nat_rt.id,
                destination_cidr_block='0.0.0.0/0',
                instance_id=get_instance_metadata()['instance-id'])
    else:
        vpc_conn.create_route(
                route_table_id=nat_rt.id,
                destination_cidr_block='0.0.0.0/0',
                instance_id=get_instance_metadata()['instance-id'])

def disable_source_dest_check(ec2_conn):
    ec2_conn.modify_instance_attribute(get_instance_metadata()['instance-id'],
            attribute='sourceDestCheck', value=False)

def find_nat_instances(ec2_conn):
    instances = ec2_conn.get_only_instances(filters= {
        'tag:Type':'NATInstance',
        'instance-state-name':'running'
        })
    return [i.private_ip_address for i in instances]

def create_session():
    session_data = json.dumps({
        'Name': 'NatLock',
        'Checks': ['serfHealth','service:nat'],
        'LockDelay': 0
            })
    response = requests.put('http://localhost:8500/v1/session/create',data=session_data)
    return json.loads(response.content)['ID']

def delete_session(session):
    response = requests.put("http://localhost:8500/v1/session/destroy/%s" % session)
    if response.status_code == 200:
        return True
    else:
        return False

def won_election():
    status_code = 0
    while status_code in (0,500):
        time.sleep(5)
        session = create_session()
        response = requests.put("http://localhost:8500/v1/kv/nat?acquire=%s" % session)
        status_code = response.status_code
    if response.content == "true":
        return True
    else:
        delete_session(session)
        return False

if args.command == 'takeover':
    if won_election():
        print "I won!"
        disable_source_dest_check(ec2_conn)
        become_gateway(vpc_conn)
        sys.exit(0)
    else:
        print "I lost"
        sys.exit(0)
if args.command == 'start':
    command = '/usr/local/bin/consul agent -bootstrap-expect=3 -server '
    command += '-data-dir=/var/run/consul -config-dir=/etc/consul.d'
    for ip in find_nat_instances(ec2_conn):
        command += " -retry-join %s" % ip
    os.system(command)
