#!/usr/bin/python
# -*- coding: iso-8859-1 -*-

# grubaker 2.0.0
# author: Alfonso E.M. alfonso@el-magnifico.org
# date: 15/Apr/2006
# last update: 22/Apr/2009 by Alfonso E.M.
# Patches by JuanJe Ojeda

import re
import pygtk
pygtk.require ('2.0')
import gtk
import gtk.glade
import os,commands,sys
import getopt
import xml.etree.ElementTree as ET


# DEBIAN MARKERS FOR AUTOMAGICALLY GENERATED KERNELS SECTION
AUTOMAGIC_START="### BEGIN AUTOMAGIC KERNELS LIST"
AUTOMAGIC_END="### END DEBIAN AUTOMAGIC KERNELS LIST"


class Menu:
	filename=""
	def __init__(self,filename):
		"""
		Grub menu file

		"""

		self.filename=filename
		self.error=''
        	self.liststore = gtk.ListStore(str,str,int)
                self.item=[]
		
		self.tree = ET.Element("menu.lst")
		self.parsefile()


### MENU.parsefile: reads menu.lst, and parses it into a tree (xml style) structure and in a gtk liststore
	def parsefile(self):

        	self.liststore.clear()

		try:
			file=open(self.filename)
		except:
			self.error="No se puede editar el menu.\nNo tiene permiso para modificar "+self.filename+". Posiblemente se encuentre en un sistema Live donde no se puede modificar el gestor de arranque. Pruebe en el sistema una vez instalado."
			print self.error
		else:
			options=""
			id=-1
                        self.default_timeout="3"
                        self.default_title=-1
                        self.default_id=-1
			mustreadoptions=False

			comments=""
			automagic=0
			elements_counter=0
			titles_counter=-1
 			self.node_before_magic=0
			self.node_magic=0
			self.node_after_magic=0
		

			while 1:
				line=file.readline()
				if line == "": 
                                       break


# As menu.lst says:
# "lines between the AUTOMAGIC KERNELS LIST markers will be modified by the debian update-grub script except for the default options below"
# So Grubaker must check (and keep) these odd markers 
                                  
				if line.find(AUTOMAGIC_START) != -1:
					i=self.liststore.append([gtk.STOCK_GOTO_BOTTOM,"ATENCION: el sistema puede modificar las siguientes líneas automáticamente",-1])
                                        self.item.append({})
					automagic=1
					if self.node_before_magic == 0:
					  self.node_before_magic=elements_counter
				if line.find(AUTOMAGIC_END) != -1:
					i=self.liststore.append([gtk.STOCK_GOTO_TOP,"Aquí termina la sección que se actualiza automáticamente",-2])
                                        self.item.append({})
					automagic=2



# buffering commend until empty line or command
				if line[:1] == "#":
                                  if line.find("Grubaker") != -1:	
     				    continue
				  comments=comments+line  
				  mustreadoptions=False
				  continue
# new comments block in tree
				else:
				  if comments != "":
   			  	    block = ET.SubElement(self.tree, "comment")
				    block.text=comments
				    elements_counter+=1
				    comments=""
                                  if line == "\n":
				    mustreadoptions=False
                                    continue

				line = line.replace("\n","")
                                if re.search("[=\s]+",line):
                                  key,value=re.split("[=\s]+",line,1)
                                else:
                                  key=line
				  value=''
                                
# title block, start subblock with kernel options
				if key == "title":
                                        id+=1
					titles_counter+=1   

                                        self.item.append({})
                                        self.item[id]['title']=value
					mustreadoptions=True

					if self.default_title == titles_counter:
						icon=gtk.STOCK_YES
						self.default_id = id
						self.item[id]['default']=True
					else:
						icon=gtk.STOCK_NO
						self.item[id]['default']=False
					i=self.liststore.append([icon,value,id])

					key=key+str(id)
	  				titleblock = ET.SubElement(self.tree, key)
				        elements_counter+=1
					titleblock.text=value

					if automagic == 1 and self.node_before_magic == 0:
					   self.node_before_magic=elements_counter

# title options
				else:
					if mustreadoptions:
					  subblock = ET.SubElement(titleblock, key)
					  subblock.text=value
					  subblock.set("id",str(id))

                                          self.item[id][key]=value

 				        else:  
	   				  block = ET.SubElement(self.tree, key)
				    	  elements_counter+=1
				  	  block.text=value
					  
					  if key == 'default':
					    self.default_title=int(value)


				if automagic == 1 and self.node_magic == 0:
					   self.node_magic=elements_counter
				if automagic == 2 and self.node_after_magic == 0:
					   self.node_after_magic=elements_counter
	
			file.close
#			ET.ElementTree(self.tree).write("/tmp/menu1.lst.xml")

### MENU.set_value: change an element property in the tree 
	def set_value(self,key,value):
		 element=self.tree.find(key)
		 if element == None:
  		   element = ET.SubElement(self.tree, key)
  	    	 element.text=value

### MENU.set_on: activates a keyword or deletes it from the tree (like 'hiddenmenu')
	def set_on(self,key,value):
		 element=self.tree.find(key)
		 if element == None and value == True:
  		   element = ET.SubElement(self.tree, key)
		 if element != None and value == False:
		     self.tree.remove(element)

### MENU.add_option: add an option to an element
	def add_option(self,id,option,parent):
		 if self.item[id].has_key(option):
		   subblock=ET.SubElement(parent,option)
                   subblock.set("id",str(id))
		   if self.item[id][option] != "":
		     subblock.text=self.item[id][option]
	   	   del self.item[id][option]		


### MENU.mod_option: changes element values
	def mod_option(self,id,option,parent):

		 subblock=parent.find(option)
		 if subblock != None:
		   if self.item[id].has_key(option):
		     subblock.text=self.item[id][option]
 		   else:
  		     parent.remove(subblock)
		 else:
		   if self.item[id].has_key(option):
		     subblock=ET.SubElement(parent,option)
		     subblock.text=self.item[id][option]
                     subblock.set("id",str(id))

### MENU.del_option: deletes element values
	def del_option(self,id,option,parent):
		 subblock=parent.find(option)
		 if subblock != None:
		   parent.remove(subblock)


### MENU.mod_title: modifies an edited title and all its suboptions (kernel, root, initrd...)
	def mod_title(self,titleblock,id):
	  titleblock.text=self.item[id]["title"]
# linux type
	  if self.item[id].has_key("kernel"):
 	    for option in ('root','kernel','module','initrd'):
              self.mod_option(id,option,titleblock)        
 	    for option in ('rootnoverify','makeactive','chainloader'):
              self.del_option(id,option,titleblock)        
# non-linux type
	  else:
 	    for option in ('rootnoverify','makeactive','chainloader'):
              self.mod_option(id,option,titleblock)
 	    for option in ('root','kernel','module','initrd'):
              self.del_option(id,option,titleblock)        
	  return

### MENU.new_title: inserts a new menu item and all its suboptions
	def new_title(self,id,magic):
# figuring out where to insert a new title element
# possible sections are: before automagic section, within or after it		
# if titles exist (in one of those sections) we know where to add a new one
# if not, we just can guess a good place!
	  if magic == 0:
	    position=self.node_before_magic
	  elif magic == 1:
	    position = self.node_magic
          else:
  	    position=self.node_after_magic

#	  print magic,position,self.node_before_magic,self.node_magic,self.node_after_magic

	  titleblock=ET.Element("title"+str(id))
          titleblock.text=self.item[id]["title"]

#	  root=self.tree.getroot()
#	  root.insert(position,titleblock)
	  self.tree.insert(position,titleblock)

# Order is important !
# First we'll write Linux options (ordered) then Windows ones
	  if self.item[id].has_key("kernel"):
	    for option in ('root','kernel','module','initrd'):
              self.add_option(id,option,titleblock)
	  else:
	    for option in ('rootnoverify','makeactive','chainloader'):
              self.add_option(id,option,titleblock)


## MENU.write: Dump elements tree to file menu.lst
	def write(self):

# insert/delete menu items in tree structure from menu items array
		titles_counter=0
		iter=self.liststore.get_iter_first()
		automagic=0
		while iter != None:
		  title=self.liststore.get_value(iter,1)
                  id=self.liststore.get_value(iter,2)

		  if id == -1:
		    automagic=1
                  if id == -2:
	 	    automagic=2

		  if id >= 0:
    		    titleelement=self.tree.find("title"+str(id))
		    if titleelement != None:
		      if title == "DELETED":
                        self.tree.remove(titleelement)
                      else:
                        self.mod_title(titleelement,id)
		    else:
		      self.new_title(id,automagic)
		
		    if id == self.default_id:
		      block=self.tree.find("default")
		      if block != None:
 		        block.text=str(titles_counter)
		    titles_counter += 1

		  iter=self.liststore.iter_next(iter)


# tree is ready to be converted
#		ET.ElementTree(self.tree).write("/tmp/menu2.lst.xml")

		tmpfilename=self.filename+".tmp"
		backupfilename=self.filename+".bak"
		automagic=0

		try:
			newfile=open(tmpfilename,"w")
		except:
			sys.exit("ERROR: "+self.filename+".tmp not writable")
			return

# converting tree structure to menu.lst format
		for element in self.tree.getiterator():
		  text=element.text

    		  if element.tag == "menu.lst":
		    newfile.write("# menu.lst edited with Grubaker\n\n")
		    continue
#   writing comments block
    		  if element.tag == "comment":
      		     newfile.write("\n"+text)
    		     continue
#   menu title block
		  if element.tag[0:5] == "title":
		    id=int(element.tag[5:])
		    if self.item[id]["title"] != "DELETED":
		      newfile.write("\ntitle "+text+"\n")                    
#   keyword block
                  else:
		      if element.get("id") == None:
			newfile.write("\n")
	      	      if element.text: 
		        newfile.write(element.tag+" "+text+"\n")
		      else:
		        newfile.write(element.tag+"\n")
		      if element.get("id") == None:
			newfile.write("\n")

		newfile.close

                os.rename(self.filename,backupfilename)
		os.rename(tmpfilename,self.filename)

		return

## MENU.set_default: sets a title as default boot
	def set_default(self,id):
		for i in range(0,len(self.item)):
		  self.item[i]['default'] = False
		self.item[id]['default'] = True
		self.default_id=id

class Appgui:
	def __init__(self, menu):
		"""
		In this init the main window is displayed
		"""
		dic = {
			 "on_bt_quit_clicked" : (self.quit),
		         "on_window_main_delete" : (self.quit), 
		         "on_window_preferences_delete_event" : self.bt_preferences_cancel_clicked, 
		         "on_treeview1_cursor_changed" : self.treeview1_cursor_changed, 
		         "on_treeview1_row_activated" : self.edit_item, 
		         "on_bt_ok_clicked" : self.bt_ok_clicked, 
		         "on_bt_preferences_clicked" : self.bt_preferences_clicked, 
		         "on_bt_edit_ok_clicked" : self.bt_edit_ok_clicked, 
		         "on_bt_edit_cancel_clicked" : self.bt_edit_cancel_clicked, 
		         "on_window_edit_delete_event" : self.bt_edit_cancel_clicked, 
		         "on_bt_preferences_ok_clicked" : self.bt_preferences_ok_clicked, 
		         "on_bt_preferences_cancel_clicked" : self.bt_preferences_cancel_clicked, 
		         "on_treeview1_drag_end" : self.treeview1_drag_end, 
		         "on_bt_delete_clicked" : self.bt_delete_clicked, 
		         "on_bt_new_clicked" : self.bt_new_clicked, 
			 "on_dialog_error_response" : self.dialog_error_response,
		}

		self.menu = menu
#		self.xml = gtk.glade.XML("./grubaker.glade")
		self.xml = gtk.glade.XML("/usr/share/grubaker/grubaker.glade")
		self.xml.signal_autoconnect (dic)
		self.treeview = self.xml.get_widget('treeview1')
		self.window_main = self.xml.get_widget('window_main')
		self.window_main.set_size_request(600,320)
		self.window_edit=self.xml.get_widget('window_edit')
		self.window_preferences=self.xml.get_widget('window_preferences')
		self.treeview.set_rules_hint(True)

		self.treeview.set_model(model=self.menu.liststore)

	        # create the TreeViewColumn to display the data
	        self.column = gtk.TreeViewColumn('Systems')

	        # add tvcolumn to treeview
       		self.treeview.append_column(self.column)

	        # create a CellRendererText to render the data
	        self.cellicon = gtk.CellRendererPixbuf()
	        self.cell = gtk.CellRendererText()
		self.cell.set_property('single-paragraph-mode',True)

        	# add the cell to the column
        	self.column.pack_start(self.cellicon, False)
        	self.column.pack_start(self.cell, True)

       		self.column.set_attributes(self.cellicon, stock_id=0)
       		self.column.add_attribute(self.cell, 'markup',1)

		if self.menu.error:
			dialog=self.xml.get_widget('dialog_error')
			errortext=self.window_edit=self.xml.get_widget('text_error')
			errortext.set_text(self.menu.error)
		        dialog.show()

	def dialog_error_response(self,dialog,response,*args):
		sys.exit('Error!')

	def run(self):
		gtk.main()

	def quit(*args):
		if hasattr(gtk, 'main_quit'):
	            gtk.main_quit()
	        else:
	            gtk.mainquit()


	            
        def bt_delete_clicked(self,widget):
                (iter,id,title)=self.get_selected_item()

                if id >= 0:	# Dummy items are not deletable
                   self.menu.liststore.remove(iter)
		   self.menu.item[id]["title"]="DELETED"
   	           self.change_item_buttons_state(False)
		   self.menu_has_changed()
                return
                
        def bt_new_clicked(self,widget):
		self.set_text('hidden_id','')
		self.set_text('entry_title','(Nuevo)')
                self.window_edit.show()
		self.menu_has_changed()
                return
                
        def get_selected_item(self):
                selection=self.treeview.get_selection()
		(model,iter)=selection.get_selected()
		title=self.menu.liststore.get_value(iter,1)
		id=self.menu.liststore.get_value(iter,2)
                return (iter,id,title)

	def bt_ok_clicked(self,widget):
		widget.set_sensitive(False)
		self.menu.write()
		self.quit()

	def treeview1_cursor_changed(self,treeview):
	        self.change_item_buttons_state(True)
	        return
	        
	def treeview1_drag_end(self,*args):
	        self.menu_has_changed()
	        return

        def change_item_buttons_state(self,value):
                widget=self.xml.get_widget('bt_delete')
        	widget.set_sensitive(value)
		return

        def menu_has_changed(self):
                widget=self.xml.get_widget('bt_ok')
        	widget.set_sensitive(True)
                return

        def bt_preferences_clicked(self,widget):
		timeout=self.menu.tree.findtext("timeout") or self.menu.defaulttimeout

		self.set_text('entry_timeout',timeout)

		if self.menu.tree.find("hiddenmenu") != None:
  		  self.set_active('check_hiddenmenu',True)

                self.window_preferences.show()
	        return

        def bt_preferences_ok_clicked(self,widget):
  	        self.menu.set_value("timeout",self.get_value('entry_timeout'))
                widget=self.xml.get_widget('check_hiddenmenu')
                if widget.get_active():
    	          self.menu.set_on("hiddenmenu",True)
		else:
    	          self.menu.set_on("hiddenmenu",False)

	        self.menu_has_changed()
	        self.window_preferences.hide()  
	        return  

        def bt_preferences_cancel_clicked(self,widget,*args):
	        self.window_preferences.hide()  
		return True


# APPGUI.edit_item: loading the menu edition windows
	def edit_item(self,treeview,TreePath,TreeViewColumn):
                (iter,id,title)=self.get_selected_item()

                if id < 0:	# Dummy items are not editable
		   return

		self.set_text('entry_title',title)
		self.set_text('hidden_id',str(id))

		if self.menu.item[id]['default']:
                    self.set_active('check_default',True)
                else:
                    self.set_active('check_default',False)

		
		if self.menu.item[id].has_key('root'):
		  self.set_text('entry_disk',self.menu.item[id]['root'])
		if self.menu.item[id].has_key('rootnoverify'):
		  self.set_text('entry_disk',self.menu.item[id]['rootnoverify'])

		if self.menu.item[id].has_key('kernel'):
		    if re.search("\s+",self.menu.item[id]['kernel']):
                       kernel,options=re.split("\s+",self.menu.item[id]['kernel'],1)         
                       self.set_text('entry_kernel',kernel)
                       self.set_text('entry_options',options)
                    else:
                       self.set_text('entry_kernel',self.menu.item[id]['kernel'])
                else:
                    self.set_text('entry_kernel','')
                    self.set_text('entry_options','')
        		
		if self.menu.item[id].has_key('initrd'):
        		self.set_text('entry_initrd',self.menu.item[id]['initrd'])
                else:
                        self.set_text('entry_initrd','')


		if self.menu.item[id].has_key('chainloader'):
        		self.set_text('entry_chainloader',self.menu.item[id]['chainloader'])
                else:
                        self.set_text('entry_chainloader','')

		if self.menu.item[id].has_key('makeactive'):
                    self.set_active('check_makeactive',True)
                else:
                    self.set_active('check_makeactive',False)
                
                self.window_edit.show()
		  
		return

# APPGUI.bt_edit_ok_clicked: new title in the menu or an old entry has been modified
	def bt_edit_ok_clicked(self,widget):

		old_id=self.get_value('hidden_id')
		title=self.get_value('entry_title')

# old_id is a hidden input, empty for new menu items (tricky)
		if old_id != '':
	                (iter,id,oldtitle)=self.get_selected_item()
        	        self.menu.liststore.set_value(iter,1,title)
	                del self.menu.item[id]
	                self.menu.item.insert(id,{})
		else:
	                self.menu.item.append({})
			id=len(self.menu.item)-1
			icon=gtk.STOCK_NO
			i=self.menu.liststore.append([icon,title,id])

	        self.menu.item[id]['title']=title


	        widget=self.xml.get_widget('check_default')
	        default=widget.get_active()
		if default:
		   self.menu.set_default(id)
                
	        kernel=self.get_value('entry_kernel')
	        options=self.get_value('entry_options')
	        if kernel:
	            self.menu.item[id]['kernel']=kernel+' '+options

	        initrd=self.get_value('entry_initrd')
	        if initrd:
	            self.menu.item[id]['initrd']=initrd
	            
	        chainloader=self.get_value('entry_chainloader')

#if chainloader field is not empty, 'rootnoverify' is used instead of 'root' (tricky)
	        if chainloader:
	            self.menu.item[id]['chainloader']=chainloader
    		    rootkey='rootnoverify'
		else:
		    rootkey='root'

	        disk=self.get_value('entry_disk')
                self.menu.item[id][rootkey]=disk

	        
	        widget=self.xml.get_widget('check_makeactive')
	        makeactive=widget.get_active()
	        if makeactive:
	            self.menu.item[id]['makeactive']=''
    
	        self.window_edit.hide()
	        self.menu_has_changed()
                return

        def set_text(self,widgetname,text):
                widget=self.xml.get_widget(widgetname)
	        widget.set_text(text)
        def get_value(self,widgetname):
                widget=self.xml.get_widget(widgetname)
	        return widget.get_text()
        def set_active(self,widgetname,value):
                widget=self.xml.get_widget(widgetname)
	        widget.set_active(value)

        def bt_edit_cancel_clicked(self,widget,*args):
	        self.window_edit.hide()  
		return True


		
def usage():
	print """
Usage:
   -h	--help		This simple help
   -f	--file=xxx	Menu file to edit (default is /boot/grub/menu.lst)
	"""


def main():
	try:
		opts, args = getopt.getopt(sys.argv[1:],"hf:",["help","file="])
	except getopt.GetoptError:
		usage()
		sys.exit(2)

	menufile="/boot/grub/menu.lst"
	for opt, arg in opts:     
	    if opt in ("-h", "--help"):
		usage()
		sys.exit()                  
	    elif opt in ("-f", "--file"):
		menufile = arg
	
	menu=Menu(menufile)
	app=Appgui(menu)
	app.run()
	  

if __name__ == '__main__':
	main()

