#!/usr/bin/env python

# Modified from
# https://github.com/f5devcentral/f5-aws-migrate/blob/master/f5-aws-migrate.py
# Assumptions: Script will be run on a BIGI-IP itself

import os
import sys
import re
import shutil
import tarfile
from tempfile import mkstemp
from optparse import OptionParser

def get_hostname(config_file):
    hostname = ""
    global_settings_token_found = False
    hostname_token_line = 0
    with open(config_file, 'r') as f:
        for line in f:
            if global_settings_token_found:
                match = re.search('\s+hostname (.+)', line)
                if match:
                    hostname = match.group(1)
                    return hostname
            else:
                match = re.search('sys global-settings', line)
                if match:
                    global_settings_token_found = True

def get_ip(config_file):
    # 1nic Autoscale IPv4 Specific
    # Looking for entry in bigip_base.conf:
    # net self /Common/10.0.11.241/24 {
    #     address 10.0.11.241/24
    #     allow-service {
    #         default
    #     }
    #     traffic-group /Common/traffic-group-local-only
    #     vlan /Common/external
    # }
    ip = ""
    line_number = 0
    net_self_token_line = 0
    with open(config_file, 'r') as f:
        for line in f:
            match = re.search('net self', line)
            if match:
                net_self_token_line = line_number
            # Search for address in that token
            match = re.search('    address (.+)/(.+)', line)
            if match and (line_number - net_self_token_line == 1):
                ip = match.group(1)
                mask = match.group(2)
            line_number += 1
    # could do a little more validation here to make sure not another self-ip configured by accident
    # octets match management ip, hostname, etc.
    return ip


def get_local_ip():
    # 1nic Autoscale IPv4 Specific
    # "privateIp entry in /shared/vadc/aws/iid-document
    # is probably most reliable
    # otherwise, getting it from dhcp
    # searching for field
    # "fixed-address 10.0.1.219;"
    ip = ""
    with open('/var/lib/dhclient/dhclient.leases', 'r') as f:
        for line in f:
            match = re.search('fixed-address (.+);', line)
            if match:
                ip = match.group(1)

    return ip


def get_gateway(config_file):
    # 1nic Autoscale IPv4 Specific
    # searching for entry
    # net route /LOCAL_ONLY/default {
    #       gw 10.0.1.1
    #       network default
    # }
    gateway = ""
    with open(config_file, 'r') as f:
        for line in f:
            match = re.search('    gw (.+)', line)
            if match:
                gateway = match.group(1)

    return gateway


def get_local_gateway():
    # 1nic Autoscale IPv4 Specific
    # searching for field
    # "option routers 10.0.1.1;"
    gateway = ""
    with open('/var/lib/dhclient/dhclient.leases', 'r') as f:
        for line in f:
            match = re.search('option routers (.+);', line)
            if match:
                gateway = match.group(1)

    return gateway


def replace(source_file_path, pattern, substring):
    fh, target_file_path = mkstemp()
    with open(target_file_path, 'w') as target_file:
        with open(source_file_path, 'r') as source_file:
            for line in source_file:
                if "-" in pattern:
                    newline = re.sub("(%s)$" %(pattern),substring,line)
                    newline = re.sub("(%s)\." %(pattern),substring+'.',newline)
                    target_file.write(newline)
                else:    
                    newline = re.sub("(%s)$" %(pattern),substring,line)
                    newline = re.sub("(%s)\/" %(pattern),substring+'/',newline)
                    target_file.write(newline)
    os.remove(source_file_path)
    shutil.move(target_file_path, source_file_path)

def removeFiles(dir, pattern):
    if os.path.exists(dir):
        for f in os.listdir(dir):
            if re.search(pattern, f):
                os.remove(os.path.join(dir, f))

def main():

    parser = OptionParser()
    parser.add_option("--debug-level", action="store", type="int", dest="debug_level", default=0, help="debug level print debug (0-9)")
    parser.add_option("--cloud-provider", action="store", type="string", dest="cloud_provider", default="" , help="Cloud being utilized, azure or aws, etc...")
    parser.add_option("--original-ucs", action="store", type="string", dest="original_ucs", help="Original UCS file name")
    parser.add_option("--updated-ucs", action="store", type="string", dest="updated_ucs", default="updated.ucs", help="Modified UCS file name")
    parser.add_option("--extract-directory", action="store", type="string", dest="extract_dir", default="ucs_extract_dir", help="name of directory to extract to")
    parser.add_option("--original-ucs-ip", action="store", type="string", dest="original_ucs_ip", help="ip in original ucs")
    parser.add_option("--original-ucs-gateway", action="store", type="string", dest="original_ucs_gateway", help="gateway in original ucs")
    parser.add_option("--dest-ip", action="store", type="string", dest="dest_ip", help="ip of destination instance")
    parser.add_option("--dest-gateway", action="store", type="string", dest="dest_gateway", help="gateway of destination instance")
    (options, args) = parser.parse_args()

    # Set variables from options
    debug_level = options.debug_level
    cloud_provider = options.cloud_provider
    original_ucs = options.original_ucs
    updated_ucs = options.updated_ucs
    extract_ucs_dir = options.extract_dir
    original_ucs_ip = options.original_ucs_ip
    original_ucs_gateway = options.original_ucs_gateway
    dest_ip = options.dest_ip
    dest_gateway = options.dest_gateway

    if not original_ucs or not cloud_provider:
        print "Usage: "
        print ("    ./%s --cloud-provider <aws | azure | gce> --original-ucs <ucs_filename>" % (sys.argv[0]))
        print ("    ./%s --cloud-provider <aws | azure | gce> --original-ucs <ucs_filename> --updated-ucs <ucs_filename>" % (sys.argv[0]))
        print "ex. "
        print ("    ./%s --cloud-provider aws --original-ucs original.ucs --updated-ucs updated.ucs" % (sys.argv[0]))
        sys.exit()

    # Open files
    tar_original = tarfile.open(original_ucs, "r:gz")
    tar_updated = tarfile.open(updated_ucs, "w:gz")

    try:
        tar_original.extractall(path=extract_ucs_dir)
    except IOError as e:
        sys.exit("I/O Error({0}): {1} - If the latest UCS file is corrupted and/or the current primary host is stuck at BECOMING_PRIMARY state, the corrupted UCS file needs to be deleted as well as primary host needs to be terminated to allow re-elect the new primary.'".format(e.errno, e.strerror))


    bigip_base_file = "/config/bigip_base.conf"
    if cloud_provider == 'aws':
        gateway_file = "/config/partitions/LOCAL_ONLY/bigip.conf"
    else:
        gateway_file = "/config/bigip.conf"

    # Grab instance's hostname from UCS
    dest_hostname = get_hostname(bigip_base_file)
    original_hostname = get_hostname(extract_ucs_dir + bigip_base_file)

    # Grab instance's IP from UCS or local config file
    if not original_ucs_ip:
        original_ucs_ip = get_ip(extract_ucs_dir + bigip_base_file)
    if not original_ucs_gateway:
        original_ucs_gateway = get_gateway(extract_ucs_dir + gateway_file)
    if not dest_ip:
        if os.path.isfile(gateway_file):
            dest_ip = get_ip(bigip_base_file)
        else:
            dest_ip = get_local_ip()
    if not dest_gateway:
        if os.path.isfile(gateway_file):
            dest_gateway = get_gateway(gateway_file)
        else:
            dest_gateway = get_local_gateway()

    if debug_level > 0:
        print "original_hostname: " + original_hostname
        print "dest_hostname: " + dest_hostname
        print "original_ucs_ip: " + original_ucs_ip
        print "original_ucs_gateway: " + original_ucs_gateway
        print "dest_ip: " + dest_ip
        print "dest_gateway: " + dest_gateway

    # Fix string version of addresses with "-". ex. ip-10-0-11-151
    original_ucs_ip_str = original_ucs_ip.replace(".", "-")
    dest_ip_str = dest_ip.replace(".", "-")

    if debug_level > 0:
        print "original_ucs_ip_str: " + original_ucs_ip_str
        print "dest_ip_str: " + dest_ip_str

    files_to_update = [
                        "/config/bigip_base.conf",
                        "/config/bigip.conf",
                        "/config/BigDB.dat",
                        "/SPEC-Manifest"
                      ]

    # Replace Gateway
    replace(extract_ucs_dir + gateway_file, original_ucs_gateway, dest_gateway)

    # Replace hostname, IP, String Versions in other files
    for f in files_to_update:
        filename = extract_ucs_dir + f
        if debug_level > 0:
            print "updating : " + filename
        replace(filename, original_ucs_ip, dest_ip)
        replace(filename, original_ucs_ip_str, dest_ip_str)
        if original_hostname and dest_hostname:
            replace(filename, original_hostname, dest_hostname)

    # Remove the cloud directory created by the template so we don't overwrite new code with old
    path_to_exclude = "/config/cloud/" + cloud_provider
    shutil.rmtree(extract_ucs_dir + path_to_exclude, ignore_errors=True)

    # Remove the cloud-libs private key info as you can't load a ucs w/ a passphrase
    shutil.rmtree(extract_ucs_dir + '/config/partitions/CloudLibsAutoscale', ignore_errors=True)
    shutil.rmtree(extract_ucs_dir + '/var/tmp/filestore_temp/files_d/CloudLibsAutoscale_d', ignore_errors=True)
    shutil.rmtree(extract_ucs_dir + '/config/partitions/CloudLibsLocal', ignore_errors=True)
    shutil.rmtree(extract_ucs_dir + '/var/tmp/filestore_temp/files_d/CloudLibsLocal_d', ignore_errors=True)

    # Remove the f5-cloud-libs local public keys as they won't match any private keys
    removeFiles(extract_ucs_dir + '/config/cloud/keys', 'cloudLocal*')

    # remove the dynad private key
    os.system("sed -i '/sys dynad key {/ { N ; /\\n[[:space:]]\+key[[:space:]]*\$M\$[^\\n]*/ { N;   /\\n[[:space:]]*}/ { d } } }' " + extract_ucs_dir + "/config/bigip_base.conf")

    tar_updated.add(extract_ucs_dir, arcname='')

    tar_original.close()
    tar_updated.close()

    shutil.rmtree(extract_ucs_dir, ignore_errors=False, onerror=None)

    print "UCS Update Complete"
    print "Load UCS with command below:"
    print "   tmsh load /sys ucs " + os.path.abspath(updated_ucs) + " no-license"

    # Leverage cfn-signal here
    # http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/cfn-signal.html

if __name__ == "__main__":
    main()