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

import re
import subprocess
import imp
import os
import logging
import CryptoBoxTools

PartTypes = {
	"linux"	: "L",
	"windows"	: "0xC"}

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):
	## TODO: also check for "is device busy" add output a warning
	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"""
	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
	## 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 = __getDeviceSize(device)
	hdf["Data.Plugins.partition.Device"] = device
	hdf["Data.Plugins.partition.Device.Size"] = size
	parts = __getPartitionsFromArgs(args, size)
	__setPartitionData(hdf, parts, size)
	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 = __getDeviceSize(device)
	hdf["Data.Plugins.partition.Device"] = device
	hdf["Data.Plugins.partition.Device.Size"] = size
	parts = __getPartitionsFromArgs(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)
	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 = __getDeviceSize(device)
	parts = __getPartitionsFromArgs(args, size)
	if parts:
		if not __runFDisk(cbox, device, parts):
			hdf["Data.Warning"] = "Plugins.partition.PartitioningFailed"
			return __actionAddPartition(hdf, cbox, args)
		else:
			hdf["Data.Success"] = "Plugins.partition.Partitioned"
			cbox.reReadContainerList()
			return "form_system"
	else:
		return __actionSelectDevice(hdf, cbox, args)


def __setPartitionData(hdf, parts, size):
	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
	for t in PartTypes.keys():
		hdf["Data.Plugins.partition.Types.%s" % t] = t


def __getPartitionsFromArgs(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: continue
			if int(size) <= 0: 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 __getDeviceSize(device):
	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):
				return int(elements[2])/1024
		except ValueError:
			pass
	return 0


def __runFDisk(cbox, device, parts):
	proc = subprocess.Popen(
		shell = False,
		stdin = 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])
	import logging
	logger = logging.getLogger("CryptoBox")
	for line in __getSFDiskLayout(parts): proc.stdin.write(line + "\n")
	(output, error) = proc.communicate()
	if error: logger.debug("partitioning failed: %s" % error)
	return proc.returncode == 0


def __getSFDiskLayout(paramParts):
	parts = paramParts[:]
	## first a primary partition
	yield ",%d,%s,*" % (parts[0]["size"], PartTypes[parts[0]["type"]])
	del parts[0]
	if not parts: return
	yield ",,E"		# extended container for the rest
	yield ";"		# empty partition in main table
	yield ";"		# another empty partition in main table
	while parts:
		yield ",%d,%s" % (parts[0]["size"], PartTypes[parts[0]["type"]])
		del parts[0]