# (C) Copyright IBM Corp. 2005
# Copyright (C) 2004 Mike Wray
# Copyright (c) 2005 XenSource Ltd
#
# Authors:
#     Sean Dague <sean at dague dot net>
#     Mike Wray <mike dot wray at hp dot com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of version 2.1 of the GNU Lesser General Public
# License as published by the Free Software Foundation.
#
# This library 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.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""Grand unified management application for Xen.
"""
import os
import os.path
import sys
import re
import getopt
import socket
import warnings
warnings.filterwarnings('ignore', category=FutureWarning)

import xen.xend.XendError
import xen.xend.XendProtocol

from xen.xend import PrettyPrint
from xen.xend import sxp
from xen.xm.opts import *

import console


# getopt.gnu_getopt is better, but only exists in Python 2.3+.  Use
# getopt.getopt if gnu_getopt is not available.  This will mean that options
# may only be specified before positional arguments.
if not hasattr(getopt, 'gnu_getopt'):
    getopt.gnu_getopt = getopt.getopt


# Strings for shorthelp
console_help = "console <DomId>                  Attach to domain DomId's console."
create_help =  """create [-c] <ConfigFile>
               [Name=Value]..       Create a domain based on Config File"""
destroy_help = "destroy <DomId>                  Terminate a domain immediately"
help_help =    "help                             Display this message"
list_help =    "list [--long] [DomId, ...]       List information about domains"
mem_max_help = "mem-max <DomId> <Mem>            Set maximum memory reservation for a domain"
mem_set_help = "mem-set <DomId> <Mem>            Adjust the current memory usage for a domain"
migrate_help = "migrate <DomId> <Host>           Migrate a domain to another machine"
pause_help =   "pause <DomId>                    Pause execution of a domain"
reboot_help =  "reboot <DomId> [-w][-a]          Reboot a domain"
restore_help = "restore <File>                   Create a domain from a saved state file"
save_help =    "save <DomId> <File>              Save domain state (and config) to file"
shutdown_help ="shutdown <DomId> [-w][-a][-R|-H] Shutdown a domain"
top_help =     "top                              Monitor system and domains in real-time"
unpause_help = "unpause <DomId>                  Unpause a paused domain"

help_spacer = """
   """

# Strings for longhelp
sysrq_help =   "sysrq   <DomId> <letter>         Send a sysrq to a domain"
domid_help =   "domid <DomName>                  Converts a domain name to a domain id"
domname_help = "domname <DomId>                  Convert a domain id to a domain name"
vcpu_set_help = """vcpu-set <DomId> <VCPUs>         Set the number of VCPUs for a domain"""
vcpu_list_help = "vcpu-list <DomId>                List the VCPUs for a domain (or all domains)"
vcpu_pin_help = "vcpu-pin <DomId> <VCPU> <CPUs>   Set which cpus a VCPU can use" 
dmesg_help =   "dmesg [-c|--clear]               Read or clear Xen's message buffer"
info_help =    "info                             Get information about the xen host"
rename_help =  "rename <DomId> <New Name>        Rename a domain"
log_help =     "log                              Print the xend log"
sched_bvt_help = """sched-bvt <Parameters>           Set Borrowed Virtual Time scheduler
                                    parameters"""
sched_bvt_ctxallow_help = """sched-bvt-ctxallow <Allow>       Set the BVT scheduler context switch
                                    allowance"""
sched_sedf_help = "sched-sedf [DOM] [OPTIONS]       Show|Set simple EDF parameters\n" + \
"              -p, --period          Relative deadline(ms).\n\
              -s, --slice           Worst-case execution time(ms)\n\
                                    (slice < period).\n\
              -l, --latency         scaled period(ms) in case the domain\n\
                                    is doing heavy I/O.\n\
              -e, --extra           flag (0/1) which controls whether the\n\
                                    domain can run in extra-time\n\
              -w, --weight          mutually exclusive with period/slice and\n\
                                    specifies another way of setting a domain's\n\
                                    cpu period/slice."

block_attach_help = """block-attach <DomId> <BackDev> <FrontDev> <Mode>
                [BackDomId]         Create a new virtual block device"""
block_detach_help = """block-detach  <DomId> <DevId>    Destroy a domain's virtual block device,
                                    where <DevId> may either be the device ID
                                    or the device name as mounted in the guest"""

block_list_help = "block-list <DomId> [--long]      List virtual block devices for a domain"
network_attach_help = """network-attach  <DomID> [script=<script>] [ip=<ip>] [mac=<mac>]
                           [bridge=<bridge>] [backend=<backDomID>]
                                    Create a new virtual network device """
network_detach_help = """network-detach  <DomId> <DevId>  Destroy a domain's virtual network
                                    device, where <DevId> is the device ID."""

network_list_help = "network-list <DomId> [--long]    List virtual network interfaces for a domain"
vnet_list_help = "vnet-list [-l|--long]            list vnets"
vnet_create_help = "vnet-create <config>             create a vnet from a config file"
vnet_delete_help = "vnet-delete <vnetid>             delete a vnet"
vtpm_list_help = "vtpm-list <DomId> [--long]       list virtual TPM devices"

short_command_list = [
    "console",
    "create",
    "destroy",
    "help",
    "list",
    "mem-set",
    "migrate",
    "pause",
    "reboot",
    "restore",
    "save",
    "shutdown",
    "top",
    "unpause",
    "vcpu-set",
    ]

domain_commands = [
    "console",
    "create",
    "destroy",
    "domid",
    "domname",
    "list",
    "mem-max",
    "mem-set",
    "migrate",
    "pause",
    "reboot",
    "rename",
    "restore",
    "save",
    "shutdown",
    "sysrq",
    "top",
    "unpause",
    "vcpu-list",
    "vcpu-pin",
    "vcpu-set",
    ]

host_commands = [
    "dmesg",
    "info",
    "log"
    ]

scheduler_commands = [
    "sched-bvt",
    "sched-bvt-ctxallow",
    "sched-sedf",
    ]

device_commands = [
    "block-attach",
    "block-detach",
    "block-list",
    "network-attach",
    "network-detach",
    "network-list",
    "vtpm-list",
    ]

vnet_commands = [
    "vnet-list",
    "vnet-create",
    "vnet-delete",
    ]

all_commands = (domain_commands + host_commands + scheduler_commands +
                device_commands + vnet_commands)


def commandToHelp(cmd):
    return eval(cmd.replace("-", "_") + "_help")


shorthelp = """Usage: xm <subcommand> [args]
    Control, list, and manipulate Xen guest instances

xm common subcommands:
   """  + help_spacer.join(map(commandToHelp, short_command_list))  + """

<DomName> can be substituted for <DomId> in xm subcommands.

For a complete list of subcommands run 'xm help --long'
For more help on xm see the xm(1) man page
For more help on xm create, see the xmdomain.cfg(5) man page"""

longhelp = """Usage: xm <subcommand> [args]
    Control, list, and manipulate Xen guest instances

xm full list of subcommands:

  Domain Commands:
   """ + help_spacer.join(map(commandToHelp,  domain_commands)) + """

  Xen Host Commands:
   """ + help_spacer.join(map(commandToHelp,  host_commands)) + """

  Scheduler Commands:
   """ + help_spacer.join(map(commandToHelp,  scheduler_commands)) + """

  Virtual Device Commands:
   """  + help_spacer.join(map(commandToHelp, device_commands)) + """

  Vnet commands:
   """ + help_spacer.join(map(commandToHelp,  vnet_commands)) + """

<DomName> can be substituted for <DomId> in xm subcommands.

For a short list of subcommands run 'xm help'
For more help on xm see the xm(1) man page
For more help on xm create, see the xmdomain.cfg(5) man page"""

# array for xm help <command>
help = {
    "--long": longhelp
    }

for command in all_commands:
    # create is handled specially
    if (command != 'create'):
        help[command] = commandToHelp(command)


####################################################################
#
#  Utility functions
#
####################################################################

def arg_check(args, name, lo, hi = -1):
    n = len(args)
    
    if hi == -1:
        if n != lo:
            err("'xm %s' requires %d argument%s.\n" % (name, lo,
                                                       lo > 1 and 's' or ''))
            usage(name)
    else:
        if n < lo or n > hi:
            err("'xm %s' requires between %d and %d arguments.\n" %
                (name, lo, hi))
            usage(name)


def unit(c):
    if not c.isalpha():
        return 0
    base = 1
    if c == 'G' or c == 'g': base = 1024 * 1024 * 1024
    elif c == 'M' or c == 'm': base = 1024 * 1024
    elif c == 'K' or c == 'k': base = 1024
    else:
        print 'ignoring unknown unit'
    return base

def int_unit(str, dest):
    base = unit(str[-1])
    if not base:
        return int(str)

    value = int(str[:-1])
    dst_base = unit(dest)
    if dst_base == 0:
        dst_base = 1
    if dst_base > base:
        return value / (dst_base / base)
    else:
        return value * (base / dst_base)

def err(msg):
    print >>sys.stderr, "Error:", msg

def handle_xend_error(cmd, args, ex):
    non_option = filter(lambda x: x[0] != '-', args)
    dom = len(non_option) > 0 and non_option[0] or None

    error = str(ex)
    if error == "Not found" and dom != None:
        err("Domain '%s' not found when running 'xm %s'" % (dom, cmd))
    else:
        err(error)

    sys.exit(1)
    

#########################################################################
#
#  Main xm functions
#
#########################################################################

def xm_save(args):
    arg_check(args, "save", 2)

    dom = args[0] # TODO: should check if this exists
    savefile = os.path.abspath(args[1])

    if not os.access(os.path.dirname(savefile), os.W_OK):
        err("xm save: Unable to create file %s" % savefile)
        sys.exit(1)
    
    from xen.xend.XendClient import server
    server.xend_domain_save(dom, savefile)
    
def xm_restore(args):
    arg_check(args, "restore", 1)

    savefile = os.path.abspath(args[0])

    if not os.access(savefile, os.R_OK):
        err("xm restore: Unable to read file %s" % savefile)
        sys.exit(1)

    from xen.xend.XendClient import server
    server.xend_domain_restore(savefile)


def getDomains(domain_names):
    from xen.xend.XendClient import server
    if domain_names:
        return map(server.xend_domain, domain_names)
    else:
        return server.xend_list_domains()


def xm_list(args):
    use_long = 0
    show_vcpus = 0
    try:
        (options, params) = getopt.gnu_getopt(args, 'lv', ['long','vcpus'])
    except getopt.GetoptError, opterr:
        err(opterr)
        sys.exit(1)
    
    for (k, v) in options:
        if k in ['-l', '--long']:
            use_long = 1
        if k in ['-v', '--vcpus']:
            show_vcpus = 1

    if show_vcpus:
        print >>sys.stderr, (
            "xm list -v is deprecated.  Please use xm vcpu-list.")
        xm_vcpu_list(params)
        return

    doms = getDomains(params)

    if use_long:
        map(PrettyPrint.prettyprint, doms)
    else:
        xm_brief_list(doms)


def parse_doms_info(info):
    def get_info(n, t, d):
        return t(sxp.child_value(info, n, d))
    
    return {
        'dom'      : get_info('domid',        int,   -1),
        'name'     : get_info('name',         str,   '??'),
        'mem'      : get_info('memory',       int,   0),
        'vcpus'    : get_info('online_vcpus', int,   0),
        'state'    : get_info('state',        str,   '??'),
        'cpu_time' : get_info('cpu_time',     float, 0),
        'ssidref'  : get_info('ssidref',      int,   0),
        }


def parse_sedf_info(info):
    def get_info(n, t, d):
        return t(sxp.child_value(info, n, d))

    return {
        'dom'      : get_info('domain',        int,   -1),
        'period'   : get_info('period',        int,   -1),
        'slice'    : get_info('slice',         int,   -1),
        'latency'  : get_info('latency',       int,   -1),
        'extratime': get_info('extratime',     int,   -1),
        'weight'   : get_info('weight',        int,   -1),
        }


def xm_brief_list(doms):
    print 'Name                              ID Mem(MiB) VCPUs State  Time(s)'
    for dom in doms:
        d = parse_doms_info(dom)
        if (d['ssidref'] != 0):
            d['ssidstr'] = (" s:%04x/p:%04x" % 
                            ((d['ssidref'] >> 16) & 0xffff,
                              d['ssidref']        & 0xffff))
        else:
            d['ssidstr'] = ""
        print ("%(name)-32s %(dom)3d %(mem)8d %(vcpus)5d %(state)5s %(cpu_time)7.1f%(ssidstr)s" % d)


def xm_vcpu_list(args):

    from xen.xend.XendClient import server
    if args:
        dominfo = map(server.xend_domain_vcpuinfo, args)
    else:
        doms = server.xend_list_domains(False)
        dominfo = map(server.xend_domain_vcpuinfo, doms)

    print 'Name                              ID  VCPU  CPU  State  Time(s)  CPU Affinity'

    for dom in dominfo:
        def get_info(n):
            return sxp.child_value(dom, n)

        #
        # convert a list of integers into a list of pairs indicating
        # continuous sequences in the list:
        #
        # [0,1,2,3]   -> [(0,3)]
        # [1,2,4,5]   -> [(1,2),(4,5)]
        # [0]         -> [(0,0)]
        # [0,1,4,6,7] -> [(0,1),(4,4),(6,7)]
        #
        def list_to_rangepairs(cmap):
            cmap.sort()
            pairs = []
            x = y = 0
            for i in range(0,len(cmap)):
                try:
                    if ((cmap[y+1] - cmap[i]) > 1):
                        pairs.append((cmap[x],cmap[y]))
                        x = y = i+1
                    else:
                        y = y + 1
                # if we go off the end, then just add x to y
                except IndexError:
                    pairs.append((cmap[x],cmap[y]))

            return pairs

        #
        # Convert pairs to range string, e.g: [(1,2),(3,3),(5,7)] -> 1-2,3,5-7
        #
        def format_pairs(pairs):
            if not pairs:
                return "no cpus"
            out = ""
            for f,s in pairs:
                if (f==s):
                    out += '%d'%f
                else:
                    out += '%d-%d'%(f,s)
                out += ','
            # trim trailing ','
            return out[:-1]

        def format_cpumap(cpumap):
            cpumap = map(lambda x: int(x), cpumap)
            cpumap.sort()

            from xen.xend.XendClient import server
            for x in server.xend_node()[1:]:
                if len(x) > 1 and x[0] == 'nr_cpus':
                    nr_cpus = int(x[1])
                    # normalize cpumap by modulus nr_cpus, and drop duplicates
                    cpumap = dict.fromkeys(
                                map(lambda x: x % nr_cpus, cpumap)).keys()
                    if len(cpumap) == nr_cpus:
                        return "any cpu"
                    break
 
            return format_pairs(list_to_rangepairs(cpumap))

        name  =     get_info('name')
        domid = int(get_info('domid'))

        for vcpu in sxp.children(dom, 'vcpu'):
            def vinfo(n, t):
                return t(sxp.child_value(vcpu, n))

            number   = vinfo('number',   int)
            cpu      = vinfo('cpu',      int)
            cpumap   = format_cpumap(vinfo('cpumap', list))
            online   = vinfo('online',   int)
            cpu_time = vinfo('cpu_time', float)
            running  = vinfo('running',  int)
            blocked  = vinfo('blocked',  int)

            if online:
                c = str(cpu)
                if running:
                    s = 'r'
                else:
                    s = '-'
                if blocked:
                    s += 'b'
                else:
                    s += '-'
                s += '-'
            else:
                c = "-"
                s = "--p"

            print (
                "%(name)-32s %(domid)3d  %(number)4d  %(c)3s   %(s)-3s   %(cpu_time)7.1f  %(cpumap)s" %
                locals())


def xm_reboot(args):
    arg_check(args, "reboot", 1, 4)
    from xen.xm import shutdown
    shutdown.main(["shutdown", "-R"] + args)

def xm_pause(args):
    arg_check(args, "pause", 1)
    dom = args[0]

    from xen.xend.XendClient import server
    server.xend_domain_pause(dom)

def xm_unpause(args):
    arg_check(args, "unpause", 1)
    dom = args[0]

    from xen.xend.XendClient import server
    server.xend_domain_unpause(dom)

def xm_rename(args):
    arg_check(args, "rename", 2)

    from xen.xend.XendClient import server
    server.xend_domain_rename(args[0], args[1])

def xm_subcommand(command, args):
    cmd = __import__(command, globals(), locals(), 'xen.xm')
    cmd.main([command] + args)


#############################################################

def cpu_make_map(cpulist):
    cpus = []
    for c in cpulist.split(','):
        if c.find('-') != -1:
            (x,y) = c.split('-')
            for i in range(int(x),int(y)+1):
                cpus.append(int(i))
        else:
            cpus.append(int(c))
    cpus.sort()
    return cpus

def xm_vcpu_pin(args):
    arg_check(args, "vcpu-pin", 3)

    dom  = args[0]
    vcpu = int(args[1])
    cpumap = cpu_make_map(args[2])
    
    from xen.xend.XendClient import server
    server.xend_domain_pincpu(dom, vcpu, cpumap)

def xm_mem_max(args):
    arg_check(args, "mem-max", 2)

    dom = args[0]
    mem = int_unit(args[1], 'm')

    from xen.xend.XendClient import server
    server.xend_domain_maxmem_set(dom, mem)
    
def xm_mem_set(args):
    arg_check(args, "mem-set", 2)

    dom = args[0]
    mem_target = int_unit(args[1], 'm')

    from xen.xend.XendClient import server
    server.xend_domain_mem_target_set(dom, mem_target)
    
def xm_vcpu_set(args):
    arg_check(args, "vcpu-set", 2)
    
    from xen.xend.XendClient import server
    server.xend_domain_set_vcpus(args[0], int(args[1]))


def xm_destroy(args):
    arg_check(args, "destroy", 1)
    from xen.xend.XendClient import server
    server.xend_domain_destroy(args[0])


def xm_domid(args):
    arg_check(args, "domid", 1)

    name = args[0]

    from xen.xend.XendClient import server
    dom = server.xend_domain(name)
    print sxp.child_value(dom, 'domid')
    
def xm_domname(args):
    arg_check(args, "domname", 1)

    name = args[0]

    from xen.xend.XendClient import server
    dom = server.xend_domain(name)
    print sxp.child_value(dom, 'name')