#!/usr/bin/env python3
# Copyright (c) 2012 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.

import sys

from xcp.bootloader import *

labels = ['xe', 'xe-serial']


def _get_menus(boot_loader):
    return [x for x in boot_loader.menu_order if x in labels]

def _match_arg(key, arg):
    return arg.startswith(key + "=") or arg == key

def _filter_args(keys, args, match=True):
    if match:
        return [arg for arg in args for key in keys if _match_arg(key, arg)]
    return [arg for arg in args for key in keys if not _match_arg(key, arg)]

def get_xen_args(b, keys):
    """Get xen args matching ks"""
    return _filter_args(keys, b.menu[b.default].getHypervisorArgs())

def replace_xen_args(b, nv):
    """Replace existing xen args with nv"""
    keys = [x.split('=')[0] for x in nv]
    for m in _get_menus(b):
        args = _filter_args(keys, b.menu[m].getHypervisorArgs(), match=False)
        b.menu[m].setHypervisorArgs(args + nv)

def delete_xen_args(b, keys):
    """Delete existing xen args match ks"""
    for m in _get_menus(b):
        args = _filter_args(keys, b.menu[m].getHypervisorArgs(), match=False)
        b.menu[m].setHypervisorArgs(args)

def get_dom0_args(b, keys):
    """Get dom0 kernel args"""
    return _filter_args(keys, b.menu[b.default].getKernelArgs())

def replace_dom0_args(b, nv):
    """Replace dom0 kernel args with nv"""
    keys = [x.split('=')[0] for x in nv]
    for m in _get_menus(b):
        args = _filter_args(keys, b.menu[m].getKernelArgs(), match=False)
        b.menu[m].setKernelArgs(args + nv)

def delete_dom0_args(b, keys):
    """Delete dom0 kernel args matching keys"""
    for m in _get_menus(b):
        args = _filter_args(keys, b.menu[m].getKernelArgs(), match=False)
        b.menu[m].setKernelArgs(args)

def list_xen_cpuid_masks(b):
    return _filter_args(["cpuid_mask"], b.menu[b.default].getHypervisorArgs())


if __name__ == "__main__":
    modified = False
    opt = None
    args = []
    operations = []

    b = Bootloader.loadExisting()

    for arg in sys.argv[1:]:
        if arg.startswith('--') and opt and len(args):
            operations.append((opt, args))
            args = []

        if arg == '--list-cpuid-masks':
            operations.append((arg, None))
        elif arg in ['--set-cpuid-masks', '--delete-cpuid-masks', '--get-xen', '--get-dom0', 
                     '--set-xen', '--set-dom0', '--delete-xen', '--delete-dom0']:
            opt = arg
        if not arg.startswith('--'):
            args.append(arg)
    if opt and len(args):
        operations.append((opt, args))

    for op in operations:
        opt, args = op

        if opt == '--list-cpuid-masks':
            print('\n'.join(list_xen_cpuid_masks(b)))
        elif opt == '--get-xen':
            print('\n'.join((get_xen_args(b, args))))
        elif opt == '--get-dom0':
            print('\n'.join((get_dom0_args(b, args))))
        elif opt in ['--set-cpuid-masks', '--set-xen']:
            replace_xen_args(b, args)
            modified = True
        elif opt in ['--delete-cpuid-masks', '--delete-xen']:
            delete_xen_args(b, args)
            modified = True
        elif opt == '--set-dom0':
            replace_dom0_args(b, args)
            modified = True
        elif opt == '--delete-dom0':
            delete_dom0_args(b, args)
            modified = True

    if modified:
        b.commit()

    sys.exit(0)
