#!/usr/bin/env python

import argparse
import ConfigParser
import logging
import boto.ec2
import os
from os.path import expanduser, join, isfile
import time
import socket
from sys import stdout
import paramiko
import uuid
import datetime
import tarfile
import glob
import pipes

class color:
    RED = '\033[91m'
    YELLOW = '\033[33m'
    CLEAR = '\033[0m'

def check_port(host, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect((host, 22))
        s.shutdown(socket.SHUT_RDWR)
        s.close()
        return True
    except socket.error as e:
        return False

# Add support for .dockerignore
def create_tar_gz():
    tar = tarfile.open(tmp_dir + '/docker-build-ami.tar.gz', 'w:gz')
    for name in glob.glob('*'):
        logger.info('Adding file to %s/docker-build-ami.tar.gz: %s' % (tmp_dir, name))
        tar.add(name)
    tar.close()

def run(c, step, skip, cmd):
    print '--> {0}Step: {1} Operation: RUN{2}\n{3}'.format(color.YELLOW, step, color.CLEAR, 'set -ex; echo ' + pipes.quote(cmd) + ' | sudo -i --')
    if skip:
        print 'Skipping'
        return
    print color.YELLOW
#    stdin, stdout, stderr = c.exec_command('sudo -s -- <<EOF\nset -ex\n' + pipes.quote(cmd) + '\nEOF', get_pty=True)
#    stdin, stdout, stderr = c.exec_command('set -ex; sudo -i ' + pipes.quote(cmd), get_pty=True)
    stdin, stdout, stderr = c.exec_command('set -ex; echo ' + pipes.quote(cmd) + ' | sudo -i --', get_pty=True)
    print stdout.read()
    print stderr.read()
    print color.CLEAR

    exit_code = stdout.channel.recv_exit_status()
    if exit_code != 0:
        print '{0}Exit Code: {1}{2}'.format(color.RED, exit_code, color.CLEAR)
        exit()

# Get arguments
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', help='Configuration file')
parser.add_argument('-d', '--debug', action='store_true', help='Print debug info')
parser.add_argument('-r', '--region', action='store_true', help='AWS region')
parser.add_argument('-t', '--instance-type', action='store_true', help='EC2 instance type')
parser.add_argument('-s', '--subnet-id', action='store_true',  help='AWS subnet id')
parser.add_argument('-n', '--image-name', action='store_true',  help='AMI image name')
parser.add_argument('-i', '--image-id', action='store_true',  help='AMI image ID')
parser.add_argument('-u', '--image-user', action='store_true',  help='AMI image user')
args = parser.parse_args()

# Create formatter
try:
    import colorlog
    formatter = colorlog.ColoredFormatter("[%(log_color)s%(levelname)-8s%(reset)s] %(log_color)s%(message)s%(reset)s")
except ImportError:
    formatter = logging.Formatter("[%(levelname)-8s] %(message)s")

# Create console handle
console = logging.StreamHandler()
console.setFormatter(formatter)

loglvl = logging.WARN
if args.debug:
    loglvl = logging.DEBUG

# Create logger
logger = logging.getLogger(__name__)
logger.setLevel(loglvl)
logger.addHandler(console)

# Get configuration
config = ConfigParser.ConfigParser()
 
# Set defaults
config.add_section('main')
config.set('main', 'host_tag', 'docker-build-ami')
config.set('main', 'key_dir', '~/.ssh')
config.set('main', 'tmp_dir', '/tmp')
config.set('main', 'image_name', 'docker-build-ami')
config.set('main', 'region', 'us-west-1')
config.set('main', 'instance_type', 'm3.medium')
config.set('main', 'subnet_id', None)
config.set('main', 'image_id', 'ami-e4ff5c93')
config.set('main', 'image_user', 'centos')
config.set('main', 'aws_access_key_id', None)
config.set('main', 'aws_secret_access_key', None)

# Load config
cfile = None
if args.config:
    if isfile(expanduser(args.config)):
        cfile = expanduser(args.config)
    else:
        logger.error('Config file doesn\'t exist: {0}'.format(init_args.config))
        exit(1)
elif isfile(expanduser('~/.docker-build-ami.conf')):
    cfile = expanduser('~/.docker-build-ami.conf')
elif isfile('/etc/docker-build-ami.conf'):
    cfile = '/etc/docker-build-ami.conf'

if cfile and isfile(cfile):
    logger.info("Loading config file: {0}".format(cfile))
    config.read(cfile)

region = config.get('main', 'region')
if args.region:
    region = args.region

instance_type = config.get('main', 'instance_type')
if args.instance_type:
    instance_type = args.instance_type

subnet_id = config.get('main', 'subnet_id')
if args.subnet_id:
    subnet_id = args.subnet_id

image_name = config.get('main', 'image_name')
if args.image_name:
    image_name = args.image_name

ami_id = config.get('main', 'image_id')
if args.image_id:
    ami_id = args.image_id

user = config.get('main', 'image_user')
if args.image_user:
    user = args.image_user

aws_access_key_id = config.get('main', 'aws_access_key_id')
if not aws_access_key_id:
    logger.critical('You need to specify a AWS Access Key ID')
    exit(1)

aws_secret_access_key = config.get('main', 'aws_secret_access_key')
if not aws_secret_access_key:
    logger.critical('You need to specify a AWS Secret Access Key')
    exit(1)

host_tag = config.get('main', 'host_tag')

tmp_dir = config.get('main', 'tmp_dir')
key_dir = config.get('main', 'tmp_dir')
key_name = str(uuid.uuid4())
key_path = expanduser(join(key_dir, key_name + '.pem'))

fn = 'Dockerfile'
step = 0        
cont = False
cmd = ''

if not isfile(fn):
    logger.critical('There needs to be a Dockerfile in the current directory')
    exit(1)

create_tar_gz()

try:
    conn = boto.ec2.connect_to_region(region, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)
except:
    logger.critical('Failed to connect to EC2')
    exit(1)

key_pair = conn.create_key_pair(key_name)
key_pair.save(key_dir)

reservation = conn.run_instances(ami_id, key_name=key_name, instance_type=instance_type, subnet_id=subnet_id)

print 'Reservation ID: %s' % reservation.id

for r in conn.get_all_instances():
    if r.id == reservation.id:
        break

instance = r.instances[0]
instance.add_tag('Name', host_tag)

print 'ID: %s' % instance.id
print 'IP: %s' % instance.private_ip_address
print 'SSH KEY: %s' % key_path

stdout.write('Waiting for instance status running.')
stdout.flush()
while instance.state != 'running':
    stdout.write('.')
    stdout.flush()
    time.sleep(5)
    instance.update()

stdout.write('\nWaiting for SSH to become ready.')
stdout.flush()
while not check_port(instance.private_ip_address, 22):
    stdout.write('.')
    stdout.flush()
    time.sleep(5)

print '\nConnect to host'
print 'User: %s, Key: %s' % (user, key_path)

key = paramiko.RSAKey.from_private_key_file(key_path)
c = paramiko.SSHClient()
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())

c.connect(hostname=instance.private_ip_address, username=user, pkey=key)

# Copy tar.gz
sftp = c.open_sftp()
print sftp
print 'src: ' + tmp_dir + '/docker-build-ami.tar.gz' + ' dst: /tmp'
sftp.put(tmp_dir + '/docker-build-ami.tar.gz', '/tmp/docker-build-ami.tar.gz')
sftp.close()

# Untar tar.gz
run(c, step, False, 'mkdir /tmp/docker-build-ami; tar -xzf /tmp/docker-build-ami.tar.gz -C /tmp/docker-build-ami')
step += 1

skip = False
with open(fn) as fp:
    content = fp.read()
    for line in content.splitlines():
        if cont:
            cmd += line.lstrip()
            if line.rstrip().rfind('\\') != (len(line.rstrip()) - 1):
                cont = False
                run(c, step, skip, cmd)
                skip = False
                cmd = ''
                step += 1

        elif line.find('RUN ') == 0:
            cmd += line.lstrip('RUN ')
            if line.rfind('\\') == len(line) - 1:
                cont = True
            else:
                run(c, step, skip, cmd)
                skip = False
                cmd = ''
                step += 1

        elif line.find('COPY ') == 0:
            run(c, step, skip, 'cp -rf /tmp/docker-build-ami/' + line.lstrip('COPY '))
            skip = False
            step += 1

        elif line.find('#SKIP') == 0:
            skip = True

c.close()
conn.delete_key_pair(key_name)
os.remove(key_path)

print "Create AMI from instance: %s" % instance.id
image_id = instance.create_image(image_name + '-%s' % datetime.datetime.now().strftime('%Y%m%d%H%M%S'))
image = conn.get_image(image_id)

while image.state == 'pending':
    stdout.write('.')
    stdout.flush()
    time.sleep(5)
    image.update()

print "Created image: %s" % image_id
print "Terminate instance: %s" % instance.id
instance.terminate()
