#!/usr/bin/env python
# CCS Computer Science
#   Common
#
# Douglas Thrift
#
# $Id$

import getpass
from M2Crypto import SSL
import os
import socket
import struct

HANDSHAKE = ('CCSCSLAB', (494, 119))
HOSTS = map(lambda host: host + '.ccs.ucsb.edu', ('zweihander', 'katana'))
CHANGE = 1
RESET = 2
ADD = 3
CHFN = 4
CHSH = 5
BOOT = 6
PORT = 2748
SSL_ROOT = '/ccs/ssl'

class Error(Exception):
	def __init__(self, error):
		self.error = error

	def __str__(self):
		return self.error

	def __unicode__(self):
		return str(self)

class HandshakeError(Error):
	def __init__(self):
		self.error = 'Invalid handshake.'

class PasswordError(Error):
	pass

_HOST = socket.gethostname().split('.', 1)[0]
HOST = socket.getfqdn(_HOST)
SYSTEMS = (
	'freebsd',
	'linux',
	'darwin',
	'solaris',
	'netbsd',
	'debian',
	'openbsd',
	'dragonfly',
)

class Shells(object):
	def __init__(self, *args):
		if args:
			assert len(SYSTEMS) == len(args)

			self.shells = dict(zip(SYSTEMS, args))
		else:
			self.shells = dict((system, None) for system in SYSTEMS)

	def __cmp__(self, other):
		return cmp(self.shells, other.shells)

	def __len__(self):
		return len(self.shells)

	def __getitem__(self, system):
		return self.shells[system]

	def __setitem__(self, system, shell):
		self.shells[system] = shell

	def __iter__(self):
		for system in SYSTEMS:
			yield self.shells[system]

SHELLS = (
	('sh', Shells(
		'/bin/sh',				# FreeBSD
		'/bin/sh',				# Linux
		'/bin/sh',				# Darwin
		'/usr/bin/sh',			# Solaris
		'/bin/sh',				# NetBSD
		'/bin/sh',				# Debian
		'/bin/sh',				# OpenBSD
		'/bin/sh',				# DragonFly
	)),
	('csh', Shells(
		'/bin/csh',				# FreeBSD
		'/bin/csh',				# Linux
		'/bin/csh',				# Darwin
		'/usr/bin/csh',			# Solaris
		'/bin/csh',				# NetBSD
		'/bin/csh',				# Debian
		'/bin/csh',				# OpenBSD
		'/bin/csh',				# DragonFly
	)),
	('tcsh', Shells(
		'/bin/tcsh',			# FreeBSD
		'/bin/tcsh',			# Linux
		'/bin/tcsh',			# Darwin
		'/usr/bin/tcsh',		# Solaris
		'/usr/pkg/bin/tcsh',	# NetBSD
		'/usr/bin/tcsh',		# Debian
		'/usr/local/bin/tcsh',	# OpenBSD
		'/bin/tcsh',			# DragonFly
	)),
	('bash', Shells(
		'/usr/local/bin/bash',	# FreeBSD
		'/bin/bash',			# Linux
		'/bin/bash',			# Darwin
		'/usr/local/bin/bash',	# Solaris
		'/usr/pkg/bin/bash',	# NetBSD
		'/bin/bash',			# Debian
		'/usr/local/bin/bash',	# OpenBSD
		'/usr/pkg/bin/bash',	# DragonFly
	)),
	('ksh', Shells(
		'/usr/local/bin/ksh',	# FreeBSD
		'/bin/ksh',				# Linux
		'/bin/ksh',				# Darwin
		'/usr/bin/ksh',			# Solaris
		'/bin/ksh',				# NetBSD
		'/usr/bin/ksh',			# Debian
		'/bin/ksh',				# OpenBSD
		'/usr/pkg/bin/ksh',		# DragonFly
	)),
	('zsh', Shells(
		'/usr/local/bin/zsh',	# FreeBSD
		'/bin/zsh',				# Linux
		'/bin/zsh',				# Darwin
		'/usr/bin/zsh',			# Solaris
		'/usr/pkg/bin/zsh',		# NetBSD
		'/usr/bin/zsh',			# Debian
		'/usr/local/bin/zsh',	# OpenBSD
		'/usr/pkg/bin/zsh',		# DragonFly
	)),
	('mksh', Shells(
		'/usr/local/bin/mksh',	# FreeBSD
		'/bin/mksh',			# Linux
		'/sw/bin/mksh',			# Darwin
		'/usr/local/bin/mksh',	# Solaris
		'/usr/pkg/bin/mksh',	# NetBSD
		'/bin/mksh',			# Debian
		'/usr/local/bin/mksh',	# OpenBSD
		'/usr/pkg/bin/mksh',	# DragonFly
	)),
	('custom', Shells()),
)

class _Encode(object):
	def __init__(self, function):
		self.function = function

	def __call__(self, *args):
		def encode(value):
			if isinstance(value, basestring):
				return value.encode('ascii')
			elif isinstance(value, Shells):
				return Shells(*map(encode, value))
			else:
				return value

		return self.function(*map(encode, args))

def _context():
	context = SSL.Context()

	context.set_verify(SSL.verify_peer | SSL.verify_fail_if_no_peer_cert, 2)
	context.load_cert(certfile = os.path.join(SSL_ROOT, _HOST + '.crt'), keyfile = os.path.join(SSL_ROOT, _HOST + '.key'))

	cacert = os.path.join(SSL_ROOT, 'ccscert.pem')

	if not os.path.exists(cacert):
		cacert = os.path.join(SSL_ROOT, 'ccscert.crt')

	context.load_verify_locations(cafile = cacert)

	return context

def _command(command, host):
	connection = SSL.Connection(_context())

	connection.connect((host, PORT))

	input, output = connection.makefile('rb', -1), connection.makefile('wb', 0)

	output.write(HANDSHAKE[0])
	output.flush()

	if struct.unpack('<HB', input.read(3)) != HANDSHAKE[1]:
		raise HandshakeError()

	command(output = output)

	output.flush()

	status = struct.unpack('<B', input.read(1))[0]

	if status != 0:
		raise Error(input.read())

	return command(input = input)

def _command_all(command, *args):
	exceptions = {}

	for host in HOSTS if command != chsh else HOSTS[:1]:
		if host != HOST:
			try:
				command(host, *args)
			except Exception, exception:
				exceptions[host] = exception

	return exceptions

@_Encode
def change(host, user, old_password, new_password):
	def _change(output = None, input = None):
		if output is None:
			return

		output.write(struct.pack('<BB', CHANGE, len(user)) + user + struct.pack('<H', len(old_password)) + old_password + struct.pack('<H', len(new_password)) + new_password)

	_command(_change, host)

def change_all(user, old_password, new_password):
	return _command_all(change, user, old_password, new_password)

@_Encode
def reset(host, user, password):
	def _reset(output = None, input = None):
		if output is None:
			return

		output.write(struct.pack('<BB', RESET, len(user)) + user + struct.pack('<H', len(password)) + password)

	_command(_reset, host)

def reset_all(user, password):
	return _command_all(reset, user, password)

@_Encode
def add(host, user, name, password):
	def _add(output = None, input = None):
		if output is None:
			return

		output.write(struct.pack('<BB', ADD, len(user)) + user + struct.pack('<B', len(name)) + name + struct.pack('<H', len(password)) + password)

	_command(_add, host)

def add_all(user, name, password):
	return _command_all(add, user, name, password)

@_Encode
def chfn(host, user, name):
	def _chfn(output = None, input = None):
		if output is None:
			return

		output.write(struct.pack('<BB', CHFN, len(user)) + user + struct.pack('<B', len(name)) + name)

	_command(_chfn, host)

def chfn_all(user, name):
	return _command_all(chfn, user, name)

@_Encode
def chsh(host, user, shell, shells):
	def _chsh(output = None, input = None):
		if output is None:
			return

		output.write(struct.pack('<BB', CHSH, len(user)) + user + struct.pack('<B', len(shell)) + shell)

		if shell != 'custom':
			output.write(struct.pack('<B', 0))
		else:
			output.write(struct.pack('<B', len(shells)))

			for _shell in shells:
				output.write(struct.pack('<B', len(_shell)) + _shell)

	_command(_chsh, host)

def chsh_all(user, shell, shells):
	return _command_all(chsh, user, shell, shells)

def boot(host):
	def _boot(output = None, input = None):
		if input is None:
			output.write(struct.pack('<B', BOOT))
		else:
			return struct.unpack('<d', input.read(8))[0]

	return _command(_boot, host)

def new_password(password = 'New Password', confirm = 'Confirm Password'):
	password = getpass.getpass('%s: ' % password)

	if not password:
		raise PasswordError('Empty password.')

	if password != getpass.getpass('%s: ' % confirm):
		raise PasswordError('Password and confirm password mismatched.')

	return password

def capitalize(system):
	if system.endswith('bsd'):
		return system[:-3].capitalize() + 'BSD'
	elif system == 'dragonfly':
		return 'DragonFly'

	return system.capitalize()

# vim: noexpandtab tabstop=4