Synnefo » snf-ganeti » ganeti-local

#!/usr/bin/python

#

# Copyright (C) 2009, 2010, 2011, 2012 Google Inc.

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# 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

# General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

# 02110-1301, USA.

"""Script to generate bash_completion script for Ganeti.

"""

# pylint: disable=C0103

# [C0103] Invalid name build-bash-completion

import os

import re

import itertools

from cStringIO import StringIO

from ganeti import constants

from ganeti import cli

from ganeti import utils

from ganeti import build

from ganeti import pathutils

# _autoconf shouldn't be imported from anywhere except constants.py, but we're

# making an exception here because this script is only used at build time.

from ganeti import _autoconf

#: Regular expression describing desired format of option names. Long names can

#: contain lowercase characters, numbers and dashes only.

_OPT_NAME_RE = re.compile(r"^-[a-zA-Z0-9]|--[a-z][-a-z0-9]+$")

def WritePreamble(sw):

  """Writes the script preamble.

  Helper functions should be written here.

  """

  sw.Write("# This script is automatically generated at build time.")

  sw.Write("# Do not modify manually.")

  sw.Write("_gnt_log() {")

  sw.IncIndent()

  try:

    sw.Write("if [[ -n \"$GANETI_COMPL_LOG\" ]]; then")

    sw.IncIndent()

    try:

      sw.Write("{")

      sw.IncIndent()

      try:

        sw.Write("echo ---")

        sw.Write("echo \"$@\"")

        sw.Write("echo")

      finally:

        sw.DecIndent()

      sw.Write("} >> $GANETI_COMPL_LOG")

    finally:

      sw.DecIndent()

    sw.Write("fi")

  finally:

    sw.DecIndent()

  sw.Write("}")

  sw.Write("_ganeti_nodes() {")

  sw.IncIndent()

  try:

    node_list_path = os.path.join(pathutils.DATA_DIR, "ssconf_node_list")

    sw.Write("cat %s 2>/dev/null || :", utils.ShellQuote(node_list_path))

  finally:

    sw.DecIndent()

  sw.Write("}")

  sw.Write("_ganeti_instances() {")

  sw.IncIndent()

  try:

    instance_list_path = os.path.join(pathutils.DATA_DIR,

                                      "ssconf_instance_list")

    sw.Write("cat %s 2>/dev/null || :", utils.ShellQuote(instance_list_path))

  finally:

    sw.DecIndent()

  sw.Write("}")

  sw.Write("_ganeti_jobs() {")

  sw.IncIndent()

  try:

    # FIXME: this is really going into the internals of the job queue

    sw.Write(("local jlist=$( shopt -s nullglob &&"

              " cd %s 2>/dev/null && echo job-* || : )"),

             utils.ShellQuote(pathutils.QUEUE_DIR))

    sw.Write('echo "${jlist//job-/}"')

  finally:

    sw.DecIndent()

  sw.Write("}")

  for (fnname, paths) in [

    ("os", pathutils.OS_SEARCH_PATH),

    ("iallocator", constants.IALLOCATOR_SEARCH_PATH),

    ]:

    sw.Write("_ganeti_%s() {", fnname)

    sw.IncIndent()

    try:

      # FIXME: Make querying the master for all OSes cheap

      for path in paths:

        sw.Write("( shopt -s nullglob && cd %s 2>/dev/null && echo * || : )",

                 utils.ShellQuote(path))

    finally:

      sw.DecIndent()

    sw.Write("}")

  sw.Write("_ganeti_nodegroup() {")

  sw.IncIndent()

  try:

    nodegroups_path = os.path.join(pathutils.DATA_DIR, "ssconf_nodegroups")

    sw.Write("cat %s 2>/dev/null || :", utils.ShellQuote(nodegroups_path))

  finally:

    sw.DecIndent()

  sw.Write("}")

  # Params: <offset> <options with values> <options without values>

  # Result variable: $first_arg_idx

  sw.Write("_ganeti_find_first_arg() {")

  sw.IncIndent()

  try:

    sw.Write("local w i")

    sw.Write("first_arg_idx=")

    sw.Write("for (( i=$1; i < COMP_CWORD; ++i )); do")

    sw.IncIndent()

    try:

      sw.Write("w=${COMP_WORDS[$i]}")

      # Skip option value

      sw.Write("""if [[ -n "$2" && "$w" == @($2) ]]; then let ++i""")

      # Skip

      sw.Write("""elif [[ -n "$3" && "$w" == @($3) ]]; then :""")

      # Ah, we found the first argument

      sw.Write("else first_arg_idx=$i; break;")

      sw.Write("fi")

    finally:

      sw.DecIndent()

    sw.Write("done")

  finally:

    sw.DecIndent()

  sw.Write("}")

  # Params: <list of options separated by space>

  # Input variable: $first_arg_idx

  # Result variables: $arg_idx, $choices

  sw.Write("_ganeti_list_options() {")

  sw.IncIndent()

  try:

    sw.Write("""if [[ -z "$first_arg_idx" ]]; then""")

    sw.IncIndent()

    try:

      sw.Write("arg_idx=0")

      # Show options only if the current word starts with a dash

      sw.Write("""if [[ "$cur" == -* ]]; then""")

      sw.IncIndent()

      try:

        sw.Write("choices=$1")

      finally:

        sw.DecIndent()

      sw.Write("fi")

      sw.Write("return")

    finally:

      sw.DecIndent()

    sw.Write("fi")

    # Calculate position of current argument

    sw.Write("arg_idx=$(( COMP_CWORD - first_arg_idx ))")

    sw.Write("choices=")

  finally:

    sw.DecIndent()

  sw.Write("}")

  # Params: <long options with equal sign> <all options>

  # Result variable: $optcur

  sw.Write("_gnt_checkopt() {")

  sw.IncIndent()

  try:

    sw.Write("""if [[ -n "$1" && "$cur" == @($1) ]]; then""")

    sw.IncIndent()

    try:

      sw.Write("optcur=\"${cur#--*=}\"")

      sw.Write("return 0")

    finally:

      sw.DecIndent()

    sw.Write("""elif [[ -n "$2" && "$prev" == @($2) ]]; then""")

    sw.IncIndent()

    try:

      sw.Write("optcur=\"$cur\"")

      sw.Write("return 0")

    finally:

      sw.DecIndent()

    sw.Write("fi")

    sw.Write("_gnt_log optcur=\"'$optcur'\"")

    sw.Write("return 1")

  finally:

    sw.DecIndent()

  sw.Write("}")

  # Params: <compgen options>

  # Result variable: $COMPREPLY

  sw.Write("_gnt_compgen() {")

  sw.IncIndent()

  try:

    sw.Write("""COMPREPLY=( $(compgen "$@") )""")

    sw.Write("_gnt_log COMPREPLY=\"${COMPREPLY[@]}\"")

  finally:

    sw.DecIndent()

  sw.Write("}")

def WriteCompReply(sw, args, cur="\"$cur\""):

  sw.Write("_gnt_compgen %s -- %s", args, cur)

  sw.Write("return")

class CompletionWriter:

  """Command completion writer class.

  """

  def __init__(self, arg_offset, opts, args):

    self.arg_offset = arg_offset

    self.opts = opts

    self.args = args

    for opt in opts:

      # While documented, these variables aren't seen as public attributes by

      # pylint. pylint: disable=W0212

      opt.all_names = sorted(opt._short_opts + opt._long_opts)

      invalid = list(itertools.ifilterfalse(_OPT_NAME_RE.match, opt.all_names))

      if invalid:

        raise Exception("Option names don't match regular expression '%s': %s" %

                        (_OPT_NAME_RE.pattern, utils.CommaJoin(invalid)))

  def _FindFirstArgument(self, sw):

    ignore = []

    skip_one = []

    for opt in self.opts:

      if opt.takes_value():

        # Ignore value

        for i in opt.all_names:

          if i.startswith("--"):

            ignore.append("%s=*" % utils.ShellQuote(i))

          skip_one.append(utils.ShellQuote(i))

      else:

        ignore.extend([utils.ShellQuote(i) for i in opt.all_names])

    ignore = sorted(utils.UniqueSequence(ignore))

    skip_one = sorted(utils.UniqueSequence(skip_one))

    if ignore or skip_one:

      # Try to locate first argument

      sw.Write("_ganeti_find_first_arg %s %s %s",

               self.arg_offset + 1,

               utils.ShellQuote("|".join(skip_one)),

               utils.ShellQuote("|".join(ignore)))

    else:

      # When there are no options the first argument is always at position

      # offset + 1

      sw.Write("first_arg_idx=%s", self.arg_offset + 1)

  def _CompleteOptionValues(self, sw):

    # Group by values

    # "values" -> [optname1, optname2, ...]

    values = {}

    for opt in self.opts:

      if not opt.takes_value():

        continue

      # Only static choices implemented so far (e.g. no node list)

      suggest = getattr(opt, "completion_suggest", None)

      # our custom option type

      if opt.type == "bool":

        suggest = ["yes", "no"]

      if not suggest:

        suggest = opt.choices

      if (isinstance(suggest, (int, long)) and

          suggest in cli.OPT_COMPL_ALL):

        key = suggest

      elif suggest:

        key = " ".join(sorted(suggest))

      else:

        key = ""

      values.setdefault(key, []).extend(opt.all_names)

    # Don't write any code if there are no option values

    if not values:

      return

    cur = "\"$optcur\""

    wrote_opt = False

    for (suggest, allnames) in values.items():

      longnames = [i for i in allnames if i.startswith("--")]

      if wrote_opt:

        condcmd = "elif"

      else:

        condcmd = "if"

      sw.Write("%s _gnt_checkopt %s %s; then", condcmd,

               utils.ShellQuote("|".join(["%s=*" % i for i in longnames])),

               utils.ShellQuote("|".join(allnames)))

      sw.IncIndent()

      try:

        if suggest == cli.OPT_COMPL_MANY_NODES:

          # TODO: Implement comma-separated values

          WriteCompReply(sw, "-W ''", cur=cur)

        elif suggest == cli.OPT_COMPL_ONE_NODE:

          WriteCompReply(sw, "-W \"$(_ganeti_nodes)\"", cur=cur)

        elif suggest == cli.OPT_COMPL_ONE_INSTANCE:

          WriteCompReply(sw, "-W \"$(_ganeti_instances)\"", cur=cur)

        elif suggest == cli.OPT_COMPL_ONE_OS:

          WriteCompReply(sw, "-W \"$(_ganeti_os)\"", cur=cur)

        elif suggest == cli.OPT_COMPL_ONE_IALLOCATOR:

          WriteCompReply(sw, "-W \"$(_ganeti_iallocator)\"", cur=cur)

        elif suggest == cli.OPT_COMPL_ONE_NODEGROUP:

          WriteCompReply(sw, "-W \"$(_ganeti_nodegroup)\"", cur=cur)

        elif suggest == cli.OPT_COMPL_INST_ADD_NODES:

          sw.Write("local tmp= node1= pfx= curvalue=\"${optcur#*:}\"")

          sw.Write("if [[ \"$optcur\" == *:* ]]; then")

          sw.IncIndent()

          try:

            sw.Write("node1=\"${optcur%%:*}\"")

            sw.Write("if [[ \"$COMP_WORDBREAKS\" != *:* ]]; then")

            sw.IncIndent()

            try:

              sw.Write("pfx=\"$node1:\"")

            finally:

              sw.DecIndent()

            sw.Write("fi")

          finally:

            sw.DecIndent()

          sw.Write("fi")

          sw.Write("_gnt_log pfx=\"'$pfx'\" curvalue=\"'$curvalue'\""

                   " node1=\"'$node1'\"")

          sw.Write("for i in $(_ganeti_nodes); do")

          sw.IncIndent()

          try:

            sw.Write("if [[ -z \"$node1\" ]]; then")

            sw.IncIndent()

            try:

              sw.Write("tmp=\"$tmp $i $i:\"")

            finally:

              sw.DecIndent()

            sw.Write("elif [[ \"$i\" != \"$node1\" ]]; then")

            sw.IncIndent()

            try:

              sw.Write("tmp=\"$tmp $i\"")

            finally:

              sw.DecIndent()

            sw.Write("fi")

          finally:

            sw.DecIndent()

          sw.Write("done")

          WriteCompReply(sw, "-P \"$pfx\" -W \"$tmp\"", cur="\"$curvalue\"")

        else:

          WriteCompReply(sw, "-W %s" % utils.ShellQuote(suggest), cur=cur)

      finally:

        sw.DecIndent()

      wrote_opt = True

    if wrote_opt:

      sw.Write("fi")

    return

  def _CompleteArguments(self, sw):

    if not (self.opts or self.args):

      return

    all_option_names = []

    for opt in self.opts:

      all_option_names.extend(opt.all_names)

    all_option_names.sort()

    # List options if no argument has been specified yet

    sw.Write("_ganeti_list_options %s",

             utils.ShellQuote(" ".join(all_option_names)))

    if self.args:

      last_idx = len(self.args) - 1

      last_arg_end = 0

      varlen_arg_idx = None

      wrote_arg = False

      sw.Write("compgenargs=")

      for idx, arg in enumerate(self.args):

        assert arg.min is not None and arg.min >= 0

        assert not (idx < last_idx and arg.max is None)

        if arg.min != arg.max or arg.max is None:

          if varlen_arg_idx is not None:

            raise Exception("Only one argument can have a variable length")

          varlen_arg_idx = idx

        compgenargs = []

        if isinstance(arg, cli.ArgUnknown):

          choices = ""

        elif isinstance(arg, cli.ArgSuggest):

          choices = utils.ShellQuote(" ".join(arg.choices))

        elif isinstance(arg, cli.ArgInstance):

          choices = "$(_ganeti_instances)"

        elif isinstance(arg, cli.ArgNode):

          choices = "$(_ganeti_nodes)"

        elif isinstance(arg, cli.ArgGroup):

          choices = "$(_ganeti_nodegroup)"

        elif isinstance(arg, cli.ArgJobId):

          choices = "$(_ganeti_jobs)"

        elif isinstance(arg, cli.ArgOs):

          choices = "$(_ganeti_os)"

        elif isinstance(arg, cli.ArgFile):

          choices = ""

          compgenargs.append("-f")

        elif isinstance(arg, cli.ArgCommand):

          choices = ""

          compgenargs.append("-c")

        elif isinstance(arg, cli.ArgHost):

          choices = ""

          compgenargs.append("-A hostname")

        else:

          raise Exception("Unknown argument type %r" % arg)

        if arg.min == 1 and arg.max == 1:

          cmpcode = """"$arg_idx" == %d""" % (last_arg_end)

        elif arg.max is None:

          cmpcode = """"$arg_idx" -ge %d""" % (last_arg_end)

        elif arg.min <= arg.max:

          cmpcode = (""""$arg_idx" -ge %d && "$arg_idx" -lt %d""" %

                     (last_arg_end, last_arg_end + arg.max))

        else:

          raise Exception("Unable to generate argument position condition")

        last_arg_end += arg.min

        if choices or compgenargs:

          if wrote_arg:

            condcmd = "elif"

          else:

            condcmd = "if"

          sw.Write("""%s [[ %s ]]; then""", condcmd, cmpcode)

          sw.IncIndent()

          try:

            if choices:

              sw.Write("""choices="$choices "%s""", choices)

            if compgenargs:

              sw.Write("compgenargs=%s",

                       utils.ShellQuote(" ".join(compgenargs)))

          finally:

            sw.DecIndent()

          wrote_arg = True

      if wrote_arg:

        sw.Write("fi")

    if self.args:

      WriteCompReply(sw, """-W "$choices" $compgenargs""")

    else:

      # $compgenargs exists only if there are arguments

      WriteCompReply(sw, '-W "$choices"')

  def WriteTo(self, sw):

    self._FindFirstArgument(sw)

    self._CompleteOptionValues(sw)

    self._CompleteArguments(sw)

def WriteCompletion(sw, scriptname, funcname,

                    commands=None,

                    opts=None, args=None):

  """Writes the completion code for one command.

  @type sw: ShellWriter

  @param sw: Script writer

  @type scriptname: string

  @param scriptname: Name of command line program

  @type funcname: string

  @param funcname: Shell function name

  @type commands: list

  @param commands: List of all subcommands in this program

  """

  sw.Write("%s_inner() {", funcname)

  sw.IncIndent()

  try:

    sw.Write("local i first_arg_idx choices compgenargs arg_idx optcur"

             ' cur="${COMP_WORDS[COMP_CWORD]}"'

             ' prev="${COMP_WORDS[COMP_CWORD-1]}"')

    sw.Write("_gnt_log cur=\"$cur\" prev=\"$prev\"")

    sw.Write("[[ -n \"$GANETI_COMPL_LOG\" ]] &&"

             " _gnt_log \"$(set | grep ^COMP_)\"")

    sw.Write("COMPREPLY=()")

    if opts is not None and args is not None:

      assert not commands

      CompletionWriter(0, opts, args).WriteTo(sw)

    else:

      sw.Write("""if [[ "$COMP_CWORD" == 1 ]]; then""")

      sw.IncIndent()

      try:

        # Complete the command name

        WriteCompReply(sw,

                       ("-W %s" %

                        utils.ShellQuote(" ".join(sorted(commands.keys())))))

      finally:

        sw.DecIndent()

      sw.Write("fi")

      # Group commands by arguments and options

      grouped_cmds = {}

      for cmd, (_, argdef, optdef, _, _) in commands.items():

        if not (argdef or optdef):

          continue

        grouped_cmds.setdefault((tuple(argdef), tuple(optdef)), set()).add(cmd)

      # We're doing options and arguments to commands

      sw.Write("""case "${COMP_WORDS[1]}" in""")

      sort_grouped = sorted(grouped_cmds.items(),

                            key=lambda (_, y): sorted(y)[0])

      for ((argdef, optdef), cmds) in sort_grouped:

        assert argdef or optdef

        sw.Write("%s)", "|".join(map(utils.ShellQuote, sorted(cmds))))

        sw.IncIndent()

        try:

          CompletionWriter(1, optdef, argdef).WriteTo(sw)

        finally:

          sw.DecIndent()

        sw.Write(";;")

      sw.Write("esac")

  finally:

    sw.DecIndent()

  sw.Write("}")

  # Wrapper function to always enable extglob (needed for advanced pattern

  # matching)

  sw.Write("%s() {", funcname)

  sw.IncIndent()

  try:

    # Get current state of extglob

    sw.Write("local -r eg=$(shopt -p extglob || :)")

    # Enable extglob

    sw.Write("shopt -s extglob")

    sw.Write("%s_inner \"$@\"", funcname)

    # Reset extglob to original value

    sw.Write("[[ -n \"$eg\" ]] && $eg")

  finally:

    sw.DecIndent()

  sw.Write("}")

  sw.Write("complete -F %s -o filenames %s",

           utils.ShellQuote(funcname),

           utils.ShellQuote(scriptname))

def GetFunctionName(name):

  return "_" + re.sub(r"[^a-z0-9]+", "_", name.lower())

def GetCommands(filename, module):

  """Returns the commands defined in a module.

  Aliases are also added as commands.

  """

  try:

    commands = getattr(module, "commands")

  except AttributeError:

    raise Exception("Script %s doesn't have 'commands' attribute" %

                    filename)

  # Add the implicit "--help" option

  help_option = cli.cli_option("-h", "--help", default=False,

                               action="store_true")

  for name, (_, _, optdef, _, _) in commands.items():

    if help_option not in optdef:

      optdef.append(help_option)

    for opt in cli.COMMON_OPTS:

      if opt in optdef:

        raise Exception("Common option '%s' listed for command '%s' in %s" %

                        (opt, name, filename))

      optdef.append(opt)

  # Use aliases

  aliases = getattr(module, "aliases", {})

  if aliases:

    commands = commands.copy()

    for name, target in aliases.items():

      commands[name] = commands[target]

  return commands

def main():

  buf = StringIO()

  sw = utils.ShellWriter(buf)

  WritePreamble(sw)

  # gnt-* scripts

  for scriptname in _autoconf.GNT_SCRIPTS:

    filename = "scripts/%s" % scriptname

    WriteCompletion(sw, scriptname,

                    GetFunctionName(scriptname),

                    commands=GetCommands(filename,

                                         build.LoadModule(filename)))

  # Burnin script

  burnin = build.LoadModule("tools/burnin")

  WriteCompletion(sw, "%s/burnin" % pathutils.TOOLSDIR, "_ganeti_burnin",

                  opts=burnin.OPTIONS, args=burnin.ARGUMENTS)

  print buf.getvalue()

if __name__ == "__main__":

  main()