#!/usr/bin/env python
# Author: Ricardo Garcia Gonzalez
# License: Public domain code

import ConfigParser
import dbus
import dbus.exceptions
import getpass
import os
import os.path
import sys

# Some constants
PROGRAM_VERSION = '4_beta'

HAL_SERVICE = 'org.freedesktop.Hal'
HAL_MANAGER_OBJ = '/org/freedesktop/Hal/Manager'
HAL_MANAGER_IFACE = 'org.freedesktop.Hal.Manager'
HAL_DEVICE_IFACE = 'org.freedesktop.Hal.Device'
HAL_VOLUME_IFACE = 'org.freedesktop.Hal.Device.Volume'
HAL_CRYPTO_IFACE = 'org.freedesktop.Hal.Device.Volume.Crypto'

CONFIG_FILENAME = '.halrv.conf'
FS_SECTION_PREFIX = 'fs'
UUID_SECTION_PREFIX = 'uuid'
MOUNT_OPTIONS_OPTION = 'mount_options'

# Return user home directory
def get_home():
	home = os.getenv('HOME')
	if home is None:
		sys.exit('ERROR: no HOME environment variable')
	return home

# Create a default config file if none is present
def maybe_create_config_file(home):
	configname = os.path.join(home, CONFIG_FILENAME)
	if os.path.exists(configname):
		return
	try:
		configfile = open(configname, 'w')
		print >>configfile, '[fs:vfat]'
		print >>configfile, 'mount_options: uid=%s' % os.getuid()
		configfile.close()
	except (IOError, OSError), err:
		sys.exit('ERROR: unable to create config file: %s' % configname)

# Parse configuration file
def parse_config_file(home):
	# These will be returned
	fs_opts = dict()
	uuid_opts = dict()

	# Read file
	configname = os.path.join(home, CONFIG_FILENAME)
	config = ConfigParser.SafeConfigParser()
	try:
		config.read([configname])
	except (ConfigParser.Error, ), err:
		sys.exit('ERROR: unable to parse config file: %s' % configname)

	# Process sections and options, populating the dictionaries
	for section in config.sections():
		splitname = section.split(':')
		if splitname[0] not in [FS_SECTION_PREFIX, UUID_SECTION_PREFIX] or len(splitname) != 2 or len(splitname[1]) == 0:
			sys.exit('ERROR: invalid section name in config file: %s' % section)
		if not config.has_option(section, MOUNT_OPTIONS_OPTION):
			sys.exit('ERROR: %s not present in config file section %s' % (MOUNT_OPTIONS_OPTION, section))
		if len(config.get(section, MOUNT_OPTIONS_OPTION)) == 0:
			sys.exit('ERROR: empty %s in config file section %s' % (MOUNT_OPTIONS_OPTION, section))

		mount_opts = config.get(section, MOUNT_OPTIONS_OPTION).split(',')
		if splitname[0] == FS_SECTION_PREFIX:
			fs_opts[splitname[1]] = mount_opts
		else:
			uuid_opts[splitname[1]] = mount_opts

	return (fs_opts, uuid_opts)

# Parse arguments; return tuple (operation, arg)
def parse_args(argv):
	operation = None
	devicename = None

	# 'ls' is the default operation
	if len(argv) == 1:
		operation = 'ls'
	else:
		operation = argv[1]

	# Find operation by unique prefix
	all_ops = ['help', '-?', '-h', '--help', 'version', '--version', 'ls', 'mount', 'umount', 'eject']
	matching_ops = [x for x in all_ops if x.startswith(operation)]
	if len(matching_ops) != 1:
		sys.exit('ERROR: %s: unknown operation; use the "help" operation for a list' % operation)
	operation = matching_ops[0]

	# Operations with no argument
	if operation in ['ls', 'help', '-?', '-h', '--help', 'version', '--version'] and len(argv) > 2:
		sys.exit('ERROR: %s does not expect any argument')

	# Operations with single argument
	if operation in ['mount', 'umount', 'eject']:
		if len(argv) != 3:
			sys.exit('ERROR: %s needs one argument' % operation)
		devicename = argv[2]

	return (operation, devicename)

# Print help text
def print_help():
	progname = os.path.basename(sys.argv[0])
	print 'Usage:'
	print '  %s                   List removable volumes'       % progname
	print '  %s ls                List removable volumes'       % progname
	print '  %s help              Print this help text'         % progname
	print '  %s -?                Print this help text'         % progname
	print '  %s -h                Print this help text'         % progname
	print '  %s --help            Print this help text'         % progname
	print '  %s version           Print program version'        % progname
	print '  %s --version         Print program version'        % progname
	print '  %s mount <device>    Mount the device'             % progname
	print '  %s umount <device>   Unmount the device'           % progname
	print '  %s eject <device>    Unmount and eject the device' % progname

# Print program version
def print_version():
	progname = os.path.basename(sys.argv[0])
	print '%s version %s' % (progname, PROGRAM_VERSION)

# Get the system bus
def get_bus():
	return dbus.SystemBus()

# Get HAL manager
def get_hal_manager(sysbus):
	return dbus.Interface(sysbus.get_object(HAL_SERVICE, HAL_MANAGER_OBJ), HAL_MANAGER_IFACE)

# Get all HAL devices in (busname, Hal.Device interface) dictionary format
def get_devices(sysbus, halman):
	devpaths = halman.GetAllDevices()
	return dict([(x, dbus.Interface(sysbus.get_object(HAL_SERVICE, x), HAL_DEVICE_IFACE)) for x in devpaths])

# Extract HAL removable volumes from device dictionary
def get_volumes(sysbus, devs):
	remvols = dict()
	for (busname, dev) in devs.iteritems():
		try:
			if 'volume' in dev.GetProperty('info.capabilities'):
				# It's a volume; check if container block device is somehow removable
				parent = devs[dev.GetProperty('block.storage_device')]
				if parent.GetProperty('storage.removable') or parent.GetProperty('storage.bus') in ['usb', 'ieee1394']:
					remvols[busname] = dbus.Interface(sysbus.get_object(HAL_SERVICE, busname), HAL_VOLUME_IFACE)
		except dbus.exceptions.DBusException:
			pass
	return remvols

# Get the bus name associated to a device name
def get_bus_name(devs, remvols, devicename):
	matching = [x for x in remvols if devs[x].GetProperty('block.device') == devicename]
	if len(matching) == 0:
		sys.exit('ERROR: no such device: %s' % devicename)
	if len(matching) > 1:
		sys.exit('ERROR: unexpected situation: %s matches more than one device' % devicename)
	return matching[0]

# Setup an encrypted volume
def crypt_setup(hal_device):
	passphrase = getpass.getpass('Encrypted volume detected. Enter passphrase: ')
	crypto_iface = dbus.Interface(hal_device, HAL_CRYPTO_IFACE)
	try:
		crypto_iface.Setup(passphrase)
		print 'Unencrypted volume created. Check device list again.'
	except (dbus.exceptions.DBusException, ), err:
		busname = err.get_dbus_name()
		if busname == 'org.freedesktop.Hal.Device.Volume.Crypto.SetupPasswordError':
			sys.exit('ERROR: wrong password')
		sys.exit('ERROR: unable to set up encrypted volume: %s' % busname)

# Tear down an encrypted volume
def crypt_teardown(hal_device):
	try:
		crypto_iface = dbus.Interface(hal_device, HAL_CRYPTO_IFACE)
		crypto_iface.Teardown()
	except (dbus.exceptions.DBusException, ), err:
		print >>sys.stderr, 'WARNING: unable to tear down encrypted volume: %s' % err.get_dbus_name()

# Perform "ls" operation
def ls_operation(devs, remvols):
	# Header

	# Content to be printed (needed to sort the lines prior to printing them)
	lines = [(	devs[x].GetProperty('block.device'),
			devs[x].GetProperty('volume.label'),
			devs[x].GetProperty('volume.uuid'),
			devs[x].GetProperty('volume.mount_point'),
		) for x in remvols]
	lines.sort()
	lines[:0] = [('Device', 'Label', 'UUID', 'Mount point')]

	# Calculate optimum field length
	lengths = [max([len(line[x]) for line in lines]) for x in xrange(len(lines[0]))]

	# Field separator and horizontal bar
	fieldsep = '  '
	horbar = '-' * (sum(lengths) + (len(fieldsep) * len(lengths) - 1))

	# Optimum line format
	format = fieldsep.join('%%-%ss' % x for x in lengths)

	# Actual printing
	print format % lines[0]
	print horbar
	for line in lines[1:]:
		print format % line

# Perform "mount" operation
def mount_operation(devs, remvols, devicename, fsopts, uuidopts):
	# Get device bus name
	busname = get_bus_name(devs, remvols, devicename)

	# Try to see if it's already mounted
	mountpoint = devs[busname].GetProperty('volume.mount_point')
	if len(mountpoint) != 0:
		sys.exit('ERROR: %s is already mounted on %s' % (devicename, mountpoint))

	# Get filesystem type
	fstype = devs[busname].GetProperty('volume.fstype')
	if len(fstype) == 0:
		sys.exit('ERROR: unable to know filesystem type on %s' % devicename)

	# Setup encrypted volume if appropriate
	if devs[busname].GetProperty('volume.fsusage') == 'crypto':
		crypt_setup(devs[busname])
		return

	# Decide mount point
	label = devs[busname].GetProperty('volume.label')
	if len(label) == 0:
		mountpoint = os.path.basename(devicename)
	else:
		mountpoint = label

	# Get mount options
	uuid = devs[busname].GetProperty('volume.uuid')
	if uuid in uuidopts:
		mountopts = uuidopts[uuid]
	elif fstype in fsopts:
		mountopts = fsopts[fstype]
	else:
		mountopts = dbus.Array([], 's')

	# Try to mount the device
	try:
		remvols[busname].Mount(mountpoint, fstype, mountopts)
		print '%s mounted on %s' % (devicename, devs[busname].GetProperty('volume.mount_point'))
	except (dbus.exceptions.DBusException, ), err:
		sys.exit('ERROR: unable to mount %s: %s' % (devicename, err.get_dbus_name()))

# Perform "umount/eject" operation
def umount_operation(devs, remvols, devicename, do_eject=False):
	last_action = 'unmount'

	# Get device bus name
	busname = get_bus_name(devs, remvols, devicename)

	# Tear down encrypted volume if appropriate
	is_crypto = (devs[busname].GetProperty('volume.fsusage') == 'crypto')
	if is_crypto:
		crypt_teardown(devs[busname])

	# Attempt to umount/eject the device
	try:
		# First, attempt to unmount it
		if not is_crypto:
			try:
				remvols[busname].Unmount(dbus.Array([],'s'))
			except (dbus.exceptions.DBusException, ), err:
				if err.get_dbus_name() == 'org.freedesktop.Hal.Device.Volume.NotMounted':
					if not do_eject:
						print '%s is not mounted' % devicename
					pass
				else:
					raise err

		# Eject it if needed
		if do_eject:
			last_action = 'eject'
			remvols[busname].Eject(dbus.Array([], 's'))
	except (dbus.exceptions.DBusException, ), err:
		print 'ERROR: unable to %s %s: %s' % (last_action, devicename, err.get_dbus_name())

### Main program ###
if __name__ == '__main__':

	# Get home dir
	home = get_home()

	# Create config file if it does not exist
	maybe_create_config_file(home)

	# Parse config file
	(fsopts, uuidopts) = parse_config_file(home)

	# Parse args
	(operation, devicename) = parse_args(sys.argv)

	# Help operation
	if operation in ['help', '-?', '-h', '--help']:
		print_help()
		sys.exit(0)

	# Version operation
	if operation in ['version', '--version']:
		print_version()
		sys.exit(0)

	# Get the system bus
	sysbus = get_bus()

	# Get the HAL manager
	halman = get_hal_manager(sysbus)

	# Get the list of devices
	devs = get_devices(sysbus, halman)

	# Get the list of volumes
	remvols = get_volumes(sysbus, devs)

	# Perform each operation
	if operation == 'ls':
		ls_operation(devs, remvols)
	elif operation == 'mount':
		mount_operation(devs, remvols, devicename, fsopts, uuidopts)
	elif operation == 'umount':
		umount_operation(devs, remvols, devicename)
	else:
		umount_operation(devs, remvols, devicename, do_eject=True)

	# Exit with success
	sys.exit(0)
