pygrub

xen 3.2.2 源码

#!/usr/bin/python
#
# pygrub - simple python-based bootloader for Xen
#
# Copyright 2005-2006 Red Hat, Inc.
# Jeremy Katz <[e-mail]>
#
# 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.6

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()

        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:
                attr = curses.A_REVERSE
            else:
                attr = 0
            self.entry_win.addstr(y + 1 - self.start_image, 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":