#!/usr/bin/env python
#
# Copyright (c) 2003-2007 Andrea Luzzardi <scox@sig11.org>
#
# This file is part of the pam_usb project. pam_usb is free software;
# you can redistribute it and/or modify it under the terms of the GNU General
# Public License version 2, as published by the Free Software Foundation.
#
# pam_usb 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., 51 Franklin
# Street, Fifth Floor, Boston, MA 02110-1301 USA.

import sys
import os
import gi

gi.require_version('UDisks', '2.0')

from gi.repository import UDisks
from xml.dom import minidom

class Device:
	def __init__(self, udi):
		self.__udi = udi
		deviceObj = udisksObjectManager.get_object(udi)
		driveObj = deviceObj.get_drive()
		if not driveObj.get_property('removable'):
                        # Workaround for removable devices with fixed media (such as SD cards)
                        if not "mmcblk" in udi:
                                raise Exception('Not a removable device')
		self.vendor = driveObj.get_property('vendor')
		self.product = driveObj.get_property('model')
		self.serialNumber = driveObj.get_property('serial')
		if len(self.volumes()) < 1:
			raise Exception('Device does not contain any volume')

	def volumes(self):
		vols = []
		for udi in [o.get_object_path() for o in udisksObjectManager.get_objects() if o.get_block()]:
			obj = udisks.get_object(udi)
			blockObj = obj.get_block()
			if blockObj.get_property('drive') != self.__udi:
				continue
			if not obj.get_filesystem():
				continue
			vols.append({'uuid' : blockObj.get_property('id-uuid'),
				'device' : blockObj.get_property('device')})
		return vols

	def __repr__(self):
		if self.product is not None:
			return "%s %s (%s)" % (self.vendor, self.product, self.serialNumber)
		return self.serialNumber

def listOptions(question, options, autodetect = True):
	if autodetect == True and len(options) == 1:
		print(question)
		print("* Using \"%s\" (only option)" % options[0])
		print()
		return 0

	while True:
		try:
			print(question)
			for i in range(len(options)):
				print( "%d) %s" % (i, options[i]))
			print()
			sys.stdout.write('[%s-%s]: ' % (0, len(options) - 1))
			optionId = int(sys.stdin.readline())
			print
			if optionId not in range(len(options)):
				raise Exception
			return optionId
		except KeyboardInterrupt: sys.exit()
		except Exception: pass
		else: break

def writeConf(options, doc):
	try:
		f = open(options['configFile'], 'w')
		doc.writexml(f)
		f.close()
	except Exception as err:
		print('Unable to save %s: %s' % (options['configFile'], err))
		sys.exit(1)
	else:
		print('Done.')

def shouldSave(options, items):
	print("\n".join(["%s\t\t: %s" % item for item in items]))
	print()
	print('Save to %s ?' % options['configFile'])
	sys.stdout.write('[Y/n] ')
	response = sys.stdin.readline().strip()
	if len(response) > 0 and response.lower() != 'y':
		sys.exit(1)

def prettifyElement(element):
	tmp = minidom.parseString(element.toprettyxml())
	return tmp.lastChild

def addUser(options):
	try:
		doc = minidom.parse(options['configFile'])
	except Exception as err:
		print('Unable to read %s: %s' % (options['configFile'], err))
		sys.exit(1)
	devSection = doc.getElementsByTagName('devices')
	if len(devSection) == 0:
		print('Malformed configuration file: No <devices> section found.')
		sys.exit(1)
	devicesObj = devSection[0].getElementsByTagName('device')
	if len(devicesObj) == 0:
		print('No devices found.')
		print('You must add a device (--add-device) before adding users')
		sys.exit(1)

	devices = []
	for device in devicesObj:
		devices.append(device.getAttribute('id'))
	device = devices[listOptions("Which device would you like to use for authentication ?",
		devices)]

	shouldSave(options, [
		('User', options['userName']),
		('Device', device)
		])

	users = doc.getElementsByTagName('users')
	user = doc.createElement('user')
	user.attributes['id'] = options['userName']
	e = doc.createElement('device')
	t = doc.createTextNode(device)
	e.appendChild(t)
	user.appendChild(e)
	users[0].appendChild(prettifyElement(user))
	writeConf(options, doc)

def addDevice(options):
	devices = []

	for udi in [o.get_object_path() for o in udisksObjectManager.get_objects() if o.get_drive()]:
		try:
			if options['verbose']:
				print('Inspecting %s' % udi)
			devices.append(Device(udi))
		except Exception as ex:
			if options['verbose']:
				print("\tInvalid: %s" % ex)
			pass
		else:
			if options['verbose']:
				print("\tValid")

	if len(devices) == 0:
		print('No devices detected. Try running in verbose (-v) mode to see what\'s going on.')
		sys.exit()
	device = devices[listOptions("Please select the device you wish to add.", devices)]

	volumes = device.volumes()
	volume = volumes[listOptions("Which volume would you like to use for " \
			"storing data ?",
			["%s (UUID: %s)" % (volume['device'],
				volume['uuid'] or "<UNDEFINED>")
				for volume in volumes]
			)]

	if volume['uuid'] == '':
		print('WARNING: No UUID detected for device %s. One time pads will be disabled.' % volume['device'])

	shouldSave(options,[
		('Name', options['deviceName']),
		('Vendor', device.vendor or "Unknown"),
		('Model', device.product or "Unknown"),
		('Serial', device.serialNumber),
		('UUID', volume['uuid'] or "Unknown")
		])

	try:
		doc = minidom.parse(options['configFile'])
	except Exception as err:
		print('Unable to read %s: %s' % (options['configFile'], err))
		sys.exit(1)

	devs = doc.getElementsByTagName('devices')

	# Check that the id of the device to add is not already present in the configFile
	for devices in devs:
		for device_ in devices.getElementsByTagName("device"):
			if device_.getAttribute("id") == options['deviceName']:
				msg =  [ '\nWARNING: A device node already exits for new device \'%s\'.',
						 '\nTo proceed re-run --add-device using a different name or remove the existing entry in %s.' ]
				print('\n'.join(msg) % (options['deviceName'], options['configFile']))
				sys.exit(2)

	dev = doc.createElement('device')
	dev.attributes['id'] = options['deviceName']

	for name, value in (('vendor', device.vendor),
			('model', device.product),
			('serial', device.serialNumber),
			('volume_uuid', volume['uuid'])):
		if value is None or value == '':
			continue
		e = doc.createElement(name)
		t = doc.createTextNode(value)
		e.appendChild(t)
		dev.appendChild(e)

	# Disable one time pads if there's no device UUID
	if volume['uuid'] == '':
		e = doc.createElement('option')
		e.setAttribute('name', 'one_time_pad')
		e.appendChild(doc.createTextNode('false'))
		dev.appendChild(e)

	devs[0].appendChild(prettifyElement(dev))
	writeConf(options, doc)

def usage():
	print('Usage: %s [--help] [--verbose] [--config=path] [--add-user=name | --add-device=name]' % os.path.basename(__file__))
	sys.exit(1)

import getopt

try:
	opts, args = getopt.getopt(sys.argv[1:], "hvd:nu:c:",
			["help", "verbose", "add-device=", "add-user=", "config="])
except getopt.GetoptError:
	usage()

if len(args) != 0:
	usage()

options = { 'deviceName' : None, 'userName' : None,
		'configFile' : '/etc/security/pam_usb.conf', 'verbose' : False }

for o, a in opts:
	if o in ("-h", "--help"):
		usage()
	if o in ("-v", "--verbose"):
		options['verbose'] = True
	if o in ("-d", "--add-device"):
		options['deviceName'] = a
	if o in ("-u", "--add-user"):
		options['userName'] = a
	if o in ("-c", "--config"):
		options['configFile'] = a

if options['deviceName'] is not None and options['userName'] is not None:
	print('You cannot use both --add-user and --add-device')
	usage()

if options['deviceName'] is None and options['userName'] is None:
	usage()

if options['deviceName'] is not None:
	udisks = UDisks.Client.new_sync()
	udisksObjectManager = udisks.get_object_manager()
	try:
		addDevice(options)
	except KeyboardInterrupt:
		sys.exit(1)

if options['userName'] is not None:
	try:
		addUser(options)
	except KeyboardInterrupt:
		sys.exit(1)