#!/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()

[((21, 10, 21, 31), 'boto3.resource', 'boto3.resource', ({(21, 25, 21, 30): '"""ec2"""'}, {}), "('ec2')", False, 'import boto3\n'), ((34, 12, 34, 31), 'boto3.client', 'boto3.client', ({(34, 25, 34, 30): '"""ec2"""'}, {}), "('ec2')", False, 'import boto3\n'), ((35, 12, 35, 33), 'boto3.resource', 'boto3.resource', ({(35, 27, 35, 32): '"""ec2"""'}, {}), "('ec2')", False, 'import boto3\n'), ((46, 10, 46, 31), 'boto3.resource', 'boto3.resource', ({(46, 25, 46, 30): '"""ec2"""'}, {}), "('ec2')", False, 'import boto3\n'), ((60, 10, 60, 29), 'boto3.client', 'boto3.client', ({(60, 23, 60, 28): '"""ec2"""'}, {}), "('ec2')", False, 'import boto3\n'), ((126, 4, 126, 12), 'time.sleep', 'sleep', ({(126, 10, 126, 11): '(3)'}, {}), '(3)', False, 'from time import sleep\n'), ((235, 4, 235, 37), 'os.chmod', 'os.chmod', ({(235, 13, 235, 28): '"""/etc/rc.local"""', (235, 30, 235, 36): '(493)'}, {}), "('/etc/rc.local', 493)", False, 'import os\n'), ((244, 4, 244, 55), 'boto3.setup_default_session', 'boto3.setup_default_session', (), '', False, 'import boto3\n'), ((37, 11, 37, 78), 'requests.get', 'requests.get', ({(37, 24, 37, 77): '"""http://169.254.169.254/latest/meta-data/instance-id"""'}, {}), "('http://169.254.169.254/latest/meta-data/instance-id')", False, 'import requests\n'), ((62, 11, 62, 78), 'requests.get', 'requests.get', ({(62, 24, 62, 77): '"""http://169.254.169.254/latest/meta-data/instance-id"""'}, {}), "('http://169.254.169.254/latest/meta-data/instance-id')", False, 'import requests\n'), ((74, 8, 74, 16), 'time.sleep', 'sleep', ({(74, 14, 74, 15): '(3)'}, {}), '(3)', False, 'from time import sleep\n'), ((102, 13, 102, 31), 'glob.glob', 'glob.glob', ({(102, 23, 102, 30): 'pattern'}, {}), '(pattern)', False, 'import glob\n'), ((129, 16, 129, 34), 'glob.glob', 'glob.glob', ({(129, 26, 129, 33): 'pattern'}, {}), '(pattern)', False, 'import glob\n'), ((140, 16, 140, 34), 'glob.glob', 'glob.glob', ({(140, 26, 140, 33): 'pattern'}, {}), '(pattern)', False, 'import glob\n'), ((160, 8, 160, 16), 'time.sleep', 'sleep', ({(160, 14, 160, 15): '(1)'}, {}), '(1)', False, 'from time import sleep\n'), ((184, 14, 184, 56), 're.match', 're.match', ({(184, 23, 184, 49): '"""(.* ){2}eth[0-9](?! $).*"""', (184, 51, 184, 55): 'line'}, {}), "('(.* ){2}eth[0-9](?! $).*', line)", False, 'import re\n'), ((202, 14, 202, 42), 're.match', 're.match', ({(202, 23, 202, 35): '"""default .*"""', (202, 37, 202, 41): 'line'}, {}), "('default .*', line)", False, 'import re\n'), ((226, 24, 226, 77), 're.sub', 're.sub', ({(226, 31, 226, 37): '"""eth."""', (226, 39, 226, 59): 'default_route_device', (226, 61, 226, 76): 'r_default_route'}, {}), "('eth.', default_route_device, r_default_route)", False, 'import re\n'), ((96, 20, 96, 38), 'glob.glob', 'glob.glob', ({(96, 30, 96, 37): 'pattern'}, {}), '(pattern)', False, 'import glob\n'), ((215, 28, 215, 75), 're.sub', 're.sub', ({(215, 35, 215, 41): '"""eth."""', (215, 43, 215, 57): "rule['device']", (215, 59, 215, 74): 'r_default_route'}, {}), "('eth.', rule['device'], r_default_route)", False, 'import re\n')]