#! /usr/bin/env python

import os
import os.path
import sys
import time
import re
from fnmatch import fnmatch
from popen2 import popen2, popen3
from glob import glob
from itertools import chain
from tempfile import mktemp

def test_x(filename):
    return os.path.isfile(filename)
def test_f(filename):
    return os.path.isfile(filename)

"""recreation of grub-install in python scripting
   used to create helper code for pygtk grub conf editor """

class GrubInstall:
    # Initialize some variables.
    prefix= "/usr"
    exec_prefix= prefix
    sbindir= exec_prefix+"/sbin"
    datadir= "/usr/lib"
    PACKAGE= "grub"
    VERSION= "0.95"
    host_cpu= "i386"
    host_os= "linux"
    host_vendor= "suse"
    pkgdatadir= datadir+"/"+PACKAGE+"/"
    
    grub_shell= sbindir+"/grub"
    rootdir= ""
    grub_prefix= "/boot/grub"

    install_device= ""
    no_floppy= ""
    force_lba= ""
    recheck= "no"
    debug= "no"

    grubdir = None
    device_map = None

    def usage(self):
        """ Print the usage. """
        print """Usage: grub-install.py [OPTION] install_device
Install GRUB on your drive.

  -h, --help              print this message and exit
  -v, --version           print the version information and exit
  --root-directory=DIR    install GRUB images under the directory DIR
                          instead of the root directory
  --grub-shell=FILE       use FILE as the grub shell
  --no-floppy             do not probe any floppy drive
  --force-lba             force GRUB to use LBA mode even for a buggy
                          BIOS
  --recheck               probe a device map even if it already exists

INSTALL_DEVICE can be a GRUB device name or a system device filename.

Report bugs to <bug-grub@gnu.org>.
"""
    def convert_linux(self, os_device):
        """ Convert an OS device to the corresponding Linux GRUB drive. """
        found = re.match( "([sh]d[a-z])([0-9]*)$", os_device)
        if found:
            return found.group(1), found.group(2)
        found = re.match( "(d[0-9]*)p([0-9]*)$", os_device)
        if found:
            return found.group(1), found.group(3)
        found = re.match( "(fd[0-9]*)$", os_device)
        if found:
            return found.group(1), ""
        found = re.match( "(/disk|/part([0-9]*))$", os_device)
        if found:
            return "/disc", found.group(2)
        found = re.match( "(c[0-7]d[0-9]*)(p*)(.*)", os_device)
        if found:
            return found.group(1), found.group(3)
        return os_device, ""
    def convert_drivemap(self, diskpart):
        """ Get the drive name """
        tmp_drive= self.stdout_readlines(
            r"grep -v '^#' '"+self.device_map+"'"+
            r"| grep '"+diskpart+" *$'"+
            r"| sed 's%.*\(([hf]d[0-9][a-g0-9,]*)\).*%\1%'")
        # If not found, print an error message and exit.
        if not tmp_drive:
            raise GrubInstall.DevicemapError(
                "%s does not have any corresponding BIOS drive."%( diskpart))
        return tmp_drive[0].strip()
        
    def convert(self, os_device):
        """ Convert an OS device to the corresponding GRUB drive.
            This part is OS-specific. """
        if "linux" in self.host_os:
            tmp_disk, tmp_part = self.convert_linux( os_device)
        tmp_drive= self.convert_drivemap( tmp_disk)
        if not tmp_part:
           # If no partition is specified, just print the drive name.
           return tmp_drive
        # If a partition is specified, we need to translate it into the
        # GRUB's syntax.
        if "linux" in self.host_os:
            return tmp_drive.replace( ")", ",%i)"%( int(tmp_part) -1))
        return None

    def resolve_symlink(self, filename):
        """ Find the real file/device that file points at """
        return os.path.realpath( filename)
                
    def find_device(self, filename):
        """ Find block device on which the file resides. """
        # For now, this uses the program `df' to get the device name, but is
        # this really portable?
        tmp_fname= self.stdout_readlines( r"df "+filename+"/ "+
                                          r"| sed -n 's%.*\(/dev/[^   ]*\).*%\1%p'")
        if not tmp_fname:
            raise GrubInstall.Error( "Could not find device for %s"%( filename))
        return self.resolve_symlink( tmp_fname[0].strip())

    def stdout_readlines(self, command):     # stdout.readlines of command
        # print command+"\n"
        child_stdout, child_stdin, child_stderr= popen3(command)
        child_stdin.close()
        stdout_lines= child_stdout.readlines()
        child_stdout.close() ; child_stderr.close()
        return stdout_lines
    def run_grub_shell(self, args, *input): # stdout.readlines of command
        print self.grub_shell+" --batch "+args,"\n"
        for line in input: print ";", line
        child_stdout, child_stdin= popen2(self.grub_shell+" --batch "+args)
        for line in input:
            child_stdin.write(line+"\n")
        child_stdin.write("quit\n")
        child_stdin.close()
        stdout_lines= child_stdout.readlines()
        child_stdout.close()
        return stdout_lines
    def has_grub_error(self, stdout_lines): # boolean
        for line in stdout_lines:
            if re.match( "Error [0-9]*: ", line):
                return True
        return False

    class Error( Exception):
        pass
    class Exit( Exception):
        pass
    class ArgumentError( Error):
        pass
    class ArgumentExit( Exit):
        pass
    class DevicemapError( Error):
        pass
    class GrubshellError( Error):
        pass
    def check_arguments(self, args):
        for option in args:
            if fnmatch( option, "-h|-help"):
                self.usage()
                raise GrubInstall.ArgumentExit()
            elif fnmatch( option, "-v|--version"):
                print "grub-install (GNU GRUB",VERSION,")"
                raise GrubInstall.ArgumentExit()
            elif fnmatch( option, "--rootdirectory=*"):
                self.rootdir= option[ len('--root-directory=') :]
            elif fnmatch( option, "--grub-shell=*"):
                self.grub_shell= option[ len('--grub-shell=') :]
            elif fnmatch( option, "--no-floppy"):
                self.no_floppy= "--no-floppy"
            elif fnmatch( option, "--force-lba"):
                self.force_lba= "--force-lba" 
            elif fnmatch( option, "--recheck"):
                self.recheck= "yes"
            # This is an undocumented feature...
            elif fnmatch( option, "--debug"):
                self.debug= "yes"
            elif fnmatch( option, "--device-map=*"):
                self.device_map= option[ len('--device-map=') :]
            elif fnmatch( option, "-*"):
                raise GrubInstall.ArgumentError( "Unrecognized option %s"%( option))
            else:
                if self.install_device:
                    raise GrubInstall.ArgumentError(
                        "More than one install_devices?"
                        +"('%s' -vs- '%s')"%(option, self.install_device ))
                self.install_device= option
        if not self.install_device:
            raise GrubInstall.ArgumentError( "install_device not specified.")
        
    def check_bootdir(self):
         " Initialize these directories here, since ROOTDIR was initialized. "
         if "netbsd" in self.host_os:
             # Because /boot is used for the boot block in NetBSD, use /grub
             # instead of /boot/grub.
             self.grub_prefix= "/grub"
             self.bootdir= self.rootdir
         else:
             # Use /boot/grub by default.
             self.bootdir= self.rootdir+"/boot"
         #
         if self.grubdir is None:
             self.grubdir= self.bootdir+"/grub"
         if self.device_map is None:
             self.device_map= self.grubdir+"/device.map"

    def check_grub_installed(self):
        """ Check if GRUB is installed. """
        arg= self.grub_shell
        if not test_f( arg):
            raise GrubInstall.Error( arg+": Not found.")
        if not test_f( self.pkgdatadir+"/stage1"):
            raise GrubInstall.Error( self.pkgdatadir+"/stage1"+": Not found.")
        if not test_f( self.pkgdatadir+"/stage2"):
            raise GrubInstall.Error( self.pkgdatadir+"/stage2"+": Not found.")
        # Don't check for *stage1_5, because it is not fatal even if any
        # Stage 1.5 does not exist.

    def mkdir_bootdir(self):
        """ Create the GRUB directory if it is not present. """
        os.path.isdir(self.bootdir) or os.mkdir(self.bootdir) # or raise Exception
        os.path.isdir(self.grubdir) or os.mkdir(self.grubdir) # or raise Exception

    def create_device_map(self):
        """ Create the device map file if it is not present. """
        if test_f(self.device_map):
            return
        log_text= self.run_grub_shell(
            self.no_floppy+" --device-map="+self.device_map)
        if self.has_grub_error(log_text):
            raise GrubInstall.GrubshellError( log_text)
        if not test_f(self.device_map):
            raise GrubInstall.DevicemapError("could not create device map "+self.device_map)
        # Make sure that there is no duplicated entry.
        tmp= self.stdout_readlines(
            r"sed -n '/^([fh]d[0-9]*)/s/\(^(.*)\).*/\1/p' "+self.device_map+
            r"| sort | uniq -d | sed -n 1p")
        if len(tmp) and len(tmp[0].strip()):
            raise GrubInstall.DevicemapError(
                "The drive %s is defined multiple times in the device map %s"
                %( tmp[0].strip(), self.device_map))
        
    def check_install_device(self):
        # Check for INSTALL_DEVICE.
        if fnmatch( self.install_device, "/dev/*"):
            self.install_device= self.resolve_symlink( self.install_device)
            self.install_drive= self.convert( self.install_device)
        elif fnmatch( self.install_device, "([hf]d[0-9]*)"):
            self.install_drive= self.install_device
        elif fnmatch( self.install_device, "[hf]d[0-9]*"):
            # The GRUB format with no parenthesis.
            self.install_drive="(%s)"%( self.install_device) 
        else:
            raise GrubInstall.Error( "Format of install_device not recognized.")

    def check_root_drive(self):
        """ Get the root drive. """
        self.root_device= self.find_device( self.rootdir)
        self.bootdir_device= self.find_device( self.bootdir)
        # Check if the boot directory is in the same device as the root dir.
        if self.root_device != self.bootdir_device:
            # Perhaps the user has a separate boot partition.
            self.root_device= self.bootdir_device
            self.grub_prefix= "/grub"
        # Convert the root device to a GRUB drive.
        self.root_drive= self.convert( self.root_device)
        # Check if the root directory exists in the same device as the grub
        # directory.
        self.grubdir_device= self.find_device( self.grubdir)
        if self.grubdir_device != self.root_device:
            # For now, cannot deal with this situation.
            raise GrubInstall.Error(
                "You must set the root directory by the option"
                " --root-directory, because\n"
                "%s does not exist in the root device %s"
                %( self.grubdir, self.root_device))
    def __copy_files(self):
        # FHS says that /usr/share is used for architecture independent data,
        # so all stage-files are directly installed to /usr/lib/grub.
        # Therefor this part is no longer needed.
        # <--cut_here-->
        # # Copy the GRUB images to the GRUB directory.
        # for file in ${grubdir}/stage1 ${grubdir}/stage2 ${grubdir}/*stage1_5
        # do  rm -f $file || exit 1
        # done
        # for file in \
        #    ${pkgdatadir}/stage1 ${pkgdatadir}/stage2 ${pkgdatadir}/*stage1_5
        # do  cp -f $file ${grubdir} || exit 1
        # done
        # <--uncut-->
        pass
    
    def check_boot_files(self):
        """ Make sure that GRUB reads the same images as the host OS. """
        img_file= mktemp( "grub-install.img")

        for file in chain( glob( self.grubdir+"/stage1"),
                           glob( self.grubdir+"/stage2"),
                           glob( self.grubdir+"/*stage1_5")):
            if file.startswith(self.grubdir):
                tmp= file.replace(self.grubdir,self.grub_prefix, 1)
            else:
                tmp= file # errornous?
            count= 5
            while count:
                log_text= self.run_grub_shell(
                    self.no_floppy+" --device-map="+self.device_map,
                    "dump %s %s"%( self.root_drive+tmp, img_file))
                if self.has_grub_error( log_text):
                    pass
                elif filecmp.cmp( file, img_file):
                    break
                time.sleep( 1)
                count -= 1
            #
            if not count:
                raise GrubInstall.Error( "The file %s not read correctly."% (file))
        #
        os.remove( img_file)

    def perform_installation(self):
        # Now perform the installation.
        log_text= self.run_grub_shell(
            self.no_floppy+" --device-map="+self.device_map,
            "root %s"%( self.root_drive),
            "setup %s --stage2=%s --prefix=%s %s"
            %( self.force_lba, self.grubdir+"/stage2",
               self.grub_prefix, self.install_drive))
        if self.has_grub_errors( log_text) or self.debug == "yes":
            raise GrubInstall.GrubshellError( log_text)
        # Prompt the user to check if the device map is correct.
        print ((
            "Installation finished. No error reported.\n"+
            "This is the contents of the device map %s.\n"+
            "Check if this is correct or not. If any of the lines is"+
            " incorrect,\n fix it and re-run the script `%s'.\n\n")
               %( self.device_map, "grub-install"))
        self.print_device_map()
    def print_device_map(self):
        for line in self.stdout_readlines("cat "+self.device_map):
            print line

    def run(self, args):
        try:
            self.check_arguments(args)
            self.check_bootdir()
            self.check_grub_installed()
            self.mkdir_bootdir()
            self.create_device_map()
            self.check_install_device()
            self.check_root_drive()
            self.check_boot_files()
            # self.perform_installation()
        except GrubInstall.ArgumentError, x:
            print x.__class__.__name__+":", str(x)+"\n"
            self.usage()
        except GrubInstall.DevicemapError, x:
            print x.__class__.__name__+":", str(x)+"\n"
            self.print_device_map()
            # raise
        except GrubInstall.Exit, x:
            return
        except GrubInstall.Error, x:
            print x.__class__.__name__+":", str(x)+"\n"
            # raise
        except Exception, x:
            raise

if __name__ == "__main__":
    GrubInstall().run(sys.argv[1:])


## --root-directory=DIR
## (meaning:) a separate partion to host /boot/grub/ files