#!/usr/bin/env python

from __future__ import with_statement
import sys, os, time, re
import threading, mutex
import SocketServer, StringIO
import ldap, ldif

debug = 0

OPERATIONS_ERROR 		= 0x01
PROTOCOL_ERROR			= 0x02
TIMELIMIT_EXCEEDED		= 0x03
SIZELIMIT_EXCEEDED		= 0x04
COMPARE_FALSE			= 0x05
COMPARE_TRUE			= 0x06
AUTH_METHOD_NOT_SUPPORTED	= 0x07
STRONG_AUTH_NOT_SUPPORTED	= 0x07
STRONG_AUTH_REQUIRED		= 0x08
STRONGER_AUTH_REQUIRED		= 0x08
PARTIAL_RESULTS			= 0x09
ADMINLIMIT_EXCEEDED		= 0x0b
CONFIDENTIALITY_REQUIRED	= 0x0d
SASL_BIND_IN_PROGRESS		= 0x0e
NO_SUCH_ATTRIBUTE		= 0x10
UNDEFINED_TYPE			= 0x11
INAPPROPRIATE_MATCHING		= 0x12
CONSTRAINT_VIOLATION		= 0x13
TYPE_OR_VALUE_EXISTS		= 0x14
INVALID_SYNTAX			= 0x15
NO_SUCH_OBJECT			= 0x20
ALIAS_PROBLEM			= 0x21
INVALID_DN_SYNTAX		= 0x22
IS_LEAF				= 0x23
ALIAS_DEREF_PROBLEM		= 0x24
X_PROXY_AUTHZ_FAILURE		= 0x2F
INAPPROPRIATE_AUTH		= 0x30
INVALID_CREDENTIALS		= 0x31
INSUFFICIENT_ACCESS		= 0x32
BUSY				= 0x33
UNAVAILABLE			= 0x34
UNWILLING_TO_PERFORM		= 0x35
LOOP_DETECT			= 0x36
NAMING_VIOLATION		= 0x40
OBJECT_CLASS_VIOLATION		= 0x41
NOT_ALLOWED_ON_NONLEAF		= 0x42
NOT_ALLOWED_ON_RDN		= 0x43
ALREADY_EXISTS			= 0x44
NO_OBJECT_CLASS_MODS		= 0x45
RESULTS_TOO_LARGE		= 0x46
AFFECTS_MULTIPLE_DSAS		= 0x47
OTHER_ERROR			= 0x50

def split_argument(line):
	parts = line.strip().split(':', 2)
	name = parts[0].strip()
	value = ""
	if len(parts) == 2:
		value = parts[1].strip()
	return (name, value)


class Parser(ldif.LDIFParser):
	def __init__(self, input):
		ldif.LDIFParser.__init__(self, input)
		self.dn = None
		self.entry = None
	def handle(self, dn, entry):
		if self.entry is None:
			self.dn = dn
			self.entry = entry


class Database:
	def __init__(self, suffix):
		self.suffix = suffix
		self.binddn = None
		self.remote = None
		self.mutex = threading.Lock()

	# Overridable to handle specific command
	def add(self, dn, entry):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Add not implemented")
	def bind(self, dn, args):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Bind not implemented")
	def compare(self, dn, entry):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Compare not implemented")
	def delete(self, dn, args):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Delete not implemented")
	def modify(self, dn, modlist):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Modify not implemented")
	def modrdn(self, dn, args):
		raise VirtualError, (UNWILLING_TO_PERFORM, "ModRDN not implemented")
	def search(self, dn, args):
		raise VirtualError, (UNWILLING_TO_PERFORM, "Search not implemented")

	# Overridable to handle all processing
	def process(self, command, dn, block):
		try:

			# This we handle specially
			if command == "MODIFY":
				return self.process_modify_internal(dn, block)

			# This we handle specially
			elif command == "ADD":
				return self.process_add_internal(dn, block)

			# All the rest we split up, multiple args go into arrays
			args = Arguments()
			while True:
				line = block.readline()
				if len(line) == 0:
					break
				(name, value) = split_argument(line)
				args.add(name, value)

			if command == "BIND":
				self.bind(dn, args)
			elif command == "COMPARE":
				result = self.compare(dn, args)
				return (result, "", None)
			elif command == "DELETE":
				self.delete(dn, args)
			elif command == "MODIFY":
				self.modify(dn, args)
			elif command == "MODRDN":
				self.modrdn(dn, args)
			elif command == "SEARCH":
				return self.process_search_internal(dn, args)
			elif command == "UNBIND":
				assert False # should have been handled in caller
			else:
				return UNWILLING_TO_PERFORM, "Unsupported operation %s" % command

			return (0, "", [])

		except Error, ex:
			return (ex.code, ex.info, None)

	def process_locked(self, command, dn, block, binddn, remote):
		with self.mutex:
			self.binddn = binddn
			self.remote = remote
			result = self.process(command, dn, block)
		return result

	def process_add_internal(self, dn, block):
		parser = Parser(block)
		parser.parse()
		self.add(dn, parser.entry)
		return (0, "", [])


	def process_compare_intersal(self, dn, block):
		parser = Parser(block)
		parser.parse()
		result = self.compare(dn, parser.entry)
		return (result, "", [])


	def process_search_internal(self, dn, args):
		results = []
		data = self.search(args["base"] or dn, args)
		for (dn, entry) in data:
			result = StringIO.StringIO()
			writer = ldif.LDIFWriter(result)
			writer.unparse(dn, entry)
			results.append(result.getvalue())
		return (0, "", results)


	def process_modify_internal(self, dn, block):

		op = None
		attr = None
		batch = 0
		mods = []

		while True:
			line = block.readline()
			if len(line) == 0:
				break
			line = line.strip()

			# A break between different mods
			if line == "-":

				# Latch batch was empty, delete/replace all
				if batch == 0:
					if op == ldap.MOD_DELETE:
						mods.append((op, attr, None))
					elif op == ldap.MOD_REPLACE:
						mods.append((op, attr, None))

				batch = 0
				op = None
				attr = None
				continue

			# The current line
			(name, value) = split_argument(line)

			# Don't have a mod type yet
			if op is None:
				attr = value
				if name == "add":
					op = ldap.MOD_ADD
				elif name == "replace":
					op = ldap.MOD_REPLACE
				elif name == "delete":
					op = ldap.MOD_DELETE
				else:
					op = None

			# Have a op, add values
			elif name == attr:
				assert attr is not None
				mods.append((op, attr, value))
				batch += 1

		self.modify(dn, mods)
		return (0, "", [])


class Error(Exception):
	"""Exception to be returned to server"""

	def __init__(self, code = OPERATIONS_ERROR, info = ""):
		self.code = code
		self.info = info

	def __str__(self):
		return "%d: %s" % (self.code, self.info)

class Arguments:
	def __init__(self):
		self.dict = {}
	def add(self, name, value):
		if self.dict.has_key(name):
			if type(self.dict[name]) != type([]):
				self.dict[name] = [self.dict[name]]
				self.dict[name].append(value)
			else:
				self.dict[name] = value

		self.dict[name] = value
	def __getitem__(self, name):
		if self.dict.has_key(name):
			return self.dict[name]
		return None
	def __len__(self):
		return len(self.dict)

class Connection(SocketServer.BaseRequestHandler):

	def __init__(self, request, client_address, server):
		server.unique += 1
		self.identifier = server.unique
		self.block_regex = re.compile("\r?\n[ \t]*\r?\n")
		SocketServer.BaseRequestHandler.__init__(self, request, client_address, server)

	def trace(self, message):
		global debug
		if debug:
			prefix = "%04d " % self.identifier
			lines = message.split("\n")
			print >> sys.stderr, prefix + lines[0]
			print >> sys.stderr, "\n".join([prefix + "*** " + line for line in lines[1:]])

	def handle(self):

		self.trace("CONNECTED")

		req = self.request
		req.setblocking(1)
		req.settimeout(None)

		extra = ""
		block = None
		while True:
			data = extra + req.recv(1024)

			# End of connection
			if len(data) == 0:
				break

			parts = self.block_regex.split(data)
			if len(parts) > 1:
				block = unicode(parts[0], "utf-8", "strict")
				break
			extra = parts[0]

		if block:
			self.trace("REQUEST\n%s\n" % block)
			self.handle_block(req, StringIO.StringIO(block))

		self.trace("DISCONNECTING")
		self.request.close()

	def handle_block(self, req, block):

		line = block.readline()
		command = line.strip()

		# Disconnect immediately on certain occasions
		if not command:
			return False
		elif command == "UNBIND":
			return False

		suffixes = []
		binddn = None
		remote = None
		ssf = None
		msgid = None
		dn = None

		while True:
			off = block.tell()

			line = block.readline()
			(name, value) = split_argument(line)

			if name == "suffix":
				suffixes.append(value)
			elif name == "msgid" and msgid is None:
				msgid = value
			elif name == "dn" and dn is None:
				dn = value.lower()
			elif name == "binddn" and binddn is None:
				binddn = value.lower()
			elif name == "peername" and remote is None:
				remote = value
			elif name == "ssf" and ssf is None:
				ssf = value
			else: # Return this line and continue
				block.seek(off)
				break

		code = 0
		info = ""
		data = None

		if len(suffixes) == 0:
			code = OPERATIONS_ERROR
			info = "No suffix specified"
		elif len(suffixes) > 1:
			code = OPERATIONS_ERROR
			info = "Multiple suffixes not supported"
		else:
			database = self.server.find_database(suffixes[0])
			(code, info, data) = database.process_locked(command, dn, block, binddn, remote)

		if data:
			for dat in data:
				self.trace("DATA\n%s\n" % dat)
				req.sendall(dat.strip("\n") + "\n\n")

		result = "RESULT"
		if info:
			result += "\ninfo: %s" % info
		# BUG: Current OpenLDAP always wants code last
		result += "\ncode: %d" % code
		self.trace("RESPONSE\n%s" % result.strip("\n"))
		req.sendall(result)

		return True


class Server(SocketServer.ThreadingMixIn, SocketServer.UnixStreamServer):
	daemon_threads = True
	allow_reuse_address = True

	def __init__(self, address, DatabaseClass, debug = False):
		if (os.path.exists(address)):
			os.unlink(address)
		SocketServer.UnixStreamServer.__init__(self, address, Connection)
		self.DatabaseClass = DatabaseClass
		self.__databases = {}
		self.unique = 0

	def find_database(self, suffix):
		suffix = suffix.lower()
		if self.__databases.has_key(suffix):
			database = self.__databases[suffix]
		else:
			database = self.DatabaseClass(suffix)
			self.__databases[suffix] = database
		return database