#!/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, reimport copyimport loggingimport platformimport curses, _curses, curses.wrapper, curses.textpad, curses.asciiimport getoptsys.path = [ '/usr/lib/python', '/usr/lib64/python' ] + sys.pathimport fsimageimport grub.GrubConfimport grub.LiloConfPYGRUB_VER = 0.6def enable_cursor(ison):    if ison:        val = 2    else:        val = 0            try:        curses.curs_set(val)    except _curses.error:        passdef 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 Falsedef 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=512DK_LABEL_LOC=1DKL_MAGIC=0xdabeV_ROOT=0x2def 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_SIZEFDISK_PART_SOLARIS=0xbfFDISK_PART_SOLARIS_OLD=0x82FDISK_PART_GPT=0xeedef 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 offsetclass 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()        maxy = self.entry_win.getmaxyx()[0]-3 # maxy - 2 for the frame + index        if self.selected_image > self.start_image + maxy:            self.start_image = self.selected_image        if self.selected_image < self.start_image:            self.start_image = self.selected_image                for y in range(self.start_image, len(self.cf.images)):            i = self.cf.images[y]            if y > self.start_image + maxy:                break            if y == self.selected_image:                self.entry_win.attron(curses.A_REVERSE)            self.entry_win.addstr(y + 1 - self.start_image, 2, i.title.ljust(70))            if y == self.selected_image:                self.entry_win.attroff(curses.A_REVERSE)        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                if idx == curline:                    self.entry_win.attron(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)                if idx == curline:                    self.entry_win.attroff(curses.A_REVERSE)            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