#!/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
import re

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_archive(tmp_dir):
    tar = tarfile.open(join(tmp_dir, 'docker-build-ami.tar.gz'), 'w:gz')
    for fn in glob.glob('*'):
        logger.info('Adding file to archive: %s' % fn)
        tar.add(fn)
    tar.close()

def ssh_cmd(ssh, cmd):
    print color.YELLOW
    stdin, stdout, stderr = ssh.exec_command('set -ex; echo %s | sudo -i --' % pipes.quote(cmd), get_pty=True)
    print stdout.read()
    print color.RED
    print stderr.read()
    print color.CLEAR

    ecode = stdout.channel.recv_exit_status()
    if ecode != 0:
        logger.error('The command %s returned a non-zero code: %s' % (cmd, ecode))
        exit(ecode)

def parse_dockerfile(ssh, fn):
    env = ''
    step = 0
    mline = False
    with open(fn) as fp:
        for line in fp.read().splitlines():
            # Multi-line command
            if mline:
                cmd += line.lstrip()
                if not re.match('.*\\\s*$', line):
                    mline = False
                    print 'Step %s : RUN %s' % (step, cmd)
                    ssh_cmd(ssh, '%s%s' % (env, cmd))
                    step += 1

            # Skip empty lines and  comments
            elif re.match('^(#.*|)$', line):
                continue

            # Environment variable
            elif re.match('^ENV +', line, re.IGNORECASE):
                op, key, val = re.split(' +', line, 2)
                env += '%s="%s"; ' % (key, val)

            # Run command
            elif re.match('^RUN +', line, re.IGNORECASE):
                op, cmd = re.split(' +', line, 1)
                if re.match('.*\\\s*$', line):
                    mline = True
                else:
                    print 'Step %s : RUN %s' % (step, cmd)
                    ssh_cmd(ssh, '%s%s' % (env, cmd))
                    step += 1

            # Copy command
            elif re.match('^COPY +', line, re.IGNORECASE):
                op, src, dst = re.split(' +', line, 2)
                print 'Step %s : COPY %s %s' % (step, src, dst)
                ssh_cmd(ssh, '%scp -rf /tmp/docker-build-ami/%s %s' % (env, src, dst))
                step += 1

            # Add command
            elif re.match('ADD +', line, re.IGNORECASE):
                op, src, dst = re.split(' +', line, 2)
                print 'Step %s : ADD %s %s' % (step, src, dst)
                if re.match('.*\.(tar|tar\.gz|tar\.bz|tar\.xz)$', src):
                    ssh_cmd(ssh, '%star -xf /tmp/docker-build-ami/%s -C %s' % (env, src, dst))
                else:
                    ssh_cmd(ssh, '%scp -rf /tmp/docker-build-ami/%s %s' % (env, src, dst))
                step += 1

# 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', help='AWS region')
parser.add_argument('-t', '--instance-type', help='EC2 instance type')
parser.add_argument('-s', '--subnet-id', help='AWS subnet id')
parser.add_argument('-n', '--image-name', help='Target AMI image name')
parser.add_argument('-i', '--image-id', help='Source AMI image ID')
parser.add_argument('-u', '--image-user', 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', '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_name = str(uuid.uuid4())
key_path = expanduser(join(tmp_dir, key_name + '.pem'))

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

create_archive(tmp_dir)

try:
    ec2 = 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 = ec2.create_key_pair(key_name)
key_pair.save(tmp_dir)

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

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

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

print 'Instance: %s' % instance.id
print 'Instance IP: %s' % instance.private_ip_address
print 'Connection 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 '\nConnecting to host'

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

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

# Copy archive
sftp = ssh.open_sftp()
sftp.put(tmp_dir + '/docker-build-ami.tar.gz', '/tmp/docker-build-ami.tar.gz')
sftp.close()

# Untar archive
ssh_cmd(ssh, 'mkdir /tmp/docker-build-ami; tar -xzf /tmp/docker-build-ami.tar.gz -C /tmp/docker-build-ami')

# Parse Dockerfile
parse_dockerfile(ssh, 'Dockerfile')

ssh.close()
ec2.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 = ec2.get_image(image_id)

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

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