#!/usr/bin/env python
"""
Warm up CloudFront edge locations.

Args:
    -p --protocol       http / https
    -d --distribution   Distribution hostname e.g. example.com
    -i --ignoressl      (Optional) Ignore SSL issues
    -f --files          Read file paths from the specified YAML config file
"""

from dns import resolver
import yaml
import os
import time
import argparse
import ssl
import requests
import sys
import subprocess
import uuid
import logging
import signal
from termcolor import colored


class ColorLog(object):

    colormap = dict(
        debug=dict(color='white'),
        info=dict(color='white', attrs=['bold']),
        warn=dict(color='yellow', attrs=['bold']),
        warning=dict(color='yellow', attrs=['bold']),
        error=dict(color='red', attrs=['bold']),
        critical=dict(color='red', attrs=['bold']),
    )

    def __init__(self, logger):
        self._log = logger

    def __getattr__(self, name):
        if name in ['debug', 'info', 'warn', 'warning', 'error', 'critical']:
            return lambda s, *args: getattr(self._log, name)(
                colored(s, **self.colormap[name]), *args)

        return getattr(self._log, name)


temp_dir='/tmp/{}'.format(uuid.uuid4())

logging.basicConfig(level=logging.INFO)
logger = ColorLog(logging.getLogger('aws-cf-warm'))

def handle_interuption(signal, frame):
    logger.error('Warming interupted, some requests may fail...')
    quit()

if __name__ == '__main__':
    original_sigint = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, handle_interuption)

try:
    with open('config/dns-servers.yml', 'r') as file:
        config = yaml.load(file)
except:
    logger.error('Error loading config/dns-servers.yml')
    quit()
    
def parse_arguments():
    PARSER = argparse.ArgumentParser(description='AWS CloudFront Warming CLI')
    PARSER.add_argument('-d', '--distribution', help="The CloudFront distribution hostname e.g. myapp.com")
    PARSER.add_argument('-p', '--protocol', help="http or https")
    PARSER.add_argument('-i', '--insecure', help="Ignore SSL issues", action='store_true')
    PARSER.add_argument('-f', '--files', help="Read config/files.yml for locating objects", default=False)
    return PARSER.parse_args()

ARGS = parse_arguments()
distribution = ARGS.distribution
protocol = ARGS.protocol
insecure = ARGS.insecure
files = ARGS.files

if files:
    try:
        with open(files, 'r') as file:
           file_paths = yaml.load(file)
    except:
        logger.error('Error loading config/files.yml')
        quit()

if insecure:
    logger.info('Ignoring SSL errors')

res = resolver.Resolver()

try:
    os.mkdir(temp_dir)
    os.mkdir('{}/downloads'.format(temp_dir))
except:
    logger.error('Couldnt make a temporary directory for the DNS host aliases')
    sys.exit(1)

# Used to avoid duplicated requests, as some DNS results may overlap
used_cf_ips=[]

for dns in config['ips']:
    ip = dns.values()
    region = dns.keys()[0]
    logger.info('Warming CloudFront distribution {}'.format(distribution))
    logger.info('Using {} DNS server {}'.format(region, ip[0]))
    res.nameservers = ip
    try:
        lookup = res.query(distribution)
    except:
        logger.error('Error resolving distribution hostname: {} using DNS server:'.format(distribution, ip))
        logger.error('Check your distribution is correct using "dig {}" and the DNS server using "dig @{} google.com"'.format(distribution, ip))
        logger.error('If the DNS server seems unhealthy, please file an issue at https://github.com/danielwhatmuff/aws-cf-warm/issues :-)')
        quit()
    else:
        host_file = '/etc/hosts'
        downloads_dir = '{}/downloads'.format(temp_dir)
        sub_env = os.environ.copy()

        wget_base_command = ['wget', '--quiet', '--no-dns-cache', '-P', downloads_dir]

        if protocol == 'https':
            logger.info('Using https mode')
            wget_base_command.insert(1, 'TLSv1')
            wget_base_command.insert(1, '--secure-protocol')
            wget_command = wget_base_command

            if insecure == True:
                logger.info('Running in insecure mode (avoid this if possible)')
                wget_command = wget_command.insert(2, '--no-check-certificate')
            else:
                logger.info('Running in secure SSL mode (recommended)')
        elif protocol == 'http':
            logger.info('Using http mode')
            wget_command = wget_base_command

        for result in lookup:
            if result.address in used_cf_ips:
                logger.info('Skipping duplicate IP {}'.format(result.address))
            else:
                used_cf_ips.append(result.address)
                logger.info('Using resultant IP {}'.format(result.address))
                # Append DNS to hosts file
                logger.info('Adding host entry')
                with open(host_file, 'a') as file:
                    file.write('{} {}'.format(result.address, distribution))

                if files == False:
                    logger.info('Running in recursive mode')
                    wget_command.insert(1, '-r')
                    wget_command.append(distribution)
                    try:
                        exit_code = subprocess.call(wget_command, stdin=subprocess.PIPE, stderr=subprocess.STDOUT)
                    except:
                        logger.error('Requests failed against {} DNS server {}'.format(region, ip[0]))
                else:
                    logger.info('Running against list of objects specified in config file {}'.format(files))
                    for file in file_paths:
                        url = '{}/{}'.format(distribution, file)
                        logger.info('Warming cache for {}'.format(url))
                        wget_command.append(url)
                        try:
                            exit_code = subprocess.call(wget_command, stdin=subprocess.PIPE, stderr=subprocess.STDOUT)
                        except:
                            logger.error('Requests failed against {} DNS server {}'.format(region, ip[0]))
                        wget_command = wget_command[:-1]

                # Remove the last hosts line added above
                logger.info('Removing host entry')
                with open(host_file, 'r') as file:
                    lines = file.readlines()[:-1]
                with open(host_file, 'w') as file:
                    for line in lines:
                        file.write(line)

                time.sleep(2)
