#!/usr/bin/env python2
#
# 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 os
import sys
import pwd
import getopt
import syslog
import gi

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

from gi.repository import GLib
from gi.repository import UDisks

try:
	# Python 2.5
	import xml.etree.ElementTree as et
except ImportError:
	# Python 2.4
	try:
		import cElementTree as et
	except ImportError:
		import elementtree.ElementTree as et

class HotPlugDevice:
	def __init__(self, serial):
		self.__udi = None
		self.__serial = serial
		self.__callbacks = []
		self.__running = False

	def run(self):
		self.__scanDevices()
		self.__registerSignals()
		self.__running = True
		GLib.MainLoop().run()
		print 'signals registered'

	def addCallback(self, callback):
		self.__callbacks.append(callback)

	def __scanDevices(self):
		for udi in udisksObjectManager.get_objects():
			if udi.get_block():
				device = udisks.get_drive_for_block(udi.get_block())
				if device:
					self.__deviceAdded(device)

	def __registerSignals(self):
		for signal, callback in (('object-added', self.__objectAdded),
				('object-removed', self.__objectRemoved)):
			udisksObjectManager.connect(signal, callback)

	def __objectAdded(self, _, udi):
		if udi.get_block():
			device = udisks.get_drive_for_block(udi.get_block())
			if device:
				self.__deviceAdded(device)

	def __objectRemoved(self, _, udi):
		if udi.get_block():
			device = udisks.get_drive_for_block(udi.get_block())
			if device:
				self.__deviceRemoved(device)

	def __deviceAdded(self, udi):
		if self.__udi is not None:
			return
		if udi.get_property('serial') != self.__serial:
			return
		self.__udi = udi
		if self.__running:
			[ cb('added') for cb in self.__callbacks ]

	def __deviceRemoved(self, udi):
		if self.__udi is None:
			return
		if self.__udi != udi:
			return
		self.__udi = None
		if self.__running:
			[ cb('removed') for cb in self.__callbacks ]

class Log:
	def __init__(self):
		syslog.openlog('pamusb-agent', syslog.LOG_PID | syslog.LOG_PERROR,
				syslog.LOG_AUTH)

	def info(self, message):
		self.__logMessage(syslog.LOG_NOTICE, message)

	def error(self, message):
		self.__logMessage(syslog.LOG_ERR, message)

	def __logMessage(self, priority, message):
		syslog.syslog(priority, message)

def usage():
	print 'Usage: %s [--help] [--config=path] [--daemon] [--check=path]' % \
			os.path.basename(__file__)
	sys.exit(1)

import getopt

try:
	opts, args = getopt.getopt(sys.argv[1:], "hc:dc:",
			["help", "config=", "daemon", "check="])
except getopt.GetoptError:
	usage()

options = {'configFile' : '/etc/pamusb.conf',
		'daemon' : False,
		'check' : '/usr/bin/pamusb-check'}

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

for o, a in opts:
	if o in ('-h', '--help'):
		usage()
	if o in ('-c', '--config'):
		options['configFile'] = a
	if o in ('-d', '--daemon'):
		options['daemon'] = True
	if o in ('-c', '--check'):
		options['check'] = a


if not os.path.exists(options['check']):
	print '%s not found.' % options['check']
	print "You might specify manually pamusb-check's location using --check."
	usage()

username = pwd.getpwuid(os.getuid())[0]

logger = Log()

doc = et.parse(options['configFile'])
users = doc.findall('users/user')
for user in users:
	if user.get('id') == username:
		break
else:
	logger.error('User %s not found in configuration file' % username)
	sys.exit(1)

events = {
		'lock' : [],
		'unlock' : []
		}

for hotplug in user.findall('agent'):
	events[hotplug.get('event')].append(hotplug.text)

deviceName = user.find('device').text.strip()

devices = doc.findall("devices/device")
for device in devices:
	if device.get('id') == deviceName:
		break
else:
	logger.error('Device %s not found in configurtion file' % deviceName)
	sys.exit(1)

serial = device.find('serial').text.strip()

def authChangeCallback(event):
	if event == 'removed':
		logger.info('Device "%s" has been removed, ' \
				'locking down user "%s"...' % (deviceName, username))
		for cmd in events['lock']:
			logger.info('Running "%s"' % cmd)
			os.system(cmd)
		logger.info('Locked.')
		return

	logger.info('Device "%s" has been inserted. ' \
			'Performing verification...' % deviceName)
	cmdLine = "%s --quiet --config=%s --service=pamusb-agent %s" % (
			options['check'], options['configFile'], username)
	logger.info('Executing "%s"' % cmdLine)
	if not os.system(cmdLine):
		logger.info('Authentication succeeded. ' \
				'Unlocking user "%s"...' % username)
		for cmd in events['unlock']:
			logger.info('Running "%s"' % cmd)
			os.system(cmd)
		logger.info('Unlocked.')
	else:
		logger.info('Authentication failed for device %s. ' \
				'Keeping user "%s" locked down.' % (deviceName, username))

udisks = UDisks.Client.new_sync()
udisksObjectManager = udisks.get_object_manager()

hpDev = HotPlugDevice(serial)
hpDev.addCallback(authChangeCallback)

if options['daemon'] and os.fork():
	sys.exit(0)

logger.info('pamusb-agent up and running.')
logger.info('Watching device "%s" for user "%s"' % (deviceName, username))

try:
	hpDev.run()
except KeyboardInterrupt:
	logger.error('Caught keyboard interruption, exiting...')