#!/usr/bin/env python3

import os
import re
import glob
import boto3
import requests
import subprocess

from time import sleep

AWS_REGION          = os.environ['AWS_REGION']
DEPLOY_UUID         = os.environ['DEPLOY_UUID']
SERVICE_NAME        = os.environ['SERVICE_NAME']
MOUNT_POINT         = "/var/lib/" + SERVICE_NAME
NIC_IP              = os.environ['NIC_IP']
TAG_KEY             = os.environ['TAG_KEY']


def retrieve_eni_ids():
    ec2 = boto3.resource('ec2')
    enis = []

    for eni in ec2.network_interfaces.all():
        for tag in eni.tag_set:
            if tag['Key'] == TAG_KEY:
                if tag['Value'] == DEPLOY_UUID:
                    enis.append(eni.network_interface_id)

    return enis if len(enis) > 0 else None


def attach_eni_ids():
    c_ec2 = boto3.client('ec2')
    r_ec2 = boto3.resource('ec2')

    i_id = requests.get('http://169.254.169.254/latest/meta-data/instance-id').text
    eni_ids = retrieve_eni_ids()
    device_number = len(r_ec2.Instance(i_id).network_interfaces) + 1

    for eni_id in eni_ids:
        c_ec2.attach_network_interface(DeviceIndex=device_number, InstanceId=i_id, NetworkInterfaceId=eni_id)


def retrieve_ebs_ids():
    ec2 = boto3.resource('ec2')
    ebss = []

    for volume in ec2.volumes.all():
        if volume.tags is not None:
            for tag in volume.tags:
                if tag['Key'] == TAG_KEY:
                    if tag['Value'] == DEPLOY_UUID:
                        ebss.append(volume.volume_id)

    return ebss if len(ebss) > 0 else None


def attach_ebs():
    ec2 = boto3.client('ec2')

    i_id = requests.get('http://169.254.169.254/latest/meta-data/instance-id').text
    volume_ids = retrieve_ebs_ids()

    i = 0
    device_char = 'z'
    while i < len(volume_ids):
        v_id = volume_ids[i]

        device = '/dev/xvd{0}'.format(device_char)
        ec2.attach_volume(Device=device, InstanceId=i_id, VolumeId=v_id)

        # Wait to ensure device is attached
        sleep(3)

        if not check_ebs(v_id):
            prepare_ebs(v_id)

        add_fstab_entries(v_id, MOUNT_POINT)

        p_mount = subprocess.Popen('mount -a'.split(), stdout=subprocess.PIPE)
        stdout, stderr = p_mount.communicate()

        p_chown = subprocess.Popen('chown -R {0}:{0} {1}'.format(SERVICE_NAME, MOUNT_POINT).split(),
                                   stdout=subprocess.PIPE)
        stdout, stderr = p_chown.communicate()

        device_char = chr(ord(device_char) - 1)
        i += 1


def check_ebs(volume_id):
    v_id = volume_id.replace('vol-', 'vol')
    pattern = '/dev/disk/by-id/*{0}-part1'.format(v_id)

    return bool(len(glob.glob(pattern)))


def prepare_ebs(volume_id):
    v_id = volume_id.replace('vol-', 'vol')
    pattern = '/dev/disk/by-id/*{0}'.format(v_id)
    device = glob.glob(pattern)[0]

    gdisk_commands = '\n'.join([
        'n',
        '1',
        '34',
        '',
        '',
        'w',
        'Y',
        ''
    ])

    p_echo = subprocess.Popen('echo -ne {0}'.format(gdisk_commands).split(' '), stdout=subprocess.PIPE)
    p_fdisk = subprocess.Popen('gdisk {0}'.format(device).split(), stdin=p_echo.stdout, stdout=subprocess.PIPE)
    stdout, stderr = p_fdisk.communicate()
    print(stdout)
    print(stderr)

    # p_partprobe = subprocess.Popen('partprobe'.split(' '), stdout=subprocess.PIPE)
    # stdout, stderr = p_partprobe.communicate()
    # print(stdout)
    # print(stderr)

    sleep(3)

    pattern = '/dev/disk/by-id/*{0}-part1'.format(v_id)
    partition = glob.glob(pattern)[0]

    p_xfs = subprocess.Popen('mkfs.xfs {0}'.format(partition).split(), stdout=subprocess.PIPE)
    stdout, stderr = p_xfs.communicate()
    print(stdout)
    print(stderr)


def add_fstab_entries(volume_id, mount_point):
    v_id = volume_id.replace('vol-', 'vol')
    pattern = '/dev/disk/by-id/*{0}-part1'.format(v_id)
    partition = glob.glob(pattern)[0]

    fstab_entries = [
        mount_point,
        'xfs',
        'defaults',
        '0',
        '0'
    ]

    with open('/etc/fstab', 'a') as f:
        f.write('{0} {1}\n'.format(partition, ' '.join(fstab_entries)))
        f.flush()
        f.close()



def wait_device_ready(timeout=3):
    c = 0
    while c < timeout:
        sleep(1)

        p_ip = subprocess.Popen('ip a'.split(), stdout=subprocess.PIPE)
        stdout, stderr = p_ip.communicate()

        for line in stdout.decode().splitlines():
            res = re.match('.*inet {0}/[0-9]{{2}}'.format(NIC_IP), line)

            if res is not None:
                return None

        c += 1

    raise Exception('Device with address {0} not ready'.format(NIC_IP))


def change_default_route():
    wait_device_ready(10)

    p_ip = subprocess.Popen('ip r'.split(), stdout=subprocess.PIPE)
    stdout, stderr = p_ip.communicate()

    r_subnet_rules = []
    for line in stdout. decode().splitlines():
        res = re.match('(.* ){2}eth[0-9](?! $).*', line)

        if res is not None:
            subnet_rule = res.group(0)
            l_subnet_rule = subnet_rule.split()
            device = l_subnet_rule[2]
            ip = l_subnet_rule[-1]

            r_subnet_rules.append(
                {
                    'device': device,
                    'ip': ip,
                    'subnet_rule': subnet_rule
                }
            )

    r_default_route = ''
    for line in stdout.decode().splitlines():
        res = re.match('default .*', line)

        if res is not None:
            r_default_route = res.group(0)

            break

    with open('/etc/rc.local', 'a') as f:
        f.write('#!/bin/bash\n\n')

        rule_index = 128
        default_route_device = ''
        for rule in r_subnet_rules:
            default_route = re.sub('eth.', rule['device'], r_default_route)

            f.write('ip rule add from {0} table {1}\n'.format(rule['ip'], rule_index))
            f.write('ip r add {0} table {1}\n'.format(default_route, rule_index))
            f.write('ip r add {0} table {1}\n\n'.format(rule['subnet_rule'], rule_index))

            if rule['ip'] == NIC_IP:
                default_route_device = rule['device']

            rule_index += 1

        default_route = re.sub('eth.', default_route_device, r_default_route)

        f.write('ip r del default\n')
        f.write('ip r add {0}\n\n'.format(default_route))
        f.write('exit 0\n')

        f.flush()
        f.close()

    os.chmod('/etc/rc.local', 0o0755)

    p_rc_local = subprocess.Popen('/etc/rc.local'.split(), stdout=subprocess.PIPE)
    stdout, stderr = p_rc_local.communicate()




if __name__ == '__main__':
    boto3.setup_default_session(region_name=AWS_REGION)

    # uses: DEPLOY_UUID, TAG_KEY
    attach_eni_ids()
    # uses: MOUNT_POINT, SERVICE_NAME, DEPLOY_UUID, TAG_KEY
    attach_ebs()
    # uses: NIC_IP
    change_default_route()

[((444, 465), 'boto3.resource', 'boto3.resource', (['"""ec2"""'], {}), "('ec2')\n", (458, 465), False, 'import boto3\n'), ((782, 801), 'boto3.client', 'boto3.client', (['"""ec2"""'], {}), "('ec2')\n", (794, 801), False, 'import boto3\n'), ((814, 835), 'boto3.resource', 'boto3.resource', (['"""ec2"""'], {}), "('ec2')\n", (828, 835), False, 'import boto3\n'), ((1197, 1218), 'boto3.resource', 'boto3.resource', (['"""ec2"""'], {}), "('ec2')\n", (1211, 1218), False, 'import boto3\n'), ((1565, 1584), 'boto3.client', 'boto3.client', (['"""ec2"""'], {}), "('ec2')\n", (1577, 1584), False, 'import boto3\n'), ((3441, 3449), 'time.sleep', 'sleep', (['(3)'], {}), '(3)\n', (3446, 3449), False, 'from time import sleep\n'), ((6408, 6438), 'os.chmod', 'os.chmod', (['"""/etc/rc.local"""', '(493)'], {}), "('/etc/rc.local', 493)\n", (6416, 6438), False, 'import os\n'), ((6607, 6658), 'boto3.setup_default_session', 'boto3.setup_default_session', ([], {'region_name': 'AWS_REGION'}), '(region_name=AWS_REGION)\n', (6634, 6658), False, 'import boto3\n'), ((848, 915), 'requests.get', 'requests.get', (['"""http://169.254.169.254/latest/meta-data/instance-id"""'], {}), "('http://169.254.169.254/latest/meta-data/instance-id')\n", (860, 915), False, 'import requests\n'), ((1597, 1664), 'requests.get', 'requests.get', (['"""http://169.254.169.254/latest/meta-data/instance-id"""'], {}), "('http://169.254.169.254/latest/meta-data/instance-id')\n", (1609, 1664), False, 'import requests\n'), ((1977, 1985), 'time.sleep', 'sleep', (['(3)'], {}), '(3)\n', (1982, 1985), False, 'from time import sleep\n'), ((2801, 2819), 'glob.glob', 'glob.glob', (['pattern'], {}), '(pattern)\n', (2810, 2819), False, 'import glob\n'), ((3523, 3541), 'glob.glob', 'glob.glob', (['pattern'], {}), '(pattern)\n', (3532, 3541), False, 'import glob\n'), ((3883, 3901), 'glob.glob', 'glob.glob', (['pattern'], {}), '(pattern)\n', (3892, 3901), False, 'import glob\n'), ((4241, 4249), 'time.sleep', 'sleep', (['(1)'], {}), '(1)\n', (4246, 4249), False, 'from time import sleep\n'), ((4892, 4934), 're.match', 're.match', (['"""(.* ){2}eth[0-9](?! $).*"""', 'line'], {}), "('(.* ){2}eth[0-9](?! $).*', line)\n", (4900, 4934), False, 'import re\n'), ((5411, 5439), 're.match', 're.match', (['"""default .*"""', 'line'], {}), "('default .*', line)\n", (5419, 5439), False, 'import re\n'), ((6187, 6240), 're.sub', 're.sub', (['"""eth."""', 'default_route_device', 'r_default_route'], {}), "('eth.', default_route_device, r_default_route)\n", (6193, 6240), False, 'import re\n'), ((2643, 2661), 'glob.glob', 'glob.glob', (['pattern'], {}), '(pattern)\n', (2652, 2661), False, 'import glob\n'), ((5733, 5780), 're.sub', 're.sub', (['"""eth."""', "rule['device']", 'r_default_route'], {}), "('eth.', rule['device'], r_default_route)\n", (5739, 5780), False, 'import re\n')]