#!/usr/bin/python
########################################################################
# Copyright (C) 2010-2018 VMWare, Inc.                                 #
# All Rights Reserved                                                  #
########################################################################

import os

from .Misc import isString

"""Provides a class for parsing and writing syslinux/isolinux/pxelinux
   configuration."""

class Label(object):
   """Represents a LABEL tag in a syslinux configuration.
         Attributes:
            * name        - A string giving the name of the label. This will be
                            the text that comes after the LABEL tag in the
                            configuration file.
            * kernel      - A string giving the name of a kernel image. The
                            type of image is assumed from the file extension.
            * append      - A string to append to the kernel image as arguments.
            * linux       - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a Linux kernel.
            * boot        - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a bootstrap program.
            * bss         - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a syslinux boot sector.
                            (The DOS superblock will be patched in.)
            * pxe         - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a PXE network bootstrap
                            program.
            * fdimage     - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a floppy disk image.
            * comboot     - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a COMBOOT program.
            * com32       - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a COM32 program.
            * config      - Equivalent to the kernel attribute, but forces the
                            image to be recognized as a configuration file.
            * localboot   - A string. For PXELINUX, valid values are "0", "4"
                            and "5". All values cause booting from a local hard
                            drive. In addition, "4" causes the UNDI driver to
                            remain in memory, and "5" causes the entire PXE
                            stack to remain in memory. For ISOLINUX, valid
                            values are "0x00", "0x80" and "-1". "0x00" causes
                            booting from a floppy, "0x80" causes booting from a
                            hard drive, and "-1" causes the next device in the
                            BIOS boot order to be tried.
            * initrd      - A string. Equivalent to adding " initrd=value" to
                            the append attribute.
            * menulabel   - A string specifying a menu label for this label.
            * menupasswd  - A string specifying a password which must be
                            entered when booting this entry.
            * menugoto    - A string. If non-empty, causes selecting this label
                            to go to a sub-menu. The sub-menu should exist and
                            have the name specified in this attribute.
            * ipappend    - Only valid on PXE linux. Causes network information
                            to be appended to the kernel command line. Valid
                            values are integers 0, 1, 2 and 3. 0 disables the
                            option. 1 causes information to be appended as
                            ip=<client-ip>:<boot-server-ip>:<gw-ip>:<netmask>.
                            2 causes information to be appended as
                            BOOTIF=<hardware-address-of-boot-interface>.
                            3 causes both the "ip" and "BOOTIF" arguments to be
                            appended.
            * menuindent  - An integer specifying how many spaces to indent the
                            menu entry in the display.
            * menudisable - A boolean. If True, disables selection of this
                            label in a menu, but the menu is still displayed.
                            Useful for organizing a menu into categories.
            * menuhide    - A boolean. If True, causes the label not to be
                            displayed in the menu.
            * menudefault - A boolean. If True, selects this label as the
                            default for a menu.
            * menuexit    - Either None or a string. If a string, causes
                            selection of this label to exit the menu. If the
                            empty string, exits to a higher menu. If a
                            non-empty string, specifies the tag name of a menu
                            to exit to.
            * menuquit    - A boolean. If True, causes selection of this label
                            to exit the menu system.
            * texthelp    - A list. Each line in the list specifies help text
                            for the label entry in the menu.
   """

   _stringelementkeys = ("kernel", "append", "linux", "boot", "bss", "pxe",
                         "fdimage", "comboot", "com32", "config", "localboot",
                         "initrd", "menulabel", "menupasswd", "menugoto")

   def __init__(self, name):
      """Class Constructor.
            Parameters:
               * name - A string giving the name of the label.
            Returns: A new Label object.
      """
      self.name = name

      for key in self._stringelementkeys:
         setattr(self, key, "")

      self.ipappend = 0
      self.menuindent = 0
      self.menudisable = False
      self.menuhide = False
      self.menudefault = False
      self.menuexit = None
      self.menuquit = False
      self.texthelp = list()

      self._intexthelp = False

   def ParseLine(self, line):
      """Parses a single line from a configuration file.
            Parameters:
               * line - A string specifying the line.
            Raises:
               * ValueError - If a valid key cannot be parsed from the line, or
                              if the key requires a value and a valid value
                              cannot be parsed.
      """
      line = line.strip()
      words = line.split()
      lenwords = len(words)
      if not lenwords:
         return # blank line?

      key = words[0].lower()
      i = 1
      if i < lenwords and key in ("menu", "text"):
         key += words[i].lower()
         i += 1

      if key == "endtext":
         if not self._intexthelp:
            raise ValueError("Unexpected endtext key.")
         self._intexthelp = False
      elif self._intexthelp:
         self.texthelp.append(line)
      elif key == "texthelp":
         self._intexthelp = True
      elif key in self._stringelementkeys:
         if i >= lenwords:
            raise ValueError("Key '%s' has no value." % key)
         setattr(self, key, line.split(None, i)[i])
      elif key in ("ipappend", "menuindent"):
         if i >= lenwords:
            raise ValueError("Key '%s' has no value." % key)
         try:
            setattr(self, key, int(words[i]))
         except Exception:
            raise ValueError("Invalid value for '%s'." % key)
      elif key in ("menudisable", "menuhide", "menudefault", "menuquit"):
         setattr(self, key, True)
      elif key == "menuexit":
         if i < lenwords:
            self.menuexit = line.split(None, i)[i]
         else:
            self.menuexit = ""
      else:
         if key[:4] in ("menu", "text"):
            raise ValueError("Unknown key '%s %s'." % (key[:4], key[4:]))
         raise ValueError("Unknown key '%s'." % key)

   def ToLines(self, level=0):
      """A generator function that outputs the object as lines of a config
         file.
            Parameters:
               * level - An optional parameter, specifying which menu level the
                         object belongs to. Causes the lines to be indented 2
                         spaces for each level.
      """
      firstindent = "  " * level
      secondindent = firstindent + "  "
      yield "%sLABEL %s\n" % (firstindent, self.name)
      for key in self._stringelementkeys:
         val = getattr(self, key)
         if not val:
            continue
         if key.startswith("menu"):
            key = " ".join((key[:4], key[4:]))
         yield "%s%s %s\n" % (secondindent, key.upper(), val)
      for key in ("ipappend", "menuindent"):
         val = getattr(self, key)
         if val:
            if key[:4] == "menu":
               key = " ".join((key[:4], key[4:]))
            yield "%s%s %s\n" % (secondindent, key.upper(), val)
      for key in ("menudisable", "menuhide", "menudefault", "menuquit"):
         val = getattr(self, key)
         if val:
            key = " ".join((key[:4], key[4:])).upper()
            yield "%s%s\n" % (secondindent, key)
      if self.menuexit is not None:
         yield "%sMENU EXIT%s\n" % (secondindent, self.menuexit)
      if self.texthelp:
         yield "%sTEXT HELP\n" % secondindent
         for line in self.texthelp:
            yield "%s  %s\n" % (secondindent, line)
         yield "%sENDTEXT\n" % secondindent

class MenuColors(object):
   """A class representing MENU COLOR attributes.
         Attributes:
            * screen
            * border
            * title
            * unsel
            * hotkey
            * sel
            * hotsel
            * disabled
            * scrollbar
            * tabmsg
            * cmdmark
            * cmdline
            * pwdborder
            * pwdheader
            * pwdentry
            * timeout_msg
            * timeout
            * help
            * msg00
              ...
              msgFF

            Each attribute is either an empty string or a color specification.
            If an empty string, the menu's colors are inherited from the parent
            (or the defaults, for the top-level menu). Values may also be None,
            which indicates to use the inherited values from a top-level menu
            or the defaults.

            See http://syslinux.zytor.com/wiki/index.php/Comboot/menu.c32 for
            a description and examples of how to specify colors for various
            menu elements.
   """
   _stringelementkeys = (("screen", "border", "title", "unsel", "hotkey",
                          "sel", "hotsel", "disabled", "scrollbar", "tabmsg",
                          "cmdmark", "cmdline", "pwdborder", "pwdheader",
                          "pwdentry", "timeout_msg", "timeout", "help") +
                         tuple("msg%02X" % i for i in range(256)))

   def __init__(self):
      for key in self._stringelementkeys:
         setattr(self, key, None)

   def ParseLine(self, line):
      """Parses a single line, setting object attributes as appropriate.
            Parameters:
               * line - A string giving a line of input.
            Raises:
               * ValueError - If a recognized key cannot be parsed from the
                              input.
      """
      line = line.strip()
      words = line.split()
      lenwords = len(words)
      i = 0
      if i < lenwords and words[i].lower() == "menu":
         i += 1
      if i < lenwords and words[i].lower() == "color":
         i += 1
      if i < lenwords:
         key = words[i].lower()
         i += 1
      if key not in self._stringelementkeys:
         raise ValueError("Unknown key '%s'." % key)
      if i < lenwords:
         setattr(self, key, line.split(None, i)[i])

   def ToLines(self, level=0):
      """A generator function that outputs the object as lines of a config
         file.
            Parameters:
               * level - An optional parameter, specifying which menu level the
                         object belongs to. Causes the lines to be indented 2
                         spaces for each level.
      """
      indent = "  " * level
      for key in self._stringelementkeys:
         value = getattr(self, key)
         if value is not None:
            yield "%sMENU COLOR %s %s\n" % (indent, key.upper(), value)

class Config(object):
   """A class representing a syslinux configuration.
         Attributes:
            * default            - A string specifying a default kernel and
                                   command line to boot automatically.
            * ui                 - A string specifying a module to use for the
                                   user interface system.
            * ontimeout          - A string specifying a kernel and command
                                   line to boot if time out occurs waiting for
                                   user input.
            * onerror            - A string specifying a kernel to boot if the
                                   selected kernel is not found.
            * serial             - A string specifying serial port options. See
                                   syslinux documentation for details.
            * font               - A string specifying a file name to load a
                                   font from.
            * kbdmap             - A string specifying a file name to load a
                                   keyboard map from.
            * display            - A string specifying a file name to display
                                   contents from before the boot: prompt.
            * f1-f12             - Each attribute is a string specifying a
                                   file name. The contents of the file are
                                   displayed when the corresponding function
                                   key is pressed.
            * menutitle          - A string specifying a title for the menu.
            * menumasterpasswd   - A string specifying a password for the menu.
            * menuresolution     - A string specifying a "height width" for the
                                   menu.
            * menubackground     - A string giving a file name to use as a
                                   background image.
            * menunotabmsg       - A string specifying the message to display
                                   when option editing is disabled, but a user
                                   attempts to use the tab key.
            * menumsgcolor       - A string. Sets all msg00-msgFF colors to the
                                   specified value. See documentation for the
                                   MenuColors class.
            * menuhidden         - A boolean. If True, the menu is not
                                   displayed unless the user presses a key.
            * menuclear          - A boolean. If True, the screen is cleared
                                   when exiting the menu.
            * menushiftkey       - A boolean. If True, the menu exits
                                   immediately unless either the Shift or Alt
                                   key is depressed, or Caps Lock or Scroll
                                   Lock is set.
            * menustart          - A string. If specified, the menu system will
                                   display the menu with this name first,
                                   instead of the top-level menu.
            * menusave           - A boolean. If True, saves the selected entry
                                   as the default for the next boot. (Valid
                                   only for extlinux.)
            * menunosave         - A boolean. If True, specifically overrides
                                   menusave. (Useful to disable the option for
                                   sub-menus.)
            * menuwidth          - An integer specifying the width (in columns)
                                   of the menu, or None for the default.
            * menumargin         - An integer specifying how many columns to
                                   use for the menu's margins, or None for the
                                   default.
            * menupasswordmargin - An integer specifying how many columns to
                                   use for the margin of the password prompt.
            * menurows           - An integer specifying how many rows to use
                                   for the menu.
            * menutabmsgrow      - An integer specifying which row to place
                                   the tab message at.
            * menucmdlinerow     - An integer specifying which row to place
                                   the command line at.
            * menuendrow         - An integer specifying at which row to end
                                   the menu.
            * menupasswordrow    - An integer specifying where to place the
                                   password prompt.
            * menutimeoutrow     - An integer specifying where to place the
                                   timeout message.
            * menuhelpmsgrow     - An integer specifying where to place the
                                   help text.
            * menuhelpmsgendrow  - An integer specifying where to end the help
                                   text.
            * menuhiddenrow
            * menuhshift
            * menuvshift
            * menuautoboot       - A string specifying a message to use to
                                   replace the "Automatic boot in # second{,s}"
                                   prompt.
            * menutabmsg         - A string specifying a message to use to
                                   replace the "Press [Tab] to edit options"
                                   prompt.
            * menupassprompt     - A string specifying a message to use to
                                   replace the "Password required" prompt.
            * implicit           - A boolean. If False, disables booting kernel
                                   images which are not specifically listed in
                                   a label.
            * nohalt             - A boolean. If True, does not halt the
                                   processor when idle.
            * prompt             - A boolean. If True, causes the boot: prompt
                                   to always be displayed.
            * noescape           - A boolean. If True, causes the escape keys
                                   (Shift/Alt/Caps Lock/Scroll Lock) to be
                                   disabled.
            * nocomplete         - A boolean. If True, disables use of Tab to
                                   display labels at the boot: prompt.
            * allowoptions       - A boolean. If False, prevents the user from
                                   overriding options to the kernel command
                                   line.
            * timeout            - An integer specifying, in tenths of a
                                   second, how long to wait for user input
                                   before booting automatically.
            * totaltimeout       - An integer specifying, in tenths of a
                                   second, how long to wait for the user to
                                   select an entry before booting
                                   automatically.
            * console            - A boolean. If False, disables output to the
                                   video console.
            * say                - A list of strings. Strings in the list are
                                   printed to the console before the boot:
                                   prompt.
            * menucolors         - An instance of MenuColors.
            * items              - A list containing Label, Menu and
                                   MenuSeparator objects. The order of the list
                                   reflects the order in which the objects are
                                   read/written to the configuration file.
            Except for say, menuitems, and menucolors, any attribute can also
            be None, which means the value will be inherited from either the
            parent menu, or the defaults. (Attributes with None values are not
            written to the configuration file.)
   """

   _stringelementkeys = ("default", "ui", "ontimeout", "onerror", "serial",
                         "font", "kbdmap", "display", "f1", "f2", "f3", "f4",
                         "f5", "f6", "f7", "f8", "f9", "f10", "f11", "f12",
                         "menutitle", "menumasterpasswd", "menuresolution",
                         "menubackground", "menunotabmsg", "menumsgcolor",
                         "menuautoboot", "menutabmsg", "menupassprompt")

   _booltrueifpresentelementkeys = ("menuhidden", "menuclear", "menushiftkey",
                                    "menustart", "menusave", "menunosave")

   _boolelementkeys = ("implicit", "allowoptions", "console", "nohalt",
                       "prompt", "noescape", "nocomplete")

   _intelementkeys = ("menuwidth", "menumargin", "menupasswordmargin",
                      "menurows", "menutabmsgrow", "menucmdlinerow",
                      "menuendrow", "menupasswordrow", "menutimeoutrow",
                      "menuhelpmsgrow", "menuhelpmsgendrow", "menuhiddenrow",
                      "menuhshift", "menuvshift", "timeout", "totaltimeout")

   def __init__(self):
      for key in self._stringelementkeys:
         setattr(self, key, None)

      for key in self._booltrueifpresentelementkeys:
         setattr(self, key, None)

      for key in self._boolelementkeys:
         setattr(self, key, None)

      for key in self._intelementkeys:
         setattr(self, key, None)

      self.say = list()
      self.menucolors = MenuColors()

      # This contains an ordered list of Label, Menu, and Separator objects.
      self.items = list()

      # If we are within a label or menu begin, these will be Label/Menu
      # objects.
      self._inlabel = None
      self._inmenu = None

      self._directory = ""

   @property
   def labels(self):
      return list(item for item in self.items if isinstance(item, Label))

   @property
   def menus(self):
      return list(item for item in self.items if isinstance(item, Menu))

   def ParseLine(self, line):
      """Parses a single line, setting object attributes as appropriate.
            Parameters:
               * line - A string giving a line of input.
            Raises:
               * ValueError - If a recognized key cannot be parsed from the
                              input.
               * IOError    - If the configuration specifies an include file
                              that cannot be read.
      """
      # If we're in a menu, let the menu do the parsing. When the menu object
      # encounters the "MENU END" tag, it returns None. Otherwise, it returns
      # itself.
      if self._inmenu is not None:
         self._inmenu = self._inmenu.ParseLine(line)
         return

      # If in a TEXT HELP block, continue parsing as mult-line text.
      if self._inlabel is not None and self._inlabel._intexthelp:
         self._inlabel.ParseLine(line)

      line = line.strip()
      words = line.split()
      lenwords = len(words)
      if not lenwords:
         return # blank line

      if words[0].startswith("#"):
         return # comment. ('#' must come at the beginning of the line.)

      key = words[0].lower()
      i = 1
      if i < lenwords and key == "menu":
         key += words[i].lower()
         i += 1

      # MENU MASTER PASSWD is the only three-word key.
      if i < lenwords and key == "menumaster":
         key += words[i].lower()
         i += 1

      # Handle special cases first.
      if key == "label":
         if i >= lenwords:
            raise ValueError("label key with no value.")
         self._inlabel = Label(line.split(None, i)[i])
         self.items.append(self._inlabel)
      elif key == "menubegin":
         if i < lenwords:
            name = line.split(None, i)[i]
         self._inmenu = Menu(name)
         # Beginning of menu implies end of label.
         self._inlabel = None
         self.items.append(self._inmenu)
      elif key in ("include", "menuinclude"):
         filename = words[i]
         i += 1
         # If there is a tag, parse everything as if it were in a "menu begin"/
         # "menu end" block.
         f = open(os.path.join(self._directory, filename), "rb")
         if i < lenwords:
            menu = Menu(line.split(None, i)[i])
            # Caller must handle any exception.
            for line in f:
               menu.ParseLine(line.strip("\n"))
            self.items.append(menu)
         else:
            for line in f:
               self.ParseLine(line.strip("\n"))
         f.close()
      elif key == "menuseparator":
         # We don't need to instantiate objects for this case.
         self.items.append(MenuSeparator)
      elif key in self._stringelementkeys:
         if i < lenwords:
            setattr(self, key, line.split(None, i)[i])
      elif key in self._booltrueifpresentelementkeys:
         setattr(self, key, True)
      elif key in self._boolelementkeys:
         if i < lenwords:
            if words[i] == "0":
               setattr(self, key, False)
            elif words[i] == "1":
               setattr(self, key, True)
      elif key in self._intelementkeys:
         if i < lenwords:
            try:
               setattr(self, key, int(words[i]))
            except Exception:
               pass
      elif key == "say":
         if i < lenwords:
            self.say.append(line.split(None, i)[i])
      elif key == "menucolors":
         self.menucolors.ParseLine(line)
      elif key == "menuend":
         # Kind of hack-ish. Needs to be different for Config and Menu.
         self._menuend()
      elif self._inlabel:
         # May raise KeyError.
         self._inlabel.ParseLine(line)
      else:
         if key[:10] == "menumaster":
            raise ValueError("Unknown key '%s %s %s'." %
                             (key[:4], key[4:10], key[10:]))
         if key[:4] == "menu":
            raise ValueError("Unknown key '%s %s'." % (key[:4], key[4:]))
         raise ValueError("Uknown key '%s'." % key)

   def _menuend(self):
      raise ValueError("Mis-matched menu end.")

   def ToLines(self, level=0):
      """A generator function that outputs the object as lines of a config
         file.
            Parameters:
               * level - An optional parameter, specifying which menu level the
                         object belongs to. Causes the lines to be indented 2
                         spaces for each level.
      """
      indent = "  " * level
      for key in self._stringelementkeys:
         value = getattr(self, key)
         if value is not None:
            if key == "menumasterpasswd":
               yield "%sMENU MASTER PASSWD %s\n" % (indent, value)
            elif key[:4] == "menu":
               key = " ".join((key[:4], key[4:]))
            yield "%s%s %s\n" % (indent, key.upper(), value)
      for key in self._booltrueifpresentelementkeys:
         value = getattr(self, key)
         if value is not None:
            # These keys all start with "menu".
            yield "%s%s %s\n" % (indent, key[:4].upper(), key[4:].upper())
      for key in self._boolelementkeys:
         value = getattr(self, key)
         if value is not None:
            yield "%s%s %s\n" % (indent, key.upper(), value and "1" or "0")
      for key in self._intelementkeys:
         value = getattr(self, key)
         if value is not None:
            if key[:4] == "menu":
               key = " ".join((key[:4], key[4:]))
            yield ("%s%s %s\n" % (indent, key.upper(), value))
      for line in self.say:
         yield "%sSAY %s\n" % (indent, line.strip("\n"))
      for line in self.menucolors.ToLines(level):
         yield line
      for item in self.items:
         for line in item.ToLines(level):
            yield line

   def Parse(self, f):
      """Populates the object's attributes by parsing values from a config
         file.
            Parameters:
               * f - Either a file name or a file(-like) object.
            Raises:
               * IOError - If the file cannot be opened or read.
      """
      if isString(f):
         fobj = open(f, "rb")
         self._directory = os.path.dirname(f)
      else:
         fobj = f
         if hasattr(f, "name") and os.path.exists(f.name):
            self._directory = os.path.dirname(f.name)
      for line in fobj:
         self.ParseLine(line)

   @classmethod
   def FromFile(cls, f):
      """Creates a new object by parsing a configuration file.
            Parameters:
               * f - Either a file name or a file(-like) object.
            Returns: A new object.
            Raises:
               * IOError - If the file cannot be opened or read.
      """
      new = cls()
      new.Parse(f)
      return new

   def Write(self, f):
      """Writes the object to a file.
            Parameters:
               * f - Either a file name or a file(-like) object.
            Raises:
               * IOError - If the file cannot be opened or written to.
      """
      if isString(f):
         fobj = open(f, "wb")
      else:
         fobj = f
      for line in self.ToLines():
         fobj.write(line.encode())
      if isString(f):
         fobj.close()

   def AddLabel(self, name):
      """Adds a new label to the configuration.
            Parameters:
               * name - A name for the new label.
            Returns: The added Label object.
      """
      label = Label(name)
      self.items.append(label)
      return label

   def AddMenu(self, name):
      """Adds a new menu to the configuration.
            Parameters:
               * name - A name for the new menu.
            Returns: The added Menu object.
      """
      menu = Menu(name)
      self.items.append(menu)
      return menu

class MenuSeparator(object):
   "A class representing a menu separator."
   @staticmethod
   def ToLines(level=0):
      yield "%sMENU SEPARATOR\n" % "  " * level

class Menu(Config):
   """A class representing a menu (between MENU BEGIN and MENU END tags) in a
      configuration. Supports the attributes defined by Config.
   """
   def __init__(self, name=""):
      """Class constructor.
            Parameters:
               * name - An optional string giving a name for the menu.
            Returns: A new Menu object.
      """
      Config.__init__(self)
      self.name = name

   def ToLines(self, level=0):
      """A generator function that outputs the object as lines of a config
         file.
            Parameters:
               * level - An optional parameter, specifying which menu level the
                         object belongs to. Causes the lines to be indented 2
                         spaces for each level.
      """
      indent = "  " * level
      if self.name:
         yield "%sMENU BEGIN %s\n" % (indent, self.name)
      else:
         yield "%sMENU BEGIN\n" % indent
      for line in Config.ToLines(self, level + 1):
         yield line
      yield "%sMENU END\n" % indent

   def _menuend(self):
      return None
