Add the install flow to grub plugin

1 files changed, 55 insertions, 22 deletions

diff --git a/plugins/plugin_grub.py b/plugins/plugin_grub.py
index 11fca0f..8137bfb 100644
--- a/plugins/plugin_grub.py
+++ b/plugins/plugin_grub.py
@@ -23,6 +23,7 @@ from pyfirstaidkit.utils import spawnvch
 import os.path
 import difflib
 import re
+import copy
 
 cfgBits = getConfigBits("firstaidkit-plugin-grub")
 import sys
@@ -34,6 +35,17 @@ class Sample1Plugin(Plugin):
     name = "GRUB plugin"
     version = "0.0.1"
     author = "Martin Sivak"
+    
+    flows = Flow.init(Plugin)
+    flows["install"] = Flow({
+                    Plugin.initial: {Return: "prepare"},
+                    "prepare"     : {ReturnSuccess: "diagnose", ReturnFailure: "clean", None: "clean"},
+                    "diagnose"    : {ReturnSuccess: "clean", ReturnFailure: "backup", None: "clean"},
+                    "backup"      : {ReturnSuccess: "install", ReturnFailure: "clean", None: "clean"},
+                    "restore"     : {ReturnSuccess: "clean", ReturnFailure: "clean", None: "clean"},
+                    "install"         : {ReturnSuccess: "clean", ReturnFailure: "restore", None: "restore"},
+                    "clean"       : {ReturnSuccess: Plugin.final, ReturnFailure: Plugin.final, None: "clean"}
+                    }, description="Install the bootloader into the partition specified in parameters")
 
     @classmethod
     def getDeps(cls):
@@ -41,18 +53,18 @@ class Sample1Plugin(Plugin):
 
     def __init__(self, *args, **kwargs):
         Plugin.__init__(self, *args, **kwargs)
-        self._partitions = [] #partitions in the system
+        self._partitions = set() #partitions in the system
         self._drives = [] #drives in the system
         self._bootable = [] #partitions with boot flag
         self._linux = [] #Linux type partitions (0x83 Linux)
-        self._grub_dir = [] #directories with stage1, menu.lst and other needed files
-        self._grub = [] #devices with possible grub instalation
+        self._grub_dir = set() #directories with stage1, menu.lst and other needed files
+        self._grub = set() #devices with possible grub instalation
         self._grub_map = {} #mapping from linux device names to grub device names
         self._grub_mask = re.compile("""\(hd[0-9]*,[0-9]*\)""")
-        from bootloaderInfo import x86BootloaderInfo
-        self._bootloaderInfo = x86BootloaderInfo()
 
     def prepare(self):
+        from bootloaderInfo import x86BootloaderInfo
+        self._bootloaderInfo = x86BootloaderInfo()
         self._result=ReturnSuccess
 
     def backup(self):
@@ -64,7 +76,7 @@ class Sample1Plugin(Plugin):
     def diagnose(self):
         #Find partitions
         self._reporting.debug(origin = self, level = PLUGIN, message = "Reading partition list")
-        self._partitions = map(lambda l: l.split()[3], file("/proc/partitions").readlines()[2:])
+        self._partitions = set(map(lambda l: l.split()[3], file("/proc/partitions").readlines()[2:]))
 
         #and select only partitions with minor number 0 -> drives
         self._drives = map(lambda l: l.split(), file("/proc/partitions").readlines()[2:])
@@ -99,27 +111,41 @@ class Sample1Plugin(Plugin):
         for l in stdout.split("\n"):
             if self._grub_mask.search(l):
                 self._reporting.info(origin = self, level = PLUGIN, message = "Grub directory found at %s" % (l.strip(),))
-                self._grub_dir.append(l.strip())
+                self._grub_dir.add(l.strip())
+
+        #TODO Mount the required partitions from self._grub_dir and read the important files from there
+        for gdrive in self._grub_dir:
+            drive, partition = gdrive[3:-1].split(",")
+            devicename = "%s%d" % (self._bootloaderInfo.drivelist[drive], int(partition)+1)
+            #XXX here, check the mount
 
         #Read the device map
         self._reporting.debug(origin = self, level = PLUGIN, message = "Reading device map")
-        for l in file(os.path.join(Config.system.root, "/boot/grub/device.map"), "r").readlines():
-            if l.startswith("#"):
-                continue
-            (v,k) = l.split()
-            self._grub_map[k] = v
+        try:
+            for l in file(os.path.join(Config.system.root, "/boot/grub/device.map"), "r").readlines():
+                if l.startswith("#"):
+                    continue
+                (v,k) = l.split()
+                self._grub_map[k] = v
+        except OSError, e:
+            self._reporting.warning(origin = self, level = PLUGIN, message = "Device map file not found: %s" % (str(e),))
+            pass #the file is not there, bad but not a reason to crash
 
         #Find out where is the grub installed
-        stage1mask = file(os.path.join(Config.system.root, "/boot/grub/stage1"), "r").read(512)
-        bootsectors = {}
-        for p in self._partitions:
-            self._reporting.debug(origin = self, level = PLUGIN, message = "Reading boot sector from %s" % (p,))
-            bootsector = file(os.path.join("/dev", p), "r").read(512)
-            bootsectors[bootsector] = self._bootloaderInfo.grubbyPartitionName(p)
-
-        for k in difflib.get_close_matches(stage1mask, bootsectors.keys()):
-            self._reporting.info(origin = self, level = PLUGIN, message = "Installed Grub probably found at %s" % (bootsectors[k],))
-            self._grub.append(bootsectors[k])
+        try:
+            stage1mask = file(os.path.join(Config.system.root, "/boot/grub/stage1"), "r").read(512)
+            bootsectors = {}
+            for p in self._partitions:
+                self._reporting.debug(origin = self, level = PLUGIN, message = "Reading boot sector from %s" % (p,))
+                bootsector = file(os.path.join("/dev", p), "r").read(512)
+                bootsectors[bootsector] = self._bootloaderInfo.grubbyPartitionName(p)
+
+            for k in difflib.get_close_matches(stage1mask, bootsectors.keys()):
+                self._reporting.info(origin = self, level = PLUGIN, message = "Installed Grub probably found at %s" % (bootsectors[k],))
+                self._grub.add(bootsectors[k])
+        except OSError, e:
+            self._reporting.warning(origin = self, level = PLUGIN, message = "Stage1 image not found: %s" % (str(e),))
+            pass #the file is not there, too bad but not a reason to crash the diagnose
 
         #if there is the grub configuration dir and the grub appears installed into MBR or bootable partition, then we are probably OK
         if len(self._grub_dir)>0 and len(self._grub)>0 and len(set(self._grub).intersection(set(self._bootable+self._drives)))>0:
@@ -131,6 +157,13 @@ class Sample1Plugin(Plugin):
     def fix(self):
         self._result=ReturnFailure
 
+    def install(self):
+        #install the grub to Config.operation.params partition
+        self._reporting.info(origin = self, level = PLUGIN, message = "Installing grub into %s (%s)" % (Config.operation.params, self._bootloaderInfo.grubbyPartitionName(Config.operation.params)))
+        grub = spawnvch(executable = "/sbin/grub", args = ["grub", "--batch"], chroot = Config.system.root)
+        (stdout, stderr) = grub.communicate("root %s\nsetup %s\n" % (self._grub_dir[0], self._bootloaderInfo.grubbyPartitionName(Config.operation.params)))
+        self._result=ReturnFailure
+
     def clean(self):
         self._result=ReturnSuccess