#!/usr/bin/env python3

"""
Copyright (C) 2006-2009 Citrix Systems Inc.
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published
by the Free Software Foundation; version 2.1 only. with the special
exception on linking described in file LICENSE.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.
"""
from __future__ import print_function

import os
import re
import sys
from contextlib import contextmanager
from optparse import OptionParser

pool_conf = '/etc/xensource/pool.conf'
inventory_file = '/etc/xensource-inventory'
management_conf = '/etc/firstboot.d/data/management.conf'
network_reset = '/var/tmp/network-reset'
RENAME_SCRIPT = '/etc/sysconfig/network-scripts/interface-rename.py'
rename_script_exists = os.path.exists(RENAME_SCRIPT)

@contextmanager
def fsync_write(filename):
    """Context manager that writes to a file and fsyncs it after writing."""

    with open(filename, "w", encoding="utf-8") as file:
        try:  # Run the context, ignoring exceptions:
            yield file
        finally:
            file.flush()  # Flush the file buffer to the OS
            os.fsync(file.fileno())  # Ask the OS to write the file to disk


def read_dict_file(fname):
    f = open(fname, 'r')
    d = {}
    for l in f.readlines():
        kv = l.split('=')
        d[kv[0].strip()] = kv[1].strip().strip("'")
    return d

def read_inventory():
    return read_dict_file(inventory_file)

def read_management_conf():
    return read_dict_file(management_conf)


def write_inventory(inventory_dict):
    with fsync_write(inventory_file) as file:
        for k in inventory_dict:
            file.write(k + "='" + inventory_dict[k] + "'\n")


def valid_vlan(vlan):
    if not re.match(r"^\d+$", vlan):
        return False
    if int(vlan)<0 or int(vlan)>4094:
        return False
    return True

def get_bridge_name(device, vlan):
    # Construct bridge name for management interface based on convention
    # NOTE: Only correct when interface-rename script exists
    if vlan != None:
        return 'xentemp'
    m = re.match(r'^eth(\d+)$', device)
    if m:
        return 'xenbr' + m.group(1)
    return 'br' + device

if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option("-m", "--master", help="Master's address", dest="address", default=None)
    parser.add_option("--device", help="Device name of new management interface", dest="device", default=None)
    parser.add_option("--mode", help='IP configuration mode for new management interface: "none", "dhcp" or "static" (default is dhcp)', dest="mode", default="dhcp")
    parser.add_option("--mode-v6", help='IPv6 configuration mode for new management interface: "none", "dhcp", "autoconf" or "static" (default is none)', dest="mode_v6", default="none")
    parser.add_option("--novlan", help="no vlan is used for new management interface", dest="novlan", action="store_const", const=True, default=False)
    parser.add_option("--vlan", help="vlanID for new management interface to be on vlan network", dest="vlan", default=None)
    parser.add_option("--ip", help="IP address for new management interface", dest="ip", default='')
    parser.add_option("--ipv6", help="IPv6 address (CIDR format) for new management interface", dest="ipv6", default='')
    parser.add_option("--netmask", help="Netmask for new management interface", dest="netmask", default='')
    parser.add_option("--gateway", help="Gateway for new management interface", dest="gateway", default='')
    parser.add_option("--gateway-v6", help="IPv6 Gateway for new management interface", dest="gateway_v6", default='')
    parser.add_option("--dns", help="DNS server for new management interface", dest="dns", default='')
    (options, args) = parser.parse_args()

    # Determine pool role
    try:
        f = open(pool_conf, 'r')
        try:
            l = f.readline()
            ls = l.split(':', maxsplit=1)
            if ls[0].strip() == 'master':
                master = True
                address = 'localhost'
            else:
                master = False
                if options.address == None:
                    address = ls[1].strip()
                else:
                    address = options.address
        finally:
            f.close()
    except Exception:
        master = None
        address = ""

    # Get the management device from the firstboot data if not specified by the user
    if options.device == None:
        try:
            conf = read_management_conf()
            device = conf['LABEL']
        except:
            print("Could not figure out which interface should become the management interface. \
                Please specify one using the --device option.")
            sys.exit(1)
    else:
        device = options.device

    # Get the VLAN if provided in the firstboot data and not specified by the user
    vlan = None
    if options.vlan:
        if options.novlan:
            parser.error('"--vlan <vlanId>" and "--novlan" should not be used together')
            sys.exit(1)
        if not valid_vlan(options.vlan):
            print("VLAN tag you gave was invalid, It must be between 0 and 4094")
            sys.exit(1)
        vlan = options.vlan
    elif not options.novlan:
        try:
            conf = read_management_conf()
            vlan = conf['VLAN']
        except KeyError:
            pass

    # Determine IP configuration for management interface
    options.mode = options.mode.lower()
    if options.mode not in ["none", "dhcp", "static"]:
        parser.error('mode should be either "none", "dhcp" or "static"')
        sys.exit(1)

    options.mode_v6 = options.mode_v6.lower()
    if options.mode not in ["none", "autoconf", "dhcp", "static"]:
        parser.error('mode-v6 should be either "none", "autoconf", "dhcp" or "static"')
        sys.exit(1)

    if options.mode == "none" and options.mode_v6 == "none":
        parser.error("Either mode or mode-v6 must be not 'none'")
        sys.exit(1)

    if options.mode == 'static' and (options.ip == '' or options.netmask == ''):
        parser.error("if static IP mode is selected, an IP address and netmask need to be specified")
        sys.exit(1)

    if options.mode_v6 == 'static':
        if options.ipv6 == '':
            parser.error("if static IPv6 mode is selected, an IPv6 address needs to be specified")
            sys.exit(1)
        if options.ipv6.find('/') == -1:
            parser.error("Invalid format: IPv6 must be specified with CIDR format: <IPv6>/<prefix>")
            sys.exit(1)

    # Warn user
    if not os.access('/tmp/fist_network_reset_no_warning', os.F_OK):
        configuration = []
        configuration.append("Management interface:   " + device)
        configuration.append("IP configuration mode:  " + options.mode)
        configuration.append("IPv6 configuration mode:" + options.mode_v6)
        if vlan != None:
            configuration.append("Vlan:                   " + vlan)
        if options.mode == "static":
            configuration.append("IP address:             " + options.ip)
            configuration.append("Netmask:                " + options.netmask)
        if options.mode_v6 == "static":
            configuration.append("IPv6/CIDR:              " + options.ipv6)
        if options.gateway != '':
            configuration.append("Gateway:                " + options.gateway)
        if options.gateway_v6 != '':
            configuration.append("IPv6 gateway:           " + options.gateway_v6)
        if options.dns != '':
            configuration.append("DNS server(s):          " + options.dns)
        if master == False:
            configuration.append("Pool master's address:  " + address)
        warning = """----------------------------------------------------------------------
!! WARNING !!

This command will reboot the host and reset its network configuration.
Any running VMs will be forcefully shutdown.

Before completing this command:
- Where possible, cleanly shutdown all VMs running on this host.
- Disable HA if this host is part of a resource pool with HA enabled.
----------------------------------------------------------------------

Your network will be re-configured as follows:\n\n"""
        confirmation = """\n\nIf you want to change any of the above settings, type 'no' and re-run
the command with appropriate arguments (use --help for a list of options).

Type 'yes' to continue.
Type 'no' to cancel.
"""
        res = input(warning + '\n'.join(configuration) + confirmation)
        if res != 'yes':
            sys.exit(1)

    # Update master's IP, if needed and given
    if master == False and options.address != None:
        print("Setting master's ip (" + address + ")...")
        with fsync_write(pool_conf) as f:
            f.write('slave:' + address)

    # Ensure xapi is not running
    print("Stopping xapi...")
    os.system('systemctl stop xapi >/dev/null 2>/dev/null')

    # Reconfigure new management interface
    print("Reconfiguring " + device + "...")
    os.system('systemctl stop xcp-networkd >/dev/null 2>/dev/null')
    try:
        os.remove('/var/lib/xcp/networkd.db')
    except Exception as e:
        print('Warning: Failed to delete networkd.db.\n%s' % e)

    # Update interfaces in inventory file
    print('Updating inventory file...')
    inventory = read_inventory()
    # If rename script does not exist, needn't to set MANAGEMENT_INTERFACE in inventory file
    # Networkd will handle it while replacing the rename script to sort interfaces
    bridge = '' if not rename_script_exists else get_bridge_name(device, vlan)
    inventory['MANAGEMENT_INTERFACE'] = bridge
    inventory['CURRENT_INTERFACES'] = ''
    write_inventory(inventory)

    # Rewrite firstboot management.conf file, which will be picked it by xcp-networkd on restart (if used)
    is_static = False
    with fsync_write(management_conf) as f:
        f.write("LABEL='" + device + "'\n")
        if options.mode != "none":
            f.write("MODE='" + options.mode + "'\n")
        if options.mode_v6 != "none":
            f.write("MODEV6='" + options.mode_v6 + "'\n")
        if vlan != None:
            f.write("VLAN='" + vlan + "'\n")
        if options.mode == 'static':
            is_static = True
            f.write("IP='" + options.ip + "'\n")
            f.write("NETMASK='" + options.netmask + "'\n")
            if options.gateway != '':
                f.write("GATEWAY='" + options.gateway + "'\n")
        if options.mode_v6 == "static":
            is_static = True
            f.write("IPv6='" + options.ipv6 + "'\n")
            if options.gateway_v6 != '':
                f.write("IPv6_GATEWAY='" + options.gateway_v6 + "'\n")
        if is_static and options.dns != '':
            f.write("DNS='" + options.dns + "'\n")

    # Write trigger file for XAPI to continue the network reset on startup
    with fsync_write(network_reset) as f:
        f.write('DEVICE=' + device + '\n')
        if options.mode != "none":
            f.write('MODE=' + options.mode + '\n')
        if options.mode_v6 != "none":
            f.write('MODE_V6=' + options.mode_v6 + '\n')
        if vlan != None:
            f.write('VLAN=' + vlan + '\n')
        if options.mode == 'static':
            f.write('IP=' + options.ip + '\n')
            f.write('NETMASK=' + options.netmask + '\n')
            if options.gateway != '':
                f.write('GATEWAY=' + options.gateway + '\n')
        if options.mode_v6 == "static":
            f.write('IPV6=' + options.ipv6 + '\n')
            if options.gateway_v6 != '':
                f.write('GATEWAY_V6=' + options.gateway_v6 + '\n')
        if is_static and options.dns != '':
            f.write('DNS=' + options.dns + '\n')
    if rename_script_exists:
        # Reset the domain 0 network interface naming configuration
        # back to a fresh-install state for the currently-installed
        # hardware.
        os.system(f"{RENAME_SCRIPT} --reset-to-install")

    # Reboot
    os.system("mount -o remount,rw / && reboot -f")
