#!/usr/bin/env python3
'''Launch an instance on ec2 to test cloud-init'''

try:
    import SoftLayer
except ImportError:
    raise RuntimeError(
        'Could not import SoftLayer.  Please install SoftLayer from pypi.')

from base64 import b64encode
import os
import re
import shutil
from subprocess import (
    CalledProcessError, DEVNULL, PIPE, STDOUT, Popen, call, check_call,
    check_output)
import sys
from time import sleep

DEFAULT_REGION = 'dal05',
KEY_PAIR_NAME = 'cloud-init-integration'
DEFAULT_SERIES = 'xenial'
SECGROUP_NAME = 'Cloud-init integration test secgroup'
SECGROUP_DESCRIPTION = 'Security group setting ssh ingress to instances'


# Directory where script artifacts can be saved for sync
DEFAULT_LOCAL_ARTIFACT_DIR = 'test-artifacts'
REMOTE_ARTIFACT_DIR = '/tmp/cloud-init-script-artifacts'

REGIONS = (
    "ams01",
    "ams03",
    "che01",
    "dal01",
    "dal05",
    "dal06",
    "dal09",
    "dal10",
    "dal12",
    "dal13",
    "fra02",
    "hkg02",
    "hou02",
    "lon02",
    "lon04",
    "lon06",
    "mel01",
    "mex01",
    "mil01",
    "mon01",
    "osl01",
    "par01",
    "sao01",
    "sea01",
    "seo01",
    "sjc01",
    "sjc03",
    "sjc04",
    "sng01",
    "syd01",
    "syd04",
    "tok02",
    "tor01",
    "wdc01",
    "wdc04",
    "wdc06",
    "wdc07",
)


# See also http://boto3.readthedocs.io/en/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
CREATE_INSTANCE_DEFAULTS = {
    'image_id': 'ami-FIXME',  # Must be set by our request
    'KeyName': 'cloud-init-integration',
    'InstanceType': 't2.micro',
    'SecurityGroups': ['FIXME'],
    #Ipv6AddressCount=123,
    'MaxCount': 1,  # Deploy 1 instance of this type
    'MinCount': 1,
    'Monitoring': {
        'Enabled': False
    },
    'Placement': {
        'AvailabilityZone': '{}{}'.format(
            DEFAULT_REGION, DEFAULT_AVAILABILITY_ZONE),
    },
    'DisableApiTermination': False,
    'EbsOptimized': False,
    'InstanceInitiatedShutdownBehavior': 'stop'
}

class SoftLayerInstance(object):
    """Wrapper around a boto ec2 instance providing methods for testing.

    Methods provided simplify integration tests and validation.
    """

    _cloud_init_ran = False   # True oncecloud-init completes on the instance

    def __init__(self, instance, pubkey, artifact_dir=REMOTE_ARTIFACT_DIR,
                 verbose=False):
        """Initialize this EC2Instance helper class.

        @param instance: ec2 boto3 instance object.
        @param pubkey: Local path to the pubkey we'll use when contacting the
            instance.
        @param artifact_dir: The remote directory in which script results can
            be stored.  This value is provided via runcmd as an environemnt
            variable $SCRIPT_ARTIFACT_DIR.
        @param verbose: Boolean, set True to see commands sent to and from
            instance.
        """
        self.instance = instance
        self.pubkey = pubkey
        self.artifact_dir = artifact_dir
        self.verbose = verbose

    def scp(self, source_path, dest_path):
        """Scp files or dirs to an instance."""
        hostname = self.instance.meta.data['PublicDnsName']
        instance_user_at_host = 'ubuntu@{}'.format(hostname)
        source_path = source_path.replace('{INSTANCE}', instance_user_at_host)
        dest_path = dest_path.replace('{INSTANCE}', instance_user_at_host)
        cmd = ['scp', '-i', self.pubkey, '-o', 'StrictHostKeyChecking=no',
               '-o', 'UserKnownHostsFile=/dev/null', source_path, dest_path]
        check_output(cmd)

    def runcmd(self, command, pipe_in=None, pipe_out=False):
        """Run a command over ssh on the instance.

        @param command: The string of command(s) to execute on the remote
            instance.
        @param pipe_in: Optional Popen process to take as stdin to the command.
        @param pipe_out: Optionally, whether to return the popen process
            spawned to allow subprocess.PIPEs to use the stdout of the command.
        """
        hostname = self.instance.meta.data['PublicDnsName']
        stdin = None
        if pipe_in:
            stdin = pipe_in.stdout
        ssh_cmd = [
            'ssh', '-i', self.pubkey, '-o', 'StrictHostKeyChecking=no', '-o',
            'UserKnownHostsFile=/dev/null', 'ubuntu@{}'.format(hostname), '--',
            'SCRIPT_ARTIFACT_DIR={}'.format(self.artifact_dir)]
        proc = Popen(ssh_cmd + [command], stdin=stdin, stdout=PIPE)
        if pipe_in:
            pipe_in.stdout.close()  # Allow p1 to receive SIGPIPE if p2 exits
        if pipe_out:
            return proc  # Caller needs to communicate and close
        return proc.communicate()

    def run_scripts(self, scripts_dir, artifacts_dir):
        """Push scripts_dir to the instance, run them and return artifacts.

        Provide SCRIPT_ARTIFACT_DIR env variable to scripts and copy any files
        in remote SCRIPT_ARTIFACT_DIR into local artifacts_dir.

        @param scripts_dir: Local path to a scripts directory which contains
            executable scripts which can be run by run-parts on the instance.
        @param artifacts_dir: Local path where script artifacts/results will be
            copied.
        """
        local_cmd = ['tar', '-czf', '-', scripts_dir]
        remote_artifact_basename = os.path.basename(self.artifact_dir)

        # Untar scripts_dir, run-parts scripts_dir, tar up self.artifact_dir
        remote_cmds = [
            'tar xzf - -C /tmp',
            'mkdir -p {0}'.format(self.artifact_dir),
            'SCRIPT_ARTIFACT_DIR={0} run-parts /tmp/{1}'.format(
                self.artifact_dir, scripts_dir),
            'tar -czf - -C /tmp {}'.format(remote_artifact_basename)]

        proc1 = Popen(local_cmd, stdout=PIPE)  # Tar up local scripts_dir
        # Perform all remote_cmds in a single ssh interaction.
        proc2 = self.runcmd(
            '; '.join(remote_cmds), pipe_in=proc1, pipe_out=True)
        # Untar self.artifact_dir locally
        proc3 = Popen(['tar', '-xzf', '-'], stdin=proc2.stdout, stdout=PIPE)
        proc2.stdout.close()
        out, err = proc3.communicate()
        # Move the remote_artifacts dirname to local artifacts_dir
        shutil.move(remote_artifact_basename, artifacts_dir)

    def wait_on_cloud_init(self):
        print("Waiting on cloud-init")
        self._cloud_init_ran = True
        out, err = self.runcmd(
            "while [ ! -f '/run/cloud-init/result.json' ]; do"
            " echo -n '.'; sleep 1; done")

    def update_proposed_cloud_init(self):
        """Update cloud-init package to the version present in proposed."""
        match = re.match(
            r'.*ubuntu-(?P<series>[^-]*).*', self.instance.image.name)
        if not match:
            raise RuntimeError(
                'Could not determine ubuntu series from image: {}'.format(
                    self.instance.image.name))
        series = match.group('series')
        self.runcmd(
            "egrep 'deb .*{0} main' /etc/apt/sources.list |"
            " sed 's/{0}/{0}-proposed/' > /tmp/cloud-init-proposed.list;"
            " sudo mv /tmp/cloud-init-proposed.list /etc/apt/sources.list.d/;"
            " sudo apt-get update;"
            " sudo  DEBIAN_FRONTEND=noninteractive apt-get install cloud-init"
            " --yes".format(series))

    def clean(self, reboot=False):
        """Clean cloud-init artifacts from the system making it look 'new'.

        A cleaned system allows cloud-init to boot as if it is seeing the
        instance on first boot.
        """
        clean_cmd = 'sudo rm -Rf /var/lib/cloud /var/log/cloud*'
        self.runcmd(clean_cmd)
        self._cloud_init_ran = True
        if reboot:
            print("Waiting for EC2 instance clean reboot")
            self.reboot()

    def destroy(self):
        """Stop and destroy the instance from EC2."""
        self.instance.terminate()

    def reboot(self):
        """Stop and destroy the instance from EC2."""
        self.instance.reboot()
        sleep(8)
        self.instance.reload()
        self.wait_on_cloud_init()

    def start(self):
        """Start the instance running in EC2."""
        self.instance.start()
        self.wait_on_cloud_init()

    def stop(self):
        """Stop the instance from running in EC2."""
        self.instance.stop()

def check_deps(packages):
    missing_packages = []
    for package in packages:
        dpkg_cmd = ['dpkg', '-l', package]
        if call(dpkg_cmd, stdout=DEVNULL, stderr=STDOUT) != 0:
            missing_packages.append(package)
    if missing_packages:
        error('Missing packages: apt install {0}'.format(
                  ' '.join(missing_packages)))

def get_image_id(series_name):
    """Return a valid image id for the given series and region_name.

    Options can be seen in 'slcli vs create-options'.
    """

    # not kidding.  precise is 'UBUNTU_12_64', xenial UBUNTU_16_64.
    s2v = {'precise': '12_64', 'trusty': '14_64', 'xenial': '16_64'}
    if series_name not in s2v:
        raise RuntimeError(
            "Couldn't find an image id for series {} in {}".format(
            series_name, region_name))
    return 'UBUNTU_%s' % s2v[series_name]


def get_parser():
    """Return an argument parser for this command."""
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        '-i', '--image-id', type=str, help='Specify the SoftLayer id to deploy')
    parser.add_argument(
        '--deb-file', type=str, required=False, dest='deb_file',
        help='Provide a local deb for install on the instance')
    parser.add_argument(
        '--script-dir', type=str, dest='script_dir',
        help=('Specify a directory of scripts to run on the deployed instance.'
              ' Scripts should use SCRIPT_ARTIFACT_DIR env variable for'
              ' output.'))
    parser.add_argument(
        '--pubkey-file', required=False, type=str, dest='pubkey_file',
        default=os.path.expanduser('~/.ssh/id_rsa.pub'),
        help=('Specify a local directory where results of script output'
              ' artifacts are saved'))
    parser.add_argument(
        '--script-artifacts-dir', type=str, dest='artifacts_dir',
        default=DEFAULT_LOCAL_ARTIFACT_DIR,
        help=('Specify a local directory where results of script output'
              ' artifacts are saved'))
    parser.add_argument(
        '--keep-alive', action='store_true', default=False,
        dest='keep_alive', help='Leave the instance running after test.')
    parser.add_argument(
        '-p', '--proposed', action='store_true', default=False,
        help='Update instance cloud-init to the version in <series>-proposed.')
    parser.add_argument(
        '-r', '--region', type=str, choices=REGIONS,
        default=DEFAULT_REGION,
        help='Specify a region to deploy to [default=%s]' % DEFAULT_REGION)
    parser.add_argument(
        '-u', '--user-data-file', dest='user_data_file', type=str,
        help='Optional user-data file to run during instance initialization')
    parser.add_argument(
        '-s', '--series', type=str, default=DEFAULT_SERIES,
        help='The default ubuntu series name to launch. [default={}]'.format(
                 DEFAULT_SERIES))
    parser.add_argument(
        '-v', '--verbose', action='store_true', default=False,
        help='Print more information about the operations performed.')
    parser.add_argument(
        '-c', '--clean', action='store_true', default=False,
        help=('Remove cloud-init artifacts and reboot system to re-run from'
              'a pseudo-clean system boot.'))
    return parser


def import_keypair(client, key_path, name=None):
    if name is None:
        name = KEY_PAIR_NAME
    keymgr = SoftLayer.SshKeyManager(client)
    keys = [
        k for k in keymgr.list_keys(label=name)
    if keys:
        return keys[0]
    with open(key_path, 'rb') as stream:
        key_material = stream.read()
    keymgr.add_key(keydata, label=name)


def ec2_import_keypair(ec2_client, key_path):
    """Create a new key in EC2 if not present using the pubkey at key_path.

    @returns: ec2.KeyParInfo
    """
    pairs = [
        k for k in ec2_client.key_pairs.filter(KeyNames=[KEY_PAIR_NAME])]
    if pairs:
        return pairs[0]  # We already have the key
    with open(key_path, 'rb') as stream:
        key_material = stream.read()
    return ec2_client.import_key_pair(
        KeyName=KEY_PAIR_NAME, PublicKeyMaterial=key_material)


def ec2_create_secgroup(ec2_client, sec_group_name=None, source_cidr=None):
    """Create cloud-init's test sec group with ssh ingress from source_cidr."""
    if sec_group_name is None:
        sec_group_name = SECGROUP_NAME
    if source_cidr is None:
        source_cidr = '0.0.0.0/0'
    try:
        existing_groups = [
            group for group in ec2_client.security_groups.filter(
                GroupNames=[SECGROUP_NAME])]
        return existing_groups[0]
    except ClientError as e:
        if e.response['Error']['Code'] == 'InvalidGroup.NotFound':
            pass
    secgroup = ec2_client.create_security_group(
        GroupName=sec_group_name, Description=SECGROUP_DESCRIPTION)
    permissions = [
        {'FromPort':22, 'ToPort': 22, 'IpProtocol': 'tcp',
         'IpRanges': [{'CidrIp': source_cidr}]}]
    secgroup.authorize_ingress(IpPermissions=permissions)
    return secgroup


def create_instance(client, image, user_data_file=None):
    """Create and instance and wait for started state."""

    opts = CREATE_INSTANCE_DEFAULTS.copy()
    try:
        # if image is an integer, its a 'image_id'. if not, an 'os_code'
        opts['image_id'] = int(image)
    except ValueError as e:
        opts['os_code'] = image
    
    kwargs['ImageId'] = ami_id
    kwargs['SecurityGroups'] = [SECGROUP_NAME]
    if user_data_file:
        if not os.path.exists(user_data_file):
            raise RuntimeError(
                'user-data file {} does not exist'.format(user_data_file))
        with open(user_data_file) as stream:
            kwargs['UserData'] = stream.read()
    [instance] = ec2_client.create_instances(**kwargs)
    # block until running
    print('Waiting for EC2 instance initialization')
    instance.wait_until_running(
        Filters=[{'Name': 'instance-state-name', 'Values': ['running']}])
    instance.reload()
    print('Found EC2 instance IP: {}'.format(
        instance.meta.data['PublicDnsName']))
    return instance


def error(message):
    sys.stderr.write('ERROR: {0}\n'.format(message))
    sys.exit(1)

def main():
    parser = get_parser()
    args = parser.parse_args()
    if not os.path.isfile(args.pubkey_file):
        error("--pubkey-file=%s: not a file\n" % args.pubkey_file)
    image_id = args.image_id
    if not image_id:
        image_id = get_image_id(args.series)
    client = SoftLayer.Client()
    try:
        import_keypair(client, args.pubkey_file)
        instance = ec2_create_instance(
            ec2_client, image_id, args.user_data_file)
    except NoCredentialsError as e:
        error('Credentials undefined or incorrect. Check ~/.aws/credentials')

    mgr = SoftLayer.VSManager(client)

    ec2_instance = EC2Instance(instance, args.pubkey_file)
    if args.deb_file:
        ec2_instance.scp(args.deb_file, '{INSTANCE}:.')
        deb_filename = os.path.basename(args.deb_file)
        ec2_instance.runcmd('sudo apt-get install --yes python3-jsonschema;'
                            ' sudo dpkg -i ./{}'.format(deb_filename))
    if args.proposed:
        ec2_instance.update_proposed_cloud_init()
    if args.clean:
        ec2_instance.clean(reboot=True)
    if args.script_dir:
        ec2_instance.run_scripts(args.script_dir, args.artifacts_dir)
    if not args.keep_alive:
        ec2_instance.destroy()
    return 0

if __name__ == "__main__":
    sys.exit(main())