#!/usr/bin/python

import ldap
import ldap.modlist as modlist
import logging

class LdapServer(object):

	uri = None
	dn = None
	pwd = None
	v2 = None

	con = 0

	def __init__(self,uri,dn=None,pwd=None,v2=None,raiseOnError=False, logger=False):
		self.uri  = uri
		self.dn   = dn
		self.pwd  = pwd
		self.raiseOnError  = raiseOnError
		if v2:
			self.v2=True
		if logger:
			self.logger = logger
		else:
			self.logger = logging.getLogger()

	def _error(self,error,level=logging.WARNING):
		if self.raiseOnError:
			raise LdapServerException(error)
		else:
			self.logger.log(level,error)

	def connect(self):
		if self.con == 0:
			try:
				con = ldap.initialize(self.uri)
				if self.v2:
					con.protocol_version = ldap.VERSION2
				else:
					con.protocol_version = ldap.VERSION3

				if self.dn:
					con.simple_bind_s(self.dn,self.pwd)

				self.con = con
				return True
			except ldap.LDAPError, e:
				self._error('LdapServer - Error connecting and binding to LDAP server : %s' % e,logging.CRITICAL)
			return False
		return True

	def get_scope(self, scope):
		if scope == 'base':
			return ldap.SCOPE_BASE
		elif scope == 'one':
			return ldap.SCOPE_ONELEVEL
		elif scope == 'sub':
			return ldap.SCOPE_SUBTREE
		raise Exception("Unknown LDAP scope '%s'" % scope)

	def search(self, basedn, filterstr, attrs, sizelimit=0, scope='sub'):
		res_id = self.con.search(basedn, self.get_scope(scope), filterstr, attrs)
		ret = {}
		c=0
		while 1:
			res_type, res_data = self.con.result(res_id,0)
			if res_data == [] or sizelimit!=0 and c>sizelimit:
				break
			else:
				if res_type == ldap.RES_SEARCH_ENTRY:
					ret[res_data[0][0]]=res_data[0][1]
			c=c+1
		return ret

	def add_object(self,dn,attrs):
		ldif = modlist.addModlist(attrs)
		try:
			self.logger.debug("LdapServer - Add %s" % dn)
			self.con.add_s(dn,ldif)
			return True
		except ldap.LDAPError, e:
			self._error("LdapServer - Error adding %s : %s" % (dn,e), logging.error)

		return False

	def update_object(self, dn, old, new, ignore_attrs=[]):
		ldif = modlist.modifyModlist(old, new, ignore_attr_types=ignore_attrs)
		if ldif == []:
			return True
		try:
			self.con.modify_s(dn,ldif)
			return True
		except ldap.LDAPError, e:
			self._error("LdapServer - Error updating %s : %s\nOld : %s\nNew : %s" % (dn, e, old, new), logging.error)
		return False

	def update_need(self, old, new, ignore_attrs=[]):
		ldif = modlist.modifyModlist(old, new, ignore_attr_types=ignore_attrs)
		if ldif == []:
			return False
		return True

	def rename_object(self,dn,new_rdn):
		try:
			self.logger.debug("LdapServer - Rename %s in %s" % (dn,new_rdn))
			self.con.rename_s(dn,new_rdn)
			return True
		except ldap.LDAPError, e:
			self._error("LdapServer - Error renaming %s in %s : %s" % (dn,new_rdn,e), logging.error)

		return False

	def drop_object(self,dn):
		try:
			self.logger.debug("LdapServer - Delete %s" % dn)
			self.con.delete_s(dn)
			return True
		except ldap.LDAPError, e:
			self._error("LdapServer - Error deleting %s : %s" % (dn,e), logging.error)

		return False

	def get_dn(self,obj):
		return obj[0][0]

	def get_attr(self,obj,attr,all=None,default=None):
		if attr not in obj:
			for k in obj:
				if k.lower() == attr.lower():
					attr = k
					break
		if all is not None:
			if attr in obj:
				return obj[attr]
			else:
				return default or []
		else:
			if attr in obj:
				return obj[attr][0]
			else:
				return default

class LdapServerException(BaseException):
	def __init__(self,msg):
		BaseException.__init__(self, msg)