#!/usr/bin/env python
#
# October 2016, Glenn F. Matthews
# Copyright (c) 2016 by Cisco Systems, Inc.
# All rights reserved.
#
# Wrapper script for the virsh command:
# 1) When creating third-party LXCs, move them to tp_app cgroup if necessary
# 2) Forward libvirtd connection to XRNNS if necessary

# XR doesn't provide argparse at present. optparse is deprecated but available.
import optparse
import os
import shutil
import sys
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import ParseError

def check_and_edit_partition(file_path):
    """Check if the given XML file places the VM in the correct cgroup.
    
    If not, fix it so that it does.
    """
    if not os.path.exists(file_path):
        return

    try:
        tree = ET.parse(file_path)
    except ParseError:
        return

    # Register 'lxc' as the preferred shorthand for the LXC namespace.
    # If we don't do this, when we write the modified XML back out,
    # it will get a default shorthand value of "ns0" - ugh. :-)
    ET.register_namespace("lxc", "http://libvirt.org/schemas/domain/lxc/1.0")

    domain = tree.getroot()

    modified = False
    # LXCs other than the system ones need to have:
    # <domain> -> <resource> -> <partition> -> "/machine/tp_app/lxc"

    resource = domain.find("./resource")
    if resource is None:
        resource = ET.SubElement(domain, "resource")
        modified = True

    partition = resource.find("./partition")
    if partition is None:
        partition = ET.SubElement(resource, "partition")
        modified = True

    if partition.text == "/machine" or not partition.text:
        # Don't allow non-system LXCs to belong to the default partition
        partition.text = "/machine/tp_app/lxc"
        modified = True
    else:
        # Some other non-default partition is already set; don't try to change it
        # However, warn the user if they're doing something dangerous:
        if not partition.text.startswith("/machine/tp_app"):
            sys.stderr.write(
"""WARNING: Unrecognized/unsupported resource partition:
{0}
Using resources beyond those allocated to the /machine/tp_app partition
may have unexpected and undesirable effects on system performance.
""".format(partition.text))

        pass

    if modified:
        backup = file_path + ".original"
        shutil.copy(file_path, backup)
        tree.write(file_path)
        sys.stderr.write(
"""XML file {0}
has been automatically modified to better fit the IOS XR environment.
A copy has been saved to {1}

To avoid this message in the future, please ensure that your XML file has
an appropriate resource partition declaration:
<domain>
  ...
  <resource>
    <partition>/machine/tp_app/lxc</partition>
  </resource>
  ...
</domain>
""".format(file_path, backup))

class IgnoreUnknownOptionParser(optparse.OptionParser):
    """OptionParser, but discard unrecognized options instead of erroring."""
    def _process_args(self, largs, rargs, values):
        while rargs:
            try:
                optparse.OptionParser._process_args(self, largs, rargs, values)
            except (optparse.BadOptionError,
                    optparse.AmbiguousOptionError) as e:
                pass

def handle_args(arguments):
    """Comb through the given arguments to look for anything interesting.

    Specifically, look for:
    (virsh) [options] create [options] file.xml [options].
    (virsh) [options] define [options] file.xml [options].
    None of the other options matter to us at this point so we can ignore them.
    """
    parser = IgnoreUnknownOptionParser(add_help_option=False)
    # Add arguments that we may encounter that take parameters
    # 'virsh' generic arguments with parameters:
    parser.add_option('-c', '--connect')
    parser.add_option('-d', '--debug')
    parser.add_option('-e', '--escape')
    parser.add_option('-k', '--keepalive-interval')
    parser.add_option('-K', '--keepalive-count')
    parser.add_option('-l', '--log')
    # 'virsh create' arguments with parameters:
    parser.add_option('--pass-fds', dest="pass_fds")
    # all other options before or after the 'create' keyword can be safely
    # ignored, as they do not take arguments
    (options, args) = parser.parse_args()

    if len(args) >= 2 and (args[0] == 'create' or args[0] == 'define'):
        check_and_edit_partition(args[1])

if __name__ == "__main__":
    handle_args(sys.argv[1:])
