#!/usr/bin/python
#
# pygrub - simple python-based bootloader for Xen
#
# Copyright 2005-2006 Red Hat, Inc.
# Jeremy Katz <katzj@redhat.com>
#
# This software may be freely redistributed under the terms of the GNU
# general public license.
#
# 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., 675 Mass Ave, Cambridge, MA 02139, USA.
#

import os, sys, string, struct, tempfile, re
import copy
import logging
import platform

import curses, _curses, curses.wrapper, curses.textpad, curses.ascii
import getopt

sys.path = [ '/usr/lib/python' ] + sys.path

import fsimage
import grub.GrubConf
import grub.LiloConf

PYGRUB_VER = 0.5

def enable_cursor(ison):
    if ison:
        val = 2
    else:
        val = 0
        
    try:
        curses.curs_set(val)
    except _curses.error:
        pass

def is_disk_image(file):
    fd = os.open(file, os.O_RDONLY)
    buf = os.read(fd, 512)
    os.close(fd)

    if len(buf) >= 512 and \
           struct.unpack("H", buf[0x1fe: 0x200]) == (0xaa55,):
        return True
    return False

def get_active_partition(file):
    """Find the offset for the start of the first active partition "
    "in the disk image file."""

    fd = os.open(file, os.O_RDONLY)
    buf = os.read(fd, 512)
    for poff in (446, 462, 478, 494): # partition offsets
        # active partition has 0x80 as the first byte
        if struct.unpack("<c", buf[poff:poff+1]) == ('\x80',):
            return buf[poff:poff+16]

    # if there's not a partition marked as active, fall back to
    # the first partition
    return buf[446:446+16]

SECTOR_SIZE=512
DK_LABEL_LOC=1
DKL_MAGIC=0xdabe
V_ROOT=0x2

def get_solaris_slice(file, offset):
    """Find the root slice in a Solaris VTOC."""

    fd = os.open(file, os.O_RDONLY)
    os.lseek(fd, offset + (DK_LABEL_LOC * SECTOR_SIZE), 0)
    buf = os.read(fd, 512)
    if struct.unpack("<H", buf[508:510])[0] != DKL_MAGIC:
        raise RuntimeError, "Invalid disklabel magic"

    nslices = struct.unpack("<H", buf[30:32])[0]

    for i in range(nslices):
        sliceoff = 72 + 12 * i
        slicetag = struct.unpack("<H", buf[sliceoff:sliceoff+2])[0]
        slicesect = struct.unpack("<L", buf[sliceoff+4:sliceoff+8])[0]
        if slicetag == V_ROOT:
            return slicesect * SECTOR_SIZE

    raise RuntimeError, "No root slice found"      

def get_fs_offset_gpt(file):
    fd = os.open(file, os.O_RDONLY)
    # assume the first partition is an EFI system partition.
    os.lseek(fd, SECTOR_SIZE * 2, 0)
    buf = os.read(fd, 512)
    return struct.unpack("<Q", buf[32:40])[0] * SECTOR_SIZE

FDISK_PART_SOLARIS=0xbf
FDISK_PART_SOLARIS_OLD=0x82
FDISK_PART_GPT=0xee

def get_fs_offset(file):
    if not is_disk_image(file):
        return 0

    partbuf = get_active_partition(file)
    if len(partbuf) == 0:
        raise RuntimeError, "Unable to find active partition on disk"

    offset = struct.unpack("<L", partbuf[8:12])[0] * SECTOR_SIZE

    type = struct.unpack("<B", partbuf[4:5])[0]

    if type == FDISK_PART_SOLARIS or type == FDISK_PART_SOLARIS_OLD:
        offset += get_solaris_slice(file, offset)

    if type == FDISK_PART_GPT:
        offset = get_fs_offset_gpt(file)
    
    return offset

class GrubLineEditor(curses.textpad.Textbox):
    def __init__(self, screen, startx, starty, line = ""):
        screen.addstr(startx, starty, "> ")
        screen.refresh()
        win = curses.newwin(1, 74, startx, starty + 2)
        curses.textpad.Textbox.__init__(self, win)
        
        self.line = list(line)
        self.pos = len(line)
        self.cancelled = False
        self.show_text()

    def show_text(self):
        """Show the text.  One of our advantages over standard textboxes
        is that we can handle lines longer than the window."""

        self.win.clear()
        p = self.pos
        off = 0
        while p > 70:
            p -= 55
            off += 55

        l = self.line[off:off+70]
        self.win.addstr(0, 0, string.join(l, ("")))
        if self.pos > 70:
            self.win.addch(0, 0, curses.ACS_LARROW)

        self.win.move(0, p)

    def do_command(self, ch):
        # we handle escape as well as moving the line around, so have
        # to override some of the default handling

        self.lastcmd = ch
        if ch == 27: # esc
            self.cancelled = True
            return 0
        elif curses.ascii.isprint(ch):
            self.line.insert(self.pos, chr(ch))
            self.pos += 1
        elif ch == curses.ascii.SOH:  # ^a
            self.pos = 0
        elif ch in (curses.ascii.STX,curses.KEY_LEFT):
            if self.pos > 0:
                self.pos -= 1
        elif ch in (curses.ascii.BS,curses.KEY_BACKSPACE):
            if self.pos > 0:
                self.pos -= 1
                if self.pos < len(self.line):
                    self.line.pop(self.pos)
        elif ch == curses.ascii.EOT:                           # ^d
            if self.pos < len(self.line):
                self.line.pop(self.pos)
        elif ch == curses.ascii.ENQ:                           # ^e
            self.pos = len(self.line)
        elif ch in (curses.ascii.ACK, curses.KEY_RIGHT):
            if self.pos < len(self.line):
                self.pos +=1
        elif ch == curses.ascii.VT:                            # ^k
            self.line = self.line[:self.pos]
        else:
            return curses.textpad.Textbox.do_command(self, ch)
        self.show_text()
        return 1

    def edit(self):
        r = curses.textpad.Textbox.edit(self)
        if self.cancelled:
            return None
        return string.join(self.line, "")
        

class Grub:
    def __init__(self, file, fs = None):
        self.screen = None
        self.entry_win = None
        self.text_win = None
        if file:
            self.read_config(file, fs)

    def draw_main_windows(self):
        if self.screen is None: #only init stuff once
            self.screen = curses.initscr()
            self.screen.timeout(1000)
            if hasattr(curses, 'use_default_colors'):
                try:
                    curses.use_default_colors()
                except:
                    pass # Not important if we can't use colour
            enable_cursor(False)
            self.entry_win = curses.newwin(10, 74, 2, 1)
            self.text_win = curses.newwin(10, 70, 12, 5)
            curses.def_prog_mode()
        
        curses.reset_prog_mode()
        self.screen.clear()
        self.screen.refresh()

        # create basic grub screen with a box of entries and a textbox
        self.screen.addstr(1, 4, "pyGRUB  version %s" %(PYGRUB_VER,))
        self.entry_win.box()
        self.screen.refresh()

    def fill_entry_list(self):
        self.entry_win.clear()
        self.entry_win.box()
        for y in range(0, len(self.cf.images)):
            i = self.cf.images[y]
            if (0, y) > self.entry_win.getmaxyx():
                break
            if y == self.selected_image:
                attr = curses.A_REVERSE
            else:
                attr = 0
            self.entry_win.addstr(y + 1, 2, i.title.ljust(70), attr)
        self.entry_win.refresh()

    def edit_entry(self, origimg):
        def draw():
            self.draw_main_windows()

            self.text_win.addstr(0, 0, "Use the U and D keys to select which entry is highlighted.")
            self.text_win.addstr(1, 0, "Press 'b' to boot, 'e' to edit the selected command in the")
            self.text_win.addstr(2, 0, "boot sequence, 'c' for a command-line, 'o' to open a new line")
            self.text_win.addstr(3, 0, "after ('O' for before) the selected line, 'd' to remove the")
            self.text_win.addstr(4, 0, "selected line, or escape to go back to the main menu.")
            self.text_win.addch(0, 8, curses.ACS_UARROW)
            self.text_win.addch(0, 14, curses.ACS_DARROW)
            (y, x) = self.text_win.getmaxyx()
            self.text_win.move(y - 1, x - 1)
            self.text_win.refresh()

        curline = 1
        img = copy.deepcopy(origimg)
        while 1:
            draw()
            self.entry_win.clear()
            self.entry_win.box()
            for idx in range(1, len(img.lines)):
                # current line should be highlighted
                attr = 0
                if idx == curline:
                    attr = curses.A_REVERSE

                # trim the line
                l = img.lines[idx].ljust(70)
                if len(l) > 70:
                    l = l[:69] + ">"
                    
                self.entry_win.addstr(idx, 2, l, attr)
            self.entry_win.refresh()

            c = self.screen.getch()
            if c in (ord('q'), 27): # 27 == esc
                break
            elif c == curses.KEY_UP:
                curline -= 1
            elif c == curses.KEY_DOWN:
                curline += 1
            elif c == ord('b'):
                self.isdone = True
                break
            elif c == ord('e'):
                l = self.edit_line(img.lines[curline])
                if l is not None:
                    img.set_from_line(l, replace = curline)
            elif c == ord('d'):
                img.lines.pop(curline)
            elif c == ord('o'):
                img.lines.insert(curline+1, "")
                curline += 1
            elif c == ord('O'):
                img.lines.insert(curline, "")
            elif c == ord('c'):
                self.command_line_mode()
                if self.isdone:
                    return
                
            # bound at the top and bottom
            if curline < 1:
                curline = 1
            elif curline >= len(img.lines):
                curline = len(img.lines) - 1

        if self.isdone:
            origimg.reset(img.lines)

    def edit_line(self, line):
        self.screen.clear()
        self.screen.addstr(1, 2, "[ Minimal BASH-like line editing is supported.  ")
        self.screen.addstr(2, 2, "  ESC at any time cancels.  ENTER at any time accepts your changes. ]")
        self.screen.refresh()

        t = GrubLineEditor(self.screen, 5, 2, line)
        enable_cursor(True)
        ret = t.edit()
        if ret:
            return ret
        return None

    def command_line_mode(self):
        self.screen.clear()
        self.screen.addstr(1, 2, "[ Minimal BASH-like line editing is supported.  ESC at any time ")
        self.screen.addstr(2, 2, "  exits.  Typing 'boot' will boot with your entered commands. ] ")
        self.screen.refresh()

        y = 5
        lines = []
        while 1:
            t = GrubLineEditor(self.screen, y, 2)
            enable_cursor(True)            
            ret = t.edit()
            if ret:
                if ret in ("quit", "return"):
                    break
                elif ret != "boot":
                    y += 1
                    lines.append(ret)
                    continue

                # if we got boot, then we want to boot the entered image 
                img = grub.GrubConf.GrubImage(lines)
                self.cf.add_image(img)
                self.selected_image = len(self.cf.images) - 1
                self.isdone = True
                break

            # else, we cancelled and should just go back
            break

    def read_config(self, fn, fs = None):
        """Read the given file to parse the config.  If fs = None, then
        we're being given a raw config file rather than a disk image."""
        
        if not os.access(fn, os.R_OK):
            raise RuntimeError, "Unable to access %s" %(fn,)

        if platform.machine() == 'ia64':
            self.cf = grub.LiloConf.LiloConfigFile()
            # common distributions
            file_list = ("/efi/debian/elilo.conf", "/efi/gentoo/elilo.conf", 
                         "/efi/redflag/elilo.conf", "/efi/redhat/elilo.conf", 
                         "/efi/SuSE/elilo.conf",)
            # fallbacks
            file_list += ("/efi/boot/elilo.conf", "/elilo.conf",)
        else:
            self.cf = grub.GrubConf.GrubConfigFile()
            file_list = ("/boot/grub/menu.lst", "/boot/grub/grub.conf",
                         "/grub/menu.lst", "/grub/grub.conf")

        if not fs:
            # set the config file and parse it
            self.cf.filename = fn
            self.cf.parse()
            return

        for f in file_list:
            if fs.file_exists(f):
                self.cf.filename = f
                break
        if self.cf.filename is None:
            raise RuntimeError, "couldn't find bootloader config file in the image provided."
        f = fs.open_file(self.cf.filename)
        buf = f.read()
        del f
        self.cf.parse(buf)

    def run(self):
        timeout = int(self.cf.timeout)

        self.selected_image = self.cf.default
        self.isdone = False
        while not self.isdone:
            self.run_main(timeout)
            timeout = -1
            
        return self.selected_image

    def run_main(self, timeout = -1):
        def draw():
            # set up the screen
            self.draw_main_windows()
            self.text_win.addstr(0, 0, "Use the U and D keys to select which entry is highlighted.")
            self.text_win.addstr(1, 0, "Press enter to boot the selected OS. 'e' to edit the")
            self.text_win.addstr(2, 0, "commands before booting, 'a' to modify the kernel arguments ")
            self.text_win.addstr(3, 0, "before booting, or 'c' for a command line.")
            self.text_win.addch(0, 8, curses.ACS_UARROW)
            self.text_win.addch(0, 14, curses.ACS_DARROW)
            (y, x) = self.text_win.getmaxyx()
            self.text_win.move(y - 1, x - 1)
            self.text_win.refresh()

        # now loop until we hit the timeout or get a go from the user
        mytime = 0
        while (timeout == -1 or mytime < int(timeout)):
            draw()
            if timeout != -1 and mytime != -1: 
                self.screen.addstr(20, 5, "Will boot selected entry in %2d seconds"
                                   %(int(timeout) - mytime))
            else:
                self.screen.addstr(20, 5, " " * 80)
            self.fill_entry_list()

            c = self.screen.getch()
            if c == -1:
                # Timed out waiting for a keypress
                if mytime != -1:
                    mytime += 1
                    if mytime >= int(timeout):
                        self.isdone = True
                        break
            else:
                # received a keypress: stop the timer
                mytime = -1
                self.screen.timeout(-1)

            # handle keypresses
            if c == ord('c'):
                self.command_line_mode()
                break
            elif c == ord('a'):
                # find the kernel line, edit it and then boot
                img = self.cf.images[self.selected_image]
                for line in img.lines:
                    if line.startswith("kernel"):
                        l = self.edit_line(line)
                        if l is not None:
                            img.set_from_line(l, replace = True)
                            self.isdone = True
                            break
                break
            elif c == ord('e'):
                img = self.cf.images[self.selected_image]
                self.edit_entry(img)
                break
            elif c in (curses.KEY_ENTER, ord('\n'), ord('\r')):
                self.isdone = True
                break
            elif c == curses.KEY_UP:
                self.selected_image -= 1
            elif c == curses.KEY_DOWN:
                self.selected_image += 1
#            elif c in (ord('q'), 27): # 27 == esc
#                self.selected_image = -1
#                self.isdone = True
#                break

            # bound at the top and bottom
            if self.selected_image < 0:
                self.selected_image = 0
            elif self.selected_image >= len(self.cf.images):
                self.selected_image = len(self.cf.images) - 1
        
def get_entry_idx(cf, entry):
    # first, see if the given entry is numeric
    try:
        idx = string.atoi(entry)
        return idx
    except ValueError:
        pass

    # it's not, now check the labels for a match
    for i in range(len(cf.images)):
        if entry == cf.images[i].title:
            return i

    return None

def run_grub(file, entry, fs):
    global g
    global sel

    def run_main(scr, *args):
        global sel
        global g
        sel = g.run()

    g = Grub(file, fs)
    if interactive:
        curses.wrapper(run_main)
    else:
        sel = g.cf.default

    # set the entry to boot as requested
    if entry is not None:
        idx = get_entry_idx(g.cf, entry)
        if idx is not None and idx > 0 and idx < len(g.cf.images):
           sel = idx

    if sel == -1:
        print "No kernel image selected!"
        sys.exit(1)

    img = g.cf.images[sel]

    grubcfg = { "kernel": None, "ramdisk": None, "args": None }

    grubcfg["kernel"] = img.kernel[1]
    if img.initrd:
        grubcfg["ramdisk"] = img.initrd[1]
    if img.args:
        grubcfg["args"] = img.args

    return grubcfg

# If nothing has been specified, look for a Solaris domU. If found, perform the
# necessary tweaks.
def sniff_solaris(fs, cfg):
    if not fs.file_exists("/platform/i86xpv/kernel/unix"):
        return cfg
    
    # darned python
    longmode = (sys.maxint != 2147483647L)
    if not longmode:
        longmode = os.uname()[4] == "x86_64"
    if not longmode:
        if (os.access("/usr/bin/isainfo", os.R_OK) and
            os.popen("/usr/bin/isainfo -b").read() == "64\n"):
            longmode = True

    if not cfg["kernel"]:
        cfg["kernel"] = "/platform/i86xpv/kernel/unix"
        cfg["ramdisk"] = "/platform/i86pc/boot_archive"
        if longmode:
            cfg["kernel"] = "/platform/i86xpv/kernel/amd64/unix"
            cfg["ramdisk"] = "/platform/i86pc/amd64/boot_archive"

    # Unpleasant. Typically we'll have 'root=foo -k' or 'root=foo /kernel -k',
    # and we need to maintain Xen properties (root= and ip=) and the kernel
    # before any user args.
    
    xenargs = ""
    userargs = ""
    
    if not cfg["args"]:
        cfg["args"] = cfg["kernel"]
    else:
        for arg in cfg["args"].split():
            if re.match("^root=", arg) or re.match("^ip=", arg):
                xenargs += arg + " "
            elif arg != cfg["kernel"]:
                userargs += arg + " "
        cfg["args"] = xenargs + " " + cfg["kernel"] + " " + userargs

    return cfg
 
if __name__ == "__main__":
    sel = None
    
    def usage():
        print >> sys.stderr, "Usage: %s [-q|--quiet] [-i|--interactive] [--output=] [--kernel=] [--ramdisk=] [--args=] [--entry=] <image>" %(sys.argv[0],)

    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], 'qih::',
                                   ["quiet", "interactive", "help", "output=",
                                    "entry=", "kernel=", "ramdisk=", "args=",
                                    "isconfig"])
    except getopt.GetoptError:
        usage()
        sys.exit(1)

    if len(args) < 1:
        usage()
        sys.exit(1)
    file = args[0]
        
    output = None
    entry = None
    interactive = True
    isconfig = False

    # what was passed in
    incfg = { "kernel": None, "ramdisk": None, "args": None }
    # what grub or sniffing chose
    chosencfg = { "kernel": None, "ramdisk": None, "args": None }
    # what to boot
    bootcfg = { "kernel": None, "ramdisk": None, "args": None }

    for o, a in opts:
        if o in ("-q", "--quiet"):
            interactive = False
        elif o in ("-i", "--interactive"):
            interactive = True
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("--output",):
            output = a
        elif o in ("--kernel",):
            incfg["kernel"] = a
        elif o in ("--ramdisk",):
            incfg["ramdisk"] = a
        elif o in ("--args",):
            incfg["args"] = a
        elif o in ("--entry",):
            entry = a
            # specifying the entry to boot implies non-interactive
            interactive = False
        elif o in ("--isconfig",):
            isconfig = True

    if output is None or output == "-":
        fd = sys.stdout.fileno()
    else:
        fd = os.open(output, os.O_WRONLY)

    # debug
    if isconfig:
        chosencfg = run_grub(file, entry)
        print "  kernel: %s" % chosencfg["kernel"]
        if img.initrd:
            print "  initrd: %s" % chosencfg["ramdisk"]
        print "  args: %s" % chosencfg["args"]
        sys.exit(0)

    fs = fsimage.open(file, get_fs_offset(file))

    chosencfg = sniff_solaris(fs, incfg)

    if not chosencfg["kernel"]:
        chosencfg = run_grub(file, entry, fs)

    data = fs.open_file(chosencfg["kernel"]).read()
    (tfd, bootcfg["kernel"]) = tempfile.mkstemp(prefix="boot_kernel.",
        dir="/var/run/xend/boot")
    os.write(tfd, data)
    os.close(tfd)

    if chosencfg["ramdisk"]:
        data = fs.open_file(chosencfg["ramdisk"],).read()
        (tfd, bootcfg["ramdisk"]) = tempfile.mkstemp(prefix="boot_ramdisk.",
            dir="/var/run/xend/boot")
        os.write(tfd, data)
        os.close(tfd)
    else:
        initrd = None

    sxp = "linux (kernel %s)" % bootcfg["kernel"]
    if bootcfg["ramdisk"]:
        sxp += "(ramdisk %s)" % bootcfg["ramdisk"]
    if chosencfg["args"]:
        sxp += "(args \"%s\")" % chosencfg["args"]

    sys.stdout.flush()
    os.write(fd, sxp)
    
