#!/usr/bin/env python3
#
# 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
import subprocess
import socket
import getopt

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 blockObj.get_property('read-only'):
				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):
	if options['yes'] is False:
		print("\n".join(["%s\t\t: %s" % item for item in items]))
		print()
		print('Save to %s? [Y/n]' % options['configFile'])
		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'))

	if options['deviceNumber'] is None:
		device = devices[listOptions("Which device would you like to use for authentication ?", devices)]
	else:
		device = devices[int(options['deviceNumber'])]

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

	# Check if user exists
	users = doc.getElementsByTagName('users')
	userElements = doc.getElementsByTagName('user')
	user = False
	for _user in userElements:
		if _user.getAttribute('id') == options['userName']:
			user = _user
			break

	if user is False: # does not exist, lets create
		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))
	else: # just add another device
		e = doc.createElement('device')
		t = doc.createTextNode(device)
		e.appendChild(t)
		user.appendChild(e)

	writeConf(options, doc)

def listAvailableDevicesAndVolumes(options):
	devices = []

	for udi in [o.get_object_path() for o in udisksObjectManager.get_objects() if o.get_drive()]:
		try:
			devices.append(Device(udi))
		except Exception as ex:
			pass

	if len(devices) == 0:
		print('No devices detected.')
		sys.exit(0)

	if options['deviceNumber'] is None:
		for i in range(len(devices)):
			print( "\"%d) %s\"" % (i, devices[i]), end=',')
	else:
		device = devices[int(options['deviceNumber'])]
		volumes = device.volumes()
		
		for i in range(len(volumes)):
			volume = volumes[i];
			print( "\"%d) %s [%s]\"" % (i, volume['uuid'], volume['device']), end=',')
	
	sys.exit(0)


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()
	
	if options['deviceNumber'] is None:
		device = devices[listOptions("Please select the device you wish to add.", devices)]
	else:
		device = devices[int(options['deviceNumber'])]

	volumes = device.volumes()

	if options['volumeNumber'] is None:
		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
				]
			)
		]
	else:
		volume = volumes[int(options['volumeNumber'])]

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

	vendor = product = 'Generic'
	if hasattr(device, 'vendor'):
		vendor = device.vendor
	else:
		msg =  [ 
			'------------------------------------------------------------------------------------------',
			'WARNING:\tThe device you want to use has no vendor name.',
			'\t\tThis makes using it less secure since the vendor won\'t be checked anymore!',
			'------------------------------------------------------------------------------------------'
		]
		
		print('\n'.join(msg))
	
	if hasattr(device, 'product'):
		product = device.product
	else:
		msg =  [ 
			'------------------------------------------------------------------------------------------',
			'WARNING:\tThe device you want to use has no product name.',
			'\t\tThis makes using it less secure since the vendor won\'t be checked anymore!',
			'------------------------------------------------------------------------------------------'
		]
		
		print('\n'.join(msg))

	shouldSave(
		options,
		[
			('Name', options['deviceName']),
			('Vendor', vendor),
			('Model', product),
			('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']

	vendor = product = 'Generic'
	if hasattr(device, 'vendor'):
		vendor = device.vendor
	
	if hasattr(device, 'product'):
		product = device.product

	for name, value in (
		('vendor', vendor),
		('model', 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)

	# @todo: can this even happen? I've never seen a device without UUID...
	# 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 resetPads():
	padFiles = []

	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.')
		sys.exit(1)

	alreadyConfiguredUsers = doc.getElementsByTagName('user')
	if len(alreadyConfiguredUsers) > 0:
		for user in alreadyConfiguredUsers:
			if user.getAttribute('id') == options['resetPads']:
				# Device specific pad for user
				userDevice = user.getElementsByTagName('device')[0].firstChild.nodeValue
				padFiles.append('/home/%s/.pamusb/%s.pad' % (options['resetPads'], userDevice))

				# User specific pad for device
				for details in devicesObj:
					if details.getAttribute('id') == userDevice:
						deviceUUID = details.getElementsByTagName('volume_uuid')[0].firstChild.nodeValue
						try:
							deviceMountPoint = subprocess.check_output(['mount | grep `readlink -f /dev/disk/by-uuid/%s` | awk \'{print $3}\'' % (deviceUUID)], shell=True)
							padFiles.append('%s/.pamusb/%s.%s.pad' % (deviceMountPoint.decode().strip(), options['resetPads'], socket.gethostname()))
						except subprocess.CalledProcessError as err:
							print('Could not get device mountpoint: ', err)

	if len(padFiles) == 0:
		print('Could not find pad files to reset, wtf?!')
		sys.exit(1)

	print('Pad files to remove: ', padFiles)
	print('Removing device specific pad for user %s...' % options['resetPads'])
	try:
		os.remove(padFiles[0])
	except FileNotFoundError as fnf:
		print('    Notice: pad file not found.')

	print('Removing user and hostname specific pad on device with UUID %s...' % deviceUUID)
	try:
		os.remove(padFiles[1])
	except FileNotFoundError as fnf:
		print('    Notice: pad file not found.')

	print('')
	print('All pad files for the given user were deleted, on next successful authentication they will be regenerated (you can force this by running pamusb-check).')
	sys.exit(0)

def usage():
	print('Version 0.8.5')
	print('Usage: %s [--help] [--verbose] [--yes] [--config=path] [--reset-pads=username] [--add-user=name | --add-device=name [[--device=number] [--volume=number]]' % os.path.basename(__file__))
	sys.exit(1)

try:
	opts, args = getopt.getopt(
		sys.argv[1:], 
		"hvd:nu:c:y",
		["help", "verbose", "add-device=", "add-user=", "config=", "reset-pads=", "yes", "device=", "volume=", "list-devices", "list-volumes"]
	)
except getopt.GetoptError:
	usage()

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

options = {
	'deviceName': None,
	'userName': None,
	'configFile': '/etc/security/pam_usb.conf',
	'verbose': False,
	'yes': False,
	'deviceNumber': None,
	'volumeNumber': None,
	'listDevices': False,
	'listVolumes': False,
	'resetPads': None
}

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 o in ("-y", "--yes"):
		options['yes'] = True
	if o in ("--device"):
		options['deviceNumber'] = a
	if o in ("--volume"):
		options['volumeNumber'] = a
	if o in ("--list-devices"):
		options['listDevices'] = True
	if o in ("--list-volumes"):
		options['listVolumes'] = True
	if o in ("--reset-pads"):
		options['resetPads'] = a

if options['resetPads'] is not None:
	resetPads()

if options['listDevices'] is True or options['listVolumes'] is True:
	udisks = UDisks.Client.new_sync()
	udisksObjectManager = udisks.get_object_manager()
	listAvailableDevicesAndVolumes(options)

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)
