cryptonas-branches/pythonrewrite/plugins/partition/partition.py

import subprocess
import os
import logging
import CryptoBoxTools

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

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

logger = logging.getLogger("CryptoBox")


def doAction(hdf, cbox, **args):
	try:
		step = args["step"]
		del args["step"]
	except KeyError:
		step = "select_device"
	if step == "add_partition":
		return __actionAddPartition(hdf, cbox, args)
	if step == "del_partition":
		return __actionDelPartition(hdf, cbox, args)
	elif step == "finish":
		return __actionFinish(hdf, cbox, args)
	else:	# for "select_device" and for invalid targets
		return __actionSelectDevice(hdf, cbox, args)


def getStatus(cbox):
	return "TODO"


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


def __isDeviceBusy(device, cbox):
	"""check if the device (or one of its partitions) is mounted"""
	# TODO: the config partition is ignored, as it is not part of the container list - that is not good
	import re
	for c in cbox.getContainerList():
		if re.match(device + "\d*$", c.getDevice()):
			if c.isMounted(): return True
	return False


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


def __actionAddPartition(hdf, cbox, args):
	try:
		device = args["block_device"]
	except KeyError:
		return __actionSelectDevice(hdf, cbox, args)
	if not __isDeviceValid(device, cbox): return __actionSelectDevice(hdf, cbox, args)
	if __isDeviceBusy(device, cbox):
		hdf["Data.Warning"] = "Plugins.partition.DiskIsBusy"
		return __actionSelectDevice(hdf, cbox, args)
	size = __getAvailableDeviceSize(device, __withConfigPartition(args))
	hdf["Data.Plugins.partition.Device"] = device
	hdf["Data.Plugins.partition.Device.Size"] = size
	parts = __getPartitionsFromArgs(hdf, args, size)
	__setPartitionData(hdf, parts, size, __withConfigPartition(args))
	return "set_partitions"


def __actionDelPartition(hdf, cbox, args):
	try:
		device = args["block_device"]
		part_num = int(args["del_num"])
	except (TypeError,KeyError):
		return __actionSelectDevice(hdf, cbox, args)
	if not __isDeviceValid(device, cbox): return __actionSelectDevice(hdf, cbox, args)
	if __isDeviceBusy(device, cbox):
		hdf["Data.Warning"] = "Plugins.partition.DiskIsBusy"
		return __actionSelectDevice(hdf, cbox, args)
	size = __getAvailableDeviceSize(device, __withConfigPartition(args))
	hdf["Data.Plugins.partition.Device"] = device
	hdf["Data.Plugins.partition.Device.Size"] = size
	parts = __getPartitionsFromArgs(hdf, args, size)
	## valid partition number to be deleted?
	if part_num < len(parts):
		del parts[part_num]
	else:
		return __actionSelectDevice(hdf, cbox, args)
	__setPartitionData(hdf, parts, size, __withConfigPartition(args))
	return "set_partitions"


def __actionFinish(hdf, cbox, args):
	try:
		device = args["block_device"]
	except KeyError:
		return __actionSelectDevice(hdf, cbox, args)
	if not __isDeviceValid(device, cbox): return __actionSelectDevice(hdf, cbox, args)
	if __isDeviceBusy(device, cbox):
		hdf["Data.Warning"] = "Plugins.partition.DiskIsBusy"
		return __actionSelectDevice(hdf, cbox, args)
	size = __getAvailableDeviceSize(device, __withConfigPartition(args))
	parts = __getPartitionsFromArgs(hdf, args, size)
	if parts:
		if not __runFDisk(cbox, device, parts, __withConfigPartition(args)):
			hdf["Data.Warning"] = "Plugins.partition.PartitioningFailed"
			return __actionAddPartition(hdf, cbox, args)
		else:
			if __formatPartitions(cbox, device, parts, __withConfigPartition(args)):
				hdf["Data.Success"] = "Plugins.partition.Partitioned"
			else:
				hdf["Data.Warning"] = "Plugins.partition.FormattingFailed"
			cbox.reReadContainerList()
			return "form_system"
	else:
		return __actionSelectDevice(hdf, cbox, args)


def __setPartitionData(hdf, parts, size, withConfigPartition):
	availSize = size
	i = 0
	for part in parts:
		logger.debug(part)
		hdf["Data.Plugins.partition.Parts.%d.Size" % i] = part["size"]
		hdf["Data.Plugins.partition.Parts.%d.Type" % i] = part["type"]
		availSize -= part["size"]
		i += 1
	hdf["Data.Plugins.partition.availSize"] = availSize
	if withConfigPartition:
		hdf["Data.Plugins.partition.CreateConfigPartition"] = "1"
	for t in PartTypes.keys():
		hdf["Data.Plugins.partition.Types.%s" % t] = t


def __getPartitionsFromArgs(hdf, args, maxSize):
	parts = []
	done = False
	availSize = maxSize
	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:
				hdf["Data.Warning"] = "Plugins.partition.PartitionTooBig"
				continue
			if int(size) < 10:
				hdf["Data.Warning"] = "Plugins.partition.PartitionTooSmall"
				continue
			if not partType in PartTypes.keys(): continue
			parts.append({"size":size, "type":partType})
			availSize -= size
		except TypeError:
			pass
		except KeyError:
			done = True
	return parts


def __getAvailableDeviceSize(device, withConfigPartition):
	"""calculate the available size (MB) of the device
	also consider a (possible) configuration partition"""
	rdev = os.stat(device).st_rdev
	minor = os.minor(rdev)
	major = os.major(rdev)
	for f in file("/proc/partitions"):
		try:
			elements = f.split()
			if len(elements) != 4: continue
			if (int(elements[0]) == major) and (int(elements[1]) == minor):
				deviceSize = int(elements[2])/1024
				if withConfigPartition: deviceSize -= ConfigPartition["size"]
				return deviceSize
		except ValueError:
			pass
	return 0


def __withConfigPartition(args):
	try:
		if args["create_config_partition"]:
			createConfig = True
		else:
			createConfig = False
	except KeyError:
		createConfig = False
	return createConfig


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


def __getSFDiskLayout(paramParts, withConfigPartition, isFilled):
	parts = paramParts[:]
	## first a (possible) configuration partition - so it will be reusable
	if withConfigPartition:
		## fill the main table (including a config partition)
		yield ",%d,%s" % (ConfigPartition["size"], ConfigPartition["type"])
	## one primary partition
	yield ",%d,%s,*" % (parts[0]["size"], 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 withConfigPartition: yield ";"
	while parts:
		if isFilled and (len(parts) == 1):
			yield ",,%s" % (PartTypes[parts[0]["type"]][0],)
		else:
			yield ",%d,%s" % (parts[0]["size"], PartTypes[parts[0]["type"]][0])
		del parts[0]


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


def __formatOnePartition(cbox, dev_name, type):
	import time, sys
	child_pid = os.fork()
	## we run formatting as a parallel thread
	## TODO: the parent thread still waits for the last child - that is not good for big harddisks
	if child_pid == 0:
		## we are the child process
		proc = subprocess.Popen(
			shell = False,
			stdout = subprocess.PIPE,
			stderr = subprocess.PIPE,
			args = [
				cbox.prefs["Programs"]["super"],
				cbox.prefs["Programs"]["CryptoBoxRootActions"],
				"plugin",
				os.path.join(os.path.dirname(__file__), "root_action.py"),
				"format",
				dev_name,
				type])
		(output, error) = proc.communicate()
		if proc.returncode != 0:
			logger.warn("failed to create filesystem on %s: %s" % (dev_name, error))
			sys.exit(1)
		else:
			sys.exit(0)
	else:
		time.sleep(1)
		(pid, exit_state) = os.waitpid(child_pid, os.WNOHANG)
		if ((pid == 0) and (exit_state == 0)) \
				or ((pid == child_pid) and (exit_state == 0)):
			return True	
		else:
			return False


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