#!/usr/bin/env python

##############################################################################
# c2005 by Paul Sherman <psherma1@rochester.rr.com> http://www.pcbypaul.com  #
#									     #
# This program is free software; you can redistribute it and/or		     #
# modify it under the terms of the GNU General Public License		     #
# as published by the Free Software Foundation; either version 2	     #
# of the License, or (at your option) any later version.		     #
#									     #
# This program is distributed in the hope that it will be useful,	     #
# but WITHOUT ANY WARRANTY; without even the implied warranty of	     #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the		     #
# GNU General Public License for more details.				     #
#									     #
# 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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.#
##############################################################################

import sys, os, time, commands, string, shutil, re, signal, popen2
import pygtk
import gobject
if gobject.pygtk_version < (2,5,96):
    print >>sys.stderr, "pyGTK version 2.6 or greater is required for this application"
    raise SystemExit
import gtk
try:
	import pango
except:
	print >> sys.stderr, "Cannot load Pango"
	raise SystemExit
	
if not len(commands.getoutput('type -path df')): 
	print >> sys.stderr, "Cannot run linHDD without the df utility"
	raise SystemExit
if not len(commands.getoutput('type -path cfdisk')): 
	print >> sys.stderr, "Cannot run linHDD without the cfdisk utility"
	raise SystemExit
if not len(commands.getoutput('type -path hdparm')): 
	print >> sys.stderr, "Cannot run linHDD without the hdparm utility"
	raise SystemExit

	
hdd_xpm = [
"42 44 17 1",
" 	g None",
".	g #FFFFFF",
"+	g #0A0A0A",
"@	g #161616",
"#	g #202020",
"$	g #404040",
"%	g #505050",
"&	g #6E6E6E",
"*	g #7D7D7D",
"=	g #2E2E2E",
"-	g #606060",
";	g #9C9C9C",
">	g #C0C0C0",
",	g #C8C8C8",
"'	g #D4D4D4",
")	g #8D8D8D",
"!	g #ACACAC",
"..........................................",
".+++++++++++++++++@@@@@@+++++++++++++++++.",
".++++++++++++++#$%&****&$@+++++++++++++++.",
".+++++++++++@=-;>,,'''',)#+++++++++++++++.",
".++++++++++#&>,''',,,,,,*@+++++++++++++++.",
".++++++++@%!''',>>,>>>>>-#%$#++++++++++++.",
".+++++++#=)'',>>>>>>>>>!$%>>)%#++++++++++.",
".++++++=**-!,>>>>>>>>>>;$&'''>)$@++++++++.",
".+++++=;'>--,,>>>>>>>>>)$;'>,,,!-@++++@@@.",
".++++=;'',!$*>>>>>>>>>,&$>,>>,,',*#+@@@@@.",
".+++#)'',,,;$;,>>>>>>>>--,,>>>>>,,*#@@@@@.",
".+++-''>>>,,&%!>>>>>>>>$)',>>>>>,,>&@@@@@.",
".++=!',>>>>,>-->,>>>>>!$!,,>>>>>>>,>%@@@@.",
".++*',>>>>>>>!%*>>>>>>)%>,>>>>>>>>>,!=+@@.",
".+=!',>>>>>>>>)%;,>>>,*&',>>>>>>>>>>,&@@@.",
".@-,,>>>>>>>>,>&%!,>>>%;'>>>>>>>>>>>>!=+@.",
".#)',>>>>>>>>>,>-->>>!%>,>>>>>>>>>>>>>-@@.",
".=!'>>>>>>>>>>>,;%*,,;-,,>>>>>>>>>>>>,)=@.",
".$>,,>>>>>>>>>>,,)$!,)*',>>>>>>>>>>>>,!$@.",
".%,,>>>>>>>>>>>>>>&%!*;',>>>>>>>>>>>>>!%@.",
".-'>>>>>>>>>>>>>>,>--->,>>>>>>>>>>>>>>>&#.",
".%!;;;;;;;;;;;;;;;;&#->,>>>>>>>>>>>>>>>&#.",
".@=====$$$$%%%%--&&&-)',>>>>>>>>>>>>>>>*#.",
".++#-**));;!!!>>>>>,,',>>>>>>>>>>>>>>>>*#.",
".++$,''''''''',,',',,,,>>>>>>>>>>>>>>>>&=.",
".@+$>',,,,,,>>,,>,>>>>>>>>>>>>>>>>>>>>>&#.",
".@@=!,,,>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>!%#.",
".@@@)'>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>!$#.",
".@@@-,>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>)=#.",
".@@+$!,>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>&##.",
".@@@@),>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>;$##.",
".@@@@%!,>>>>>>>>>>>>>>>>>>>>>>>>>>>>>*=##.",
".+@@@#*>,>>>>>>>>>>>>>>>>>>>>>>>>>>>;$#=#.",
".@@@@@$;>>>>>>>>>>>>>>>>>>>>>>>>>>,!-###=.",
".@@@@@#%>>>>>>>>>>>>>>>>>>>>>>>>>>>*===#=.",
".@@@##@@->>>>>>>>>>>>>>>>>>>>>>>>>*$#=#==.",
".##@@@#@#-!>>>>>>>>>>>>>>>>>>>>>!)==#==#=.",
".@@@#@@#@=%;>,>>>>>>>>>>>>>>>>>!&$#==#===.",
".@#@@#@####$*>>>>>>>>>>>>>>>,>;%===#=====.",
".#@#####@####$*!>>>>>>>>>>>!)-===#=======.",
".#########@####$-);;!!!!;)&%$=#==========.",
".@##@#@#@########=$$%%%$$===#============.",
".#@##################========#===========.",
".........................................."]

class DriveWin:


	def info(self, message):
		box = gtk.MessageDialog(self.window, gtk.DIALOG_MODAL, gtk.MESSAGE_INFO, gtk.BUTTONS_OK, message)
		box.set_position(gtk.WIN_POS_CENTER)
		box.set_title('Information')
		box.run()
		box.destroy()	
		
		
	def destroy_cb(self, widget, main_window):
		main_window.hide()
		gtk.main_quit()


	def get_free(self, prt):
		myline = ""
		myperc = ""
		r = commands.getoutput("df")
		if len(r):
			r1 = string.split(r, "\n")
			for l in r1:
				if re.search(prt, l):
					myline = l
					break
			if len(myline):
				r2 = string.split(myline, " ")
				for l in r2:
					if re.search("%", l):
						myperc = l
						break
				return myperc
				
				
	def get_prt_lst(self, drv):
		self.prt_lst = []
		prts_raw = commands.getoutput('fdisk -l %s' % drv)
		prt_lines = prts_raw.split("\n")
		for l in prt_lines:
			if re.search(drv, l):
				if not re.search("swap", l):
					d = l.split(" ")[0]
					if d != "Disk":
						self.prt_lst.append(d)
		return self.prt_lst


	def get_drives(self):
		drives = []
		for ide_dev in "ide0", "ide1":
			for disk_dev in "hda", "hdb", "hdc", "hdd", "hde", "hdf", "hdg", "hdh":
				if commands.getoutput('cat /proc/ide/%s/%s/media' % (ide_dev, disk_dev))  == "disk":
					drives.append(disk_dev)		
		return drives
	
		
	def change_drive(self, cb):
		drv = cb.get_active_text()
		self.d1 = self.get_params(drv)

		self.prt_box.remove(self.cb_part)
		self.cb_part = gtk.combo_box_new_text()
		plist = self.get_prt_lst(drv)
		for item in plist:
			self.cb_part.append_text(item)
                self.cb_part.set_active(0)			
		self.prt_box.pack_start(self.cb_part, False, False, 0)	
		self.cb_part.connect('changed', self.change_part)
		self.cb_part.show()
		b, e = self.txt.get_bounds()
		self.txt.delete(b, e)
		self.txt.insert_at_cursor(self.d1)
		####################################################################
		act = self.cb_part.get_active_text()
		if act == None:
			if os.getuid() == 0:
				self.btn_part.set_sensitive(True)
				self.cb_format.set_sensitive(False)
				self.tog_check.set_sensitive(False)
			self.perc = "0%"
			self.frac = 0
			self.prg1.set_fraction(0)
			self.prg1.set_text('%s  does not appear to be Partitioned' % self.combo.get_active_text())
		else:
			self.reset_btns()
			try:
				self.perc = self.get_free(plist[0])
				self.frac = float(self.perc.rstrip("%")) / 100
				self.prg1.set_fraction(self.frac)
				self.prg1.set_text("Used Space is " + self.perc)
			except:
				self.perc = "0%"
				self.prg1.set_fraction(0)
				if os.getuid() == 0:
					self.prg1.set_text('%s  is not mounted. See "Partition" for details' % self.cb_part.get_active_text())
				else:
					self.prg1.set_text('%s  is not mounted. You will have to be ROOT to change that.' % self.cb_part.get_active_text())
		#####################################################################


	def reset_btns(self):
		if os.getuid() == 0:
			ret = commands.getoutput('df')
			act = self.cb_part.get_active_text()
			if re.search(act, ret):# partition is mounted
				self.cb_format.set_sensitive(False)
				self.btn_part.set_sensitive(False)
				self.tog_check.set_sensitive(False)
			else:
				self.cb_format.set_sensitive(True)
				self.btn_part.set_sensitive(True)
				self.tog_check.set_sensitive(True)					


	def change_part(self, cb):
		self.reset_btns()
		try:
			self.perc = self.get_free(cb.get_active_text())
			self.frac = float(self.perc.rstrip("%")) / 100	
			self.prg1.set_fraction(self.frac)
			self.prg1.set_text('Used Space in %s is %s' % (cb.get_active_text(), self.perc))
		except:
			self.perc = "0%"
			self.prg1.set_fraction(0)
			if os.getuid() == 0:
				self.prg1.set_text('%s  is not mounted. See "Partition" for details' % cb.get_active_text())
			else:
				self.prg1.set_text('%s  is not mounted. You will have to be ROOT to change that.' % cb.get_active_text())				



	def speed_test(self, btn):
		drv = self.combo.get_active_text()
		self.btn_test.set_label('Testing %s' % drv)
		while gtk.events_pending():
			gtk.main_iteration()
		r = commands.getoutput('hdparm -t %s'% drv)
		if len(r):
			if r[0] == "\n":
				f = r[1:-1]
		self.btn_test.set_label(" Test Drive Speed")
		self.lbl_speed.set_text(f)
			
			
	def get_params(self, drv):
		d_info = commands.getoutput('hdparm -i %s' % drv)
		return d_info

		
	def mk_prt(self, btn):
		# put in warning to nmount is any partition on drv is mounted
		ret = commands.getoutput('df')
		drv = self.combo.get_active_text()
		first = False
		umount_str = ""
		if re.search(drv, ret):	
			r_mounted = []
			r = commands.getoutput('df')
			str_list = []
			str_list.append('You should unmount any partitions on %s \nbefore you partition the drive. Among other reasons, "umount" will not allow you to unmount a drive that is busy, so you do not have to worry about destroying the system by attempting to partition your "boot" drive by accident.\n The following partition(s) are currently mounted:\n\n' % drv)
			p_lines = r.split('\n')
			for l in p_lines:
				if re.search(drv, l): 
					pp = l.split(" ")
					mp = pp[-1]
					for m in pp:
						if re.search(drv, m):
							str_list.append(m + '  is mounted on  %s\n' % mp)
							if not first:
								umount_str = '\nUNMOUNTING: As an example, to unmount %s,\nissue the following command into a term:\n\numount %s' % (m, mp)
								first = True
			
			str_list.append(umount_str)
			ret_str = string.join(str_list, '')
			self.info(ret_str)
			return
	
		try:
			os.system('xterm -e cfdisk %s' % drv)
			self.combo.emit("changed")
			while gtk.events_pending():
				gtk.main_iteration()
		except:
			pass
			
			
	def format_drive(self, cb):
		r = cb.get_active_text()
		if self.tog_check.get_active() == True:
			c = " -c"
		else:
			c = ""
		drv = self.cb_part.get_active_text()
		if r == "Format":
			return
		elif r == "ext2":
			t = "ext2"
			f = 'mke2fs -q%s %s' % (c, drv)
		elif r == "ext3":
			t = "ext3"
			f = 'mke2fs -q%s -j %s' % (c, drv)
		elif r == "reiserfs":
			t = "reiserfs"
			f = 'mkreiserfs -q%s %s' % (c, drv)
		elif r == "xfs":
			t = "xfs"
			f = 'mkfs.xfs -q%s %s' % (c, drv)
		elif r == "jfs":
			t = "jfs"
			f = 'mkfs.jfs -q%s %s' % (c, drv)			
		else:
			return

		if self.do_format(f, t):
			self.info("Format Complete")
			cb.set_active(0)
		else:
			self.info("Drive was not formatted")
			cb.set_active(0)
			
					
			
			
	def do_format(self, cmd, ftype):
		prt = self.cb_part.get_active_text()
		dlg = gtk.MessageDialog(self.window, gtk.DIALOG_MODAL, gtk.MESSAGE_QUESTION, gtk.BUTTONS_YES_NO, 'Are you sure you want to format %s as %s?' % (prt, ftype))
		dlg.set_position(gtk.WIN_POS_CENTER)
		response = dlg.run()
		if response == gtk.RESPONSE_NO:
			dlg.destroy()
			return False
		else:
			dlg.destroy()
			b, e = self.txt.get_bounds()
			self.txt.delete(b, e)
			self.txt.insert_at_cursor('Formatting %s as %s\nPlease Wait...\n' % (prt, ftype))
			while gtk.events_pending():
				gtk.main_iteration()
			pipe = popen2.Popen4(cmd)
			i = 0
			while pipe.poll() < 0:
				gtk.main_iteration()
				time.sleep(1)
				i += 1
				if i > 160:
					self.txt.insert_at_cursor('\n')
					i = 0
				else:
					self.txt.insert_at_cursor('.')
				while gtk.events_pending():
					gtk.main_iteration()
			self.combo.emit("changed")
			while gtk.events_pending():
				gtk.main_iteration()
			return True
			

	def close(self, btn):
		gtk.main_quit()


	def __init__(self):
		self.window = None
		HDDs = self.get_drives()
		if len(HDDs) == 0:
			print >>sys.stderr, "No ide hard drives were found, exiting"
			raise SystemExit
		init_drv = "/dev/" + HDDs[0]
		self.prt_lst = []
		self.prt_lst = self.get_prt_lst(init_drv)
		self.d1 = self.get_params(init_drv)
		self.perc = self.get_free(self.prt_lst[0])
		if self.perc == None:
			self.perc = "0%"
			self.frac = 0
		else:
		
			self.frac = float(self.perc.rstrip("%")) / 100
	
		hdico = gtk.gdk.pixbuf_new_from_xpm_data(hdd_xpm)
		hdd_img = gtk.Image()
		hdd_img.set_from_pixbuf(hdico)


		self.window = gtk.Window( type=gtk.WINDOW_TOPLEVEL)
		self.window.set_title("linHDD - 0.2")
		self.window.set_default_size(530, 426)
		self.window.set_border_width(10)
		self.window.set_position(gtk.WIN_POS_CENTER)
		self.window.connect("destroy", self.destroy_cb, self.window)
		self.window.modify_bg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#4d4d4d"))

		vbox = gtk.VBox(False, 10)
		self.window.add(vbox)
		box1 = gtk.HBox(False, 12)
		box2 = gtk.HBox(False, 0)
		box3 = gtk.HBox(False, 0)
		box4 = gtk.HBox(False, 0)
		box5 = gtk.HBox(False, 0)
		vbox.pack_start(box1, False, False, 0)
		vbox.pack_start(box2, False, False, 0)
		vbox.pack_start(box3, True, True, 0)
		vbox.pack_start(box4, False, False, 0)
		sep = gtk.HSeparator()
		vbox.pack_start(sep, False, False, 0)
		vbox.pack_start(box5, False, False, 0)
		
		box1.pack_start(hdd_img, False, False, 16)
		lbl1 = gtk.Label("")
		lbl1.set_markup('<span size="16384" background="#4d4d4d" foreground="#FFFFFF"><b>Hard Drive Information</b></span>')
		box1.pack_start(lbl1, False, False, 4)


		self.prt_box = gtk.VBox(False, 0)
		lbl_prt = gtk.Label("Partition")
		lbl_prt.modify_fg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#ffffff"))		
		self.cb_part = gtk.combo_box_new_text()
		prts_raw = commands.getoutput('fdisk -l %s' % init_drv)
		prt_lines = prts_raw.split("\n")
		for l in prt_lines:
			if re.search(init_drv, l):
				if not re.search("swap", l):
					d = l.split(" ")[0]
					if d != "Disk":
						self.cb_part.append_text(d)
                self.cb_part.set_active(0)		
		self.prt_box.pack_start(lbl_prt, False, False, 8)		
		self.prt_box.pack_start(self.cb_part, False, False, 0)
		box1.pack_end(self.prt_box, False, False, 0)
				
		
		
		hdd_box = gtk.VBox(False, 0)
		lbl_hdd = gtk.Label("Hard Drive")
		lbl_hdd.modify_fg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#ffffff"))
		self.combo = gtk.combo_box_new_text()
		for l in HDDs:
			nHDD = "/dev/" + l
			self.combo.append_text(nHDD)
                self.combo.set_active(0)
		hdd_box.pack_start(lbl_hdd, False, False, 8)		
		hdd_box.pack_start(self.combo, False, False, 0)
		box1.pack_end(hdd_box, False, False, 0)
		

		self.prg1 = gtk.ProgressBar()
		box2.pack_start(self.prg1, True, True, 0)
		self.prg1.set_fraction(self.frac)
		self.prg1.set_text("Used Space is " + self.perc)
		
		self.txt = gtk.TextBuffer(table=None)
		self.txt_drive_info = gtk.TextView(buffer=self.txt)
		self.txt_drive_info.set_left_margin(8)
		
		swin = gtk.ScrolledWindow()
		swin.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
		swin.add(self.txt_drive_info)
		box3.pack_start(swin, True, True, 0)
		self.txt.insert_at_cursor(self.d1)
		
		
		self.btn_test = gtk.Button(" Test Drive Speed")
		self.btn_test.connect("clicked", self.speed_test)
		self.btn_test.modify_bg(gtk.STATE_PRELIGHT, gtk.gdk.color_parse("#FF3333"))
		self.btn_test.modify_bg(gtk.STATE_ACTIVE, gtk.gdk.color_parse("#ffff66"))
		box4.pack_start(self.btn_test, False, False, 12)
		self.lbl_speed = gtk.Label("Block Device\nResults of read speed here")
		self.lbl_speed.modify_fg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#ffffff"))
		box4.pack_start(self.lbl_speed, False, False, 12)
		
		
		if os.getuid() == 0:
			lbl_avail = gtk.Label("Partition and Format are not\navailable on mounted partitions")
		else:
			lbl_avail = gtk.Label("Partition and Format are not\navailable for non-root users")
		lbl_avail.modify_fg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#ffffff"))
		box5.pack_start(lbl_avail, False, True, 16)


		btns = gtk.HButtonBox()
		box5.pack_end(btns, False, True, 12)

		self.btn_close = gtk.Button(" Close ")
		self.btn_close.connect("clicked", self.close)
		btns.pack_start(self.btn_close, False, True, 12)
		
		self.tog_check = gtk.CheckButton(label="Check for\nbad blocks")
		tog_lbl = self.tog_check.get_children()[0]
		tog_lbl.modify_fg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#FFFFFF"))
		self.tog_check.set_sensitive(False)
		box5.pack_end(self.tog_check, False, True, 12)
		
		self.cb_format = gtk.combo_box_new_text()
		self.cb_format.append_text("Format")
		self.cb_format.append_text("ext2")
		self.cb_format.append_text("ext3")
		self.cb_format.append_text("reiserfs")
		if len(commands.getoutput('type -path mkfs.xfs')):
			self.cb_format.append_text("xfs")
		if len(commands.getoutput('type -path mkfs.jfs')):
			self.cb_format.append_text("jfs")
		self.cb_format.set_active(0)
		self.cb_format.connect('changed', self.format_drive)
		box5.pack_end(self.cb_format, False, True, 12)

		self.btn_part = gtk.Button("Partition")
		self.btn_part.connect("clicked", self.mk_prt)
		box5.pack_end(self.btn_part, False, True, 12)		

		self.cb_format.set_sensitive(False)
		self.btn_part.set_sensitive(False)
		self.combo.connect('changed', self.change_drive)
		self.cb_part.connect('changed', self.change_part)
		
		
		self.window.show_all()

	
def main():
	DriveWin()
	gtk.main()


if __name__ == '__main__':
	main()
