#!/usr/bin/env python

#   FILE: gnome-luks-format -- 
# AUTHOR: W. Michael Petullo <mike@flyn.org>
#   DATE: 01 April 2005
#
# Copyright (C) 2005 W. Michael Petullo <mike@flyn.org>
# All rights reserved.
#
# 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

try:
        import os
	import sys
	import time
	import getopt
	import gtk
	import gtk.glade
	import gobject
	import dbus
except ImportError:
	print "Error importing dbus, pygtk2 and pygtk2-libglade"
	sys.exit(1)

class Cipher:
	def name(self):
		pass
	def keyBitCount(self):
		pass

class AES256(Cipher):
	def name(self):
		return "aes"
	def keyBitCount(self):
		return 256

class AES128(Cipher):
	def name(self):
		return "aes"
	def keyBitCount(self):
		return 128

class FS:
	def name(self):
		pass

class Ext3(FS):
	def name(self):
		return "ext3"

class Ext2(FS):
	def name(self):
		return "ext2"

class Vfat(FS):
	def name(self):
		return "vfat"

class Format:
	def asyncFormat(self):
		pass

class NullFormat(Format):
	def __init__(self, fs, cipher, passphrase, volName, device,
		     checkBlocks = False, isVerbose = False):
		self.hadError = False

		self.fs = fs
		self.ciper = cipher
		self.passphrase = passphrase
		self.volName = volName
		self.device = device
		self.checkBlocks = checkBlocks
		self.isVerbose = isVerbose

	def asyncFormat(self):
		pid = os.fork()
		if pid == 0:
			time.sleep(5)
			sys.exit(0)
		else:
			return pid

class LuksFormat(Format):
	def __init__(self, fs, cipher, passphrase, volName, device,
  		     checkBlocks = False, isVerbose = False):
		self.hadError = False

		self.fs = fs
		self.cipher = cipher
		self.passphrase = passphrase
		self.volName = volName
		self.device = device
		self.checkBlocks = checkBlocks
		self.isVerbose = isVerbose

	def asyncFormat(self):
		# Execute format process and provide password.
		pipefd = os.pipe()
		pid = os.fork()
		if pid == 0:
			try:
				os.close(0)
				os.dup(pipefd[0])
				os.close(pipefd[0])
				os.close(pipefd[1])

				os.execv("/usr/sbin/luks-format", self.buildCommand())
			except:
				sys.exit(1)
		else:
			try:
				os.close(pipefd[0])
				os.write(pipefd[1], self.passphrase)
				os.close(pipefd[1]);

				return pid
			except:
				return -1

	def buildCommand(self):

		cmd = [ "luks-format", "--fs-type=" + self.fs.name(), 
		        "--fs-cipher=" + self.cipher.name(), "--fs-keylen=" +
			str(self.cipher.keyBitCount()) ]
		if self.checkBlocks:
			cmd.append("-c")
		if self.isVerbose:
			cmd.append("-v")
		cmd.append("--volume-name=" + self.volName)
		cmd.append(self.device)

		return cmd

class GnomeLuksFormat:
	def __init__(self, d = None):
		self.glade = gtk.glade.XML("/usr/share/luks-tools/gnome-luks-format.glade")
		handlers = {
			"onTopLevelDeleteEvent": self.onTopLevelDeleteEvent,
			"onAES128Activate": self.onAES128Activate,
			"onAES256Activate": self.onAES256Activate,
			"onExt2Activate": self.onExt2Activate,
			"onExt3Activate": self.onExt3Activate,
			"onVfatActivate": self.onVfatActivate,
			"onCloseButtonClicked": self.onCloseButtonClicked,
			"onHelpButtonClicked": self.onHelpButtonClicked,
			"onErrOkButtonClicked": self.onErrOkButtonClicked,
			"onDevSelCancelButtonClicked": self.onDevSelCancelButtonClicked,
			"onFormatButtonClicked": self.onFormatButtonClicked,
			"onOpenButtonClicked": self.onOpenButtonClicked,
			"onDevSelOkButtonClicked": self.onDevSelOkButtonClicked
		} 

		self.glade.signal_autoconnect(handlers)

		# Default values.
		self.isDryRun = False
		self.isVerbose = False
		self.cipher = AES256()
		self.fs = Ext3()
		self.passphrase = None
		self.volName = None

		# References to various Windows used by application.
		self.topLevel = self.glade.get_widget("topLevel")
		self.devSelUI = self.glade.get_widget("devSelUI")
		self.errUI = self.glade.get_widget("errUI")
		self.progressUI = self.glade.get_widget("progressUI")
		self.device = d

		deviceMenu = self.glade.get_widget("optionMenuDevice")

		if d != None:
			menu = gtk.Menu()
			self.addToAvailableDevices(menu, d)
			deviceMenu.set_menu(menu)

			l = self.glade.get_widget("labelDevice")
			l.set_sensitive(False)

			deviceMenu.set_sensitive(False)

			l = self.glade.get_widget("displayedDevice")
			l.set_text(d)
			l.set_sensitive(False)
		else:
			menu = gtk.Menu()
			self.setupAvailableDevices(menu)
			deviceMenu.set_menu(menu)

			deviceMenu.set_sensitive(True)
		
			l = self.glade.get_widget("displayedDevice")
			l.set_text("none selected")

	def onDeviceSelected(self, item, data):
		self.device = data

	def addToAvailableDevices(self, menu, device):
		item = gtk.MenuItem(device)
		item.connect("activate", self.onDeviceSelected, device)
		item.show()

		menu.insert(item, 0)

	# See FindDeviceByCapability in system-config-packages for similar
	# code
	def setupAvailableDevices(self, menu):
		# Find available block devices.
		bus = dbus.Bus.get_system()
		hal_manager = bus.get_object("org.freedesktop.Hal",
		                             "/org/freedesktop/Hal/Manager")
		blockDevs = hal_manager.FindDeviceByCapability("block",
			           dbus_interface="org.freedesktop.Hal.Manager")

		for devUri in blockDevs:
		        dev = bus.get_object ("org.freedesktop.Hal", devUri)
			devIface = dbus.Interface(dev,
			                          "org.freedesktop.Hal.Device")
			devicePath = devIface.GetPropertyString ("block.device")
			self.addToAvailableDevices(menu, devicePath)

		# FIXME: what should be default?  do we need to ask are you sure?
		self.device = devicePath

	def error(self, msg):
		l = self.glade.get_widget("labelErr")
		l.set_text(msg)
		self.topLevel.set_sensitive(False)
		self.errUI.show()

	def onTopLevelDeleteEvent(self, event, notKnown = None):
		gtk.main_quit()
		sys.exit(0)

	def onAES128Activate(self, event):
		self.cipher = AES128()

	def onAES256Activate(self, event):
		self.cipher = AES256()

	def onExt2Activate(self, event):
		self.fs = Ext2()

	def onExt3Activate(self, event):
		self.fs = Ext3()

	def onVfatActivate(self, event):
		self.fs = Vfat()

	def onCloseButtonClicked(self, event):
		gtk.main_quit()
		sys.exit(0)

	def onHelpButtonClicked(self, event):
		# FIXME: implement help
		pass

	def onErrOkButtonClicked(self, event):
		self.errUI.hide()
		self.topLevel.set_sensitive(True)

	def onDevSelCancelButtonClicked(self, event):
		self.devSelUI.hide()
		self.topLevel.set_sensitive(True)

	def topLevelInputOk(self):
		fnval = False;

		if self.device == None:
			self.error("Device not selected")
		elif self.passphrase == "":
			self.error("Passphrase not entered")
		elif self.volName == "":
			self.error("Volume name not entered")
		else:
			fnval = True;

		return fnval

	def pulser(self, pid):
		# This is ugly, but fork/waitpid does not seem to work if both
		# are used in a non-main thread
		try:
			waitPid, self.formatStatus = os.waitpid(pid, os.WNOHANG)
			# Process still exists
			p = self.glade.get_widget("progressBarFormat")
			p.pulse();
			# As this is a timeout function, return True so that it
			# continues to get called
			return True
		except OSError:
			# Process gone, formatStatus should have been set in
			# previous call.  Should have received a no more children
			# exception.
			if os.WEXITSTATUS(self.formatStatus):
				self.error("Error executing luks-format")
			self.progressUI.hide()
			self.topLevel.set_sensitive(True)
			return False

	def onFormatButtonClicked(self, event):
		entry = self.glade.get_widget("entryPassphrase")
		self.passphrase = entry.get_text()

		entry = self.glade.get_widget("entryVolumeName")
		self.volName = entry.get_text()

		if self.topLevelInputOk():
			self.topLevel.set_sensitive(False)
			self.progressUI.show()

			l = self.glade.get_widget("labelFormat")

			if not self.isDryRun:
				l.set_text("Formatting " + self.device)
				c = self.glade.get_widget("checkButtonBadBlocks");
				formatter = LuksFormat(self.fs, self.cipher, 
						       self.passphrase,
						       self.volName,
						       self.device,
						       c.get_active(),
						       self.isVerbose)
			else:
				l.set_text("[Simulated] Formatting " + self.device)
				formatter = NullFormat(self.fs, self.cipher,
						       self.passphrase,
						       self.volName,
						       self.device)

			pid = formatter.asyncFormat()

			if pid == -1:
				self.error("Error executing luks-format")
			else:
				p = gobject.idle_add(self.pulser, pid)

	def onOpenButtonClicked(self, event):
		self.topLevel.set_sensitive(False)
		self.devSelUI.show()

	def devSelInputOk(self, path):
		fnval = False

		try:
			os.stat(path)
			fnval = True
		except:
			pass

		return fnval

	def onDevSelOkButtonClicked(self, event):
		tmp = self.devSelUI.get_filename()
		if self.devSelInputOk(tmp):
			self.device = tmp

			l = self.glade.get_widget("displayedDevice");
			l.set_text(self.device);

			self.devSelUI.hide();
			self.topLevel.set_sensitive(True);

	def setDryRun(self, setting):
		self.isDryRun = setting;

	def setVerbose(self, setting):
		self.isVerbose = setting;

def printUsage(status, error=None, more=None):
	sys.stderr.write("gnome-luks-setup [options] [device]\n\n" +
	  "-h, --help\n" +
	  "      print a list of options\n\n" +
	  "-v, --verbose\n" +
	  "      print information to stdout\n\n" +
	  "-d, --dry-run\n" +
	  "      do not actually perform any operations on device\n\n")
	if error != None:
		sys.stderr.write(error + ": " + more + "\n");
	sys.exit(status);

def main():
	optDryRun = False
	optVerbose = False
	try:
		opts, args = getopt.getopt(sys.argv[1:], "hvd", 
		                           ["help", "verbose", "dry-run"])
	except getopt.GetoptError:
		# print help information and exit:
		printUsage(1)

	for o, a in opts:
		if o in ("-d", "--dry-run"):
			optDryRun = True
		elif o in ("-v", "--verbose"):
			optVerbose = True
		elif o in ("-h", "--help"):
			printUsage(1)
		else:
			printUsage(1)
	
	if len(args) == 0:
		gui = GnomeLuksFormat()
	elif len(args) == 1:
		gui = GnomeLuksFormat(args[0])
	else:
		printUsage(1)

	gui.setDryRun(optDryRun)
	gui.setVerbose(optVerbose)
	gtk.main()

if __name__ == "__main__":
	main()
