#
# Copyright 2006 sense.lab e.V.
# 
# This file is part of the CryptoBox.
#
# The CryptoBox 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.
# 
# The CryptoBox 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 the CryptoBox; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

import subprocess
import os
import logging
import cryptobox.core.tools as cbxTools
import cryptobox.plugins.base

class partition(cryptobox.plugins.base.CryptoBoxPlugin):

	pluginCapabilities = [ "system" ]
	pluginVisibility = [ "preferences" ]
	requestAuth = True
	rank = 80

	PartTypes = {
		"windows"	: ["0xC", "vfat"],
		"linux"	: ["L", "ext3"]}

	ConfigPartition = {
		"size"	: 5,	# size of configuration partition (if necessary) in MB
		"type"	: "L",
		"fs"	: "ext2"}


	def doAction(self, **args):
		## load default hdf values
		self.__prepareDataset()
		## retrieve some values from 'args' - defaults are empty
		self.device = self.__getSelectedDevice(args)
		self.withConfigPartition = self.__isWithConfigPartition()
		self.cbox.log.debug("partition plugin: selected device=%s" % str(self.device))
		self.deviceSize = self.__getAvailableDeviceSize(self.device)
		try:
			step = args["step"]
			del args["step"]
		except KeyError:
			step = "select_device"
		try:
			## this way of selecting the easy setup is necessary: see select_device.py for details
			if args["easy"]: step = "easy"
		except KeyError:
			pass
		## no (or invalid) device was supplied
		if not self.device:
			step = "select_device"
		if step == "add_partition":
			return self.__actionAddPartition(args)
		elif step == "del_partition":
			return self.__actionDelPartition(args)
		elif step == "finish":
			return self.__actionFinish(args)
		elif step == "easy":
			return self.__actionEasySetup(args)
		else:	# for "select_device" and for invalid targets
			return self.__actionSelectDevice(args)


	def getStatus(self):
		return "%s / %s / %s" % (self.device, self.deviceSize, self.withConfigPartition)


	def __prepareDataset(self):
		self.hdf[self.hdf_prefix + "StyleSheetFile"] = os.path.join(self.pluginDir, "partition.css")


	def __getSelectedDevice(self, args):
		try:
			device = args["block_device"]
		except KeyError:
			return None
		if not self.__isDeviceValid(device):
			return None
		if self.__isDeviceBusy(device):
			self.hdf["Data.Warning"] = "Plugins.partition.DiskIsBusy"
			return None
		return device


	def __isDeviceValid(self, device):
		if not device:
			return False
		if not self.cbox.isDeviceAllowed(device):
			return False
		if not device in cbxTools.getParentBlockDevices():
			return False
		return True


	def __isDeviceBusy(self, device):
		"""check if the device (or one of its partitions) is mounted"""
		# the config partition is ignored, as it will get unmounted if necessary
		import re
		for c in self.cbox.getContainerList():
			if re.match(device + "\d*$", c.getDevice()):
				if c.isMounted(): return True
		return False


	def __actionSelectDevice(self, args):
		block_devices = [e
			for e in cbxTools.getParentBlockDevices()
			if self.cbox.isDeviceAllowed(e)]
		counter = 0
		for a in block_devices:
			self.hdf[self.hdf_prefix + "BlockDevices.%d.name" % counter] = a
			self.hdf[self.hdf_prefix + "BlockDevices.%d.size" % counter] = cbxTools.getBlockDeviceSizeHumanly(a)
			self.cbox.log.debug("found a suitable block device: %s" % a)
			counter += 1
		if self.withConfigPartition:
			self.hdf[self.hdf_prefix + "CreateConfigPartition"] = "1"
		## there is no disk available
		if not block_devices:
			self.hdf["Data.Warning"] = "Plugins.partition.NoDisksAvailable"
		return "select_device"


	def __actionAddPartition(self, args):
		self.hdf[self.hdf_prefix + "Device"] = self.device
		self.hdf[self.hdf_prefix + "Device.Size"] = self.deviceSize
		parts = self.__getPartitionsFromArgs(args)
		self.__setPartitionData(parts)
		return "set_partitions"
	

	def __actionDelPartition(self, args):
		try:
			part_num = int(args["del_num"])
		except (TypeError,KeyError):
			return self.__actionAddPartition(args)
		self.hdf[self.hdf_prefix + "Device"] = self.device
		self.hdf[self.hdf_prefix + "Device.Size"] = self.deviceSize
		parts = self.__getPartitionsFromArgs(args)
		## valid partition number to be deleted?
		if part_num < len(parts):
			del parts[part_num]
		else:
			return self.__actionAddPartition(args)
		self.__setPartitionData(parts)
		return "set_partitions"


	def __actionFinish(self, args):
		parts = self.__getPartitionsFromArgs(args)
		if parts:
			self.__setPartitionData(parts)
			if cbxTools.isPartOfBlockDevice(self.device, self.cbox.prefs.getActivePartition()):
				self.cbox.prefs.umountPartition()
			if not self.__runFDisk(parts):
				self.hdf["Data.Warning"] = "Plugins.partition.PartitioningFailed"
				self.cbox.log.warn("partition: failed to partition device: %s" % self.device)
				return self.__actionAddPartition(args)
			else:
				"""
				tricky problem: if the device was partitioned, then a created config partition is still part of the containerlist, as the label is not checked again - very ugly!!! So we will call reReadContainerList after formatting the last partition - see below
				"""
				self.cbox.reReadContainerList()
				def result_generator():
					counter = 0
					## initialize the generator
					formatPart_gen = self.__formatPartitions(parts)
					while counter < len(parts):
						## first part: get the device name
						yield formatPart_gen.next()
						counter += 1
						## second part: do the real formatting of a partition
						result = formatPart_gen.next()
						## after the first partiton, we can reRead the containerList (as the possible config partition was already created)
						if self.withConfigPartition and (counter == 1):
							## important: reRead the containerList - but somehow it breaks the flow (hanging process)
							#self.cbox.reReadContainerList()
							## write config data
							self.cbox.prefs.mountPartition()
							self.cbox.prefs.write()
							self.cbox.log.info("settings stored on config partition")
						## return the result
						if result:
							yield "OK"
						else:
							yield "<b>Error</b>"
				return {
					"template": "show_format_progress",
					"generator": result_generator}
		else:
			return self.__actionAddPartition(args)


	def __actionEasySetup(self, args):
		import types
		## we do not have to take special care for a possible config partition
		parts = [ { "size": self.deviceSize, "type": "windows" } ]
		## umount partition if necessary
		if cbxTools.isPartOfBlockDevice(self.device, self.cbox.prefs.getActivePartition()):
			self.cbox.prefs.umountPartition()
		## partition it
		if not self.__runFDisk(parts):
			self.hdf["Data.Warning"] = "Plugins.partition.PartitioningFailed"
			return None
		## "formatPartitions" is a generator, returning device names and bolean values
		result = [e for e in self.__formatPartitions(parts) if type(e) == types.BooleanType]
		if self.withConfigPartition:
			self.cbox.prefs.mountPartition()
			self.cbox.prefs.write()
		## check if there is a "False" return value
		if False in result:
			## operation failed
			self.hdf["Data.Warning"] = "Plugins.partition.FormattingFailed"
			self.cbox.log.info("easy partitioning failed")
			return "select_partitions"
		else:
			## operation was successful
			self.hdf["Data.Success"] = "Plugins.partition.EasySetup"
			self.cbox.log.info("easy partitioning succeeded")
			## do not show the disk overview immediately - it does not get updated that fast
			return { "plugin":"system_preferences", "values":[] }


	def __setPartitionData(self, parts):
		availSize = self.deviceSize
		i = 0
		for part in parts:
			self.cbox.log.debug(part)
			self.hdf[self.hdf_prefix + "Parts.%d.Size" % i] = part["size"]
			self.hdf[self.hdf_prefix + "Parts.%d.Type" % i] = part["type"]
			availSize -= part["size"]
			i += 1
		self.hdf[self.hdf_prefix + "availSize"] = availSize
		if self.withConfigPartition:
			self.hdf[self.hdf_prefix + "CreateConfigPartition"] = "1"
		for t in self.PartTypes.keys():
			self.hdf[self.hdf_prefix + "Types.%s" % t] = t
		## store the currently existing partitions of the choosen block device
		current_containers = [ e for e in self.cbox.getContainerList() if cbxTools.isPartOfBlockDevice(self.device, e.getDevice()) ]
		for (index, t) in enumerate(current_containers):
			self.hdf[self.hdf_prefix + "ExistingContainers.%d.name" % index] = t.getName()
			self.hdf[self.hdf_prefix + "ExistingContainers.%d.size" % index] = cbxTools.getBlockDeviceSizeHumanly(t.getDevice())


	def __getPartitionsFromArgs(self, args):
		parts = []
		done = False
		availSize = self.deviceSize
		i = -1
		while not done:
			i += 1
			try:
				size = int(args["part%d_size" % i])
				partType = args["part%d_type" % i]
				if int(size) > availSize:
					self.hdf["Data.Warning"] = "Plugins.partition.PartitionTooBig"
					continue
				if int(size) < 10:
					self.hdf["Data.Warning"] = "Plugins.partition.PartitionTooSmall"
					continue
				if not partType in self.PartTypes.keys(): continue
				parts.append({"size":size, "type":partType})
				availSize -= size
			except TypeError:
				pass
			except KeyError:
				done = True
		return parts


	def __getAvailableDeviceSize(self, device):
		"""calculate the available size (MB) of the device
		also consider a (possible) configuration partition"""
		deviceSize = cbxTools.getBlockDeviceSize(device)
		if deviceSize < 0: return 0
		if self.withConfigPartition:
			deviceSize -= self.ConfigPartition["size"]
		return deviceSize


	def __isWithConfigPartition(self):
		"""check if we have to create a configuration partition"""
		if self.cbox.prefs.requiresPartition():
			active = self.cbox.prefs.getActivePartition()
			## we need a partition, if there is no active one
			if not active: return True
			## check if the active one is part of the current device
			return cbxTools.isPartOfBlockDevice(self.device, active)
		return False


	def __runFDisk(self, parts):
		## check if the device is completely filled (to avoid some empty last blocks)
		avail_size = self.deviceSize
		for d in parts: avail_size -= d["size"]
		self.cbox.log.debug("remaining size: %d" % avail_size)
		isFilled = avail_size == 0
		proc = subprocess.Popen(
			shell = False,
			stdin = subprocess.PIPE,
			stdout = subprocess.PIPE,
			stderr = subprocess.PIPE,
			args = [
				self.cbox.prefs["Programs"]["super"],
				self.cbox.prefs["Programs"]["CryptoBoxRootActions"],
				"plugin",
				os.path.join(self.pluginDir, "root_action.py"),
				"partition",
				self.device])
		for line in self.__getSFDiskLayout(parts, isFilled):
			proc.stdin.write(line + "\n")
		(output, error) = proc.communicate()
		if proc.returncode != 0: self.cbox.log.debug("partitioning failed: %s" % error)
		return proc.returncode == 0


	def __getSFDiskLayout(self, paramParts, isFilled):
		"""this generator returns the input lines for sfdisk"""
		parts = paramParts[:]
		## first a (possible) configuration partition - so it will be reusable
		if self.withConfigPartition:
			## fill the main table (including a config partition)
			yield ",%d,%s" % (self.ConfigPartition["size"], self.ConfigPartition["type"])
		## one primary partition
		if isFilled and (len(parts) == 1):
			## fill the rest of the device
			yield ",,%s,*" % self.PartTypes[parts[0]["type"]][0]
		else:
			## only use the specified size
			yield ",%d,%s,*" % (parts[0]["size"], self.PartTypes[parts[0]["type"]][0])
		del parts[0]
		## no extended partition, if there is only one disk
		if not parts: return
		## an extended container for the rest
		yield ",,E"
		## an empty partition in main table
		yield ";"
		## maybe another empty partition if there is no config partition
		if not self.withConfigPartition: yield ";"
		while parts:
			if isFilled and (len(parts) == 1):
				yield ",,%s" % (self.PartTypes[parts[0]["type"]][0],)
			else:
				yield ",%d,%s" % (parts[0]["size"], self.PartTypes[parts[0]["type"]][0])
			del parts[0]


	def __formatPartitions(self, paramParts):
		parts = paramParts[:]
		part_num = 1
		## maybe a config partition?
		if self.withConfigPartition:
			dev_name = self.device + str(part_num)
			self.cbox.log.info("formatting config partition (%s)" % dev_name)
			if self.__formatOnePartition(dev_name, self.ConfigPartition["fs"]):
				self.__setLabelOfPartition(dev_name, self.cbox.prefs["Main"]["ConfigVolumeLabel"])
			part_num += 1
		## the first data partition
		dev_name = self.device + str(part_num)
		partType = self.PartTypes[parts[0]["type"]][1]
		self.cbox.log.info("formatting partition (%s) as '%s'" % (dev_name, partType))
		yield dev_name
		yield self.__formatOnePartition(dev_name, partType)
		del parts[0]
		## other data partitions
		part_num = 5
		while parts:
			dev_name = self.device + str(part_num)
			partType = self.PartTypes[parts[0]["type"]][1]
			self.cbox.log.info("formatting partition (%s) as '%s'" % (dev_name, partType))
			yield dev_name
			yield self.__formatOnePartition(dev_name, partType)
			part_num += 1
			del parts[0]
		return


	def __formatOnePartition(self, dev_name, type):
		## first: retrieve UUID - it can be removed from the database afterwards
		volDB = self.cbox.prefs.volumesDB
		prev_name = [e.getName() for e in self.cbox.getContainerList() if e.getDevice() == dev_name]
		## call "mkfs"
		proc = subprocess.Popen(
			shell = False,
			args = [
				self.cbox.prefs["Programs"]["super"],
				self.cbox.prefs["Programs"]["CryptoBoxRootActions"],
				"plugin",
				os.path.join(self.pluginDir, "root_action.py"),
				"format",
				dev_name,
				type])
		(output, error) = proc.communicate()
		if proc.returncode != 0:
			self.cbox.log.warn("failed to create filesystem on %s: %s" % (dev_name, error))
			return False
		else:
			## remove unused volume entry
			if prev_name:
				self.cbox.prefs.volumesDB[prev_name[0]]
			return True


	def __setLabelOfPartition(self, dev_name, label):
		proc = subprocess.Popen(
			shell = False,
			stdout = subprocess.PIPE,
			stderr = subprocess.PIPE,
			args = [
				self.cbox.prefs["Programs"]["super"],
				self.cbox.prefs["Programs"]["CryptoBoxRootActions"],
				"plugin",
				os.path.join(self.pluginDir, "root_action.py"),
				"label",
				dev_name,
				label])
		(output, error) = proc.communicate()
		if proc.returncode == 0:
			return True
		else:
			self.cbox.log.warn("failed to create filesystem on %s: %s" % (dev_name, error))
			return False