import ldap
from struct import pack
from binascii import hexlify

class AmbiguousResult(Exception):
	pass
	
class DoesNotExist(Exception):
	pass

class ADConnection:
	"""
	Creates a connection to the Active Directory, and provides methods to search.
	
	Author: Matt Magin
	Contact: matt.azmoo@gmail.com
	
	Args:
	domain: String, Name of the active directory domain (eg, mydomain.com)
	username: String, Username to connect with.
	password: String, Duh.
	server: String, AD Server to connect to. TODO: Find the most relevant ad server automatically.
	port: Int, Port to connect to AD. Almost always 389.
	timeout: Connection timeout. 3 by default, if it's not set then it will never time out. FAIL.
	"""	
	def __init__(self, domain, username, password, server, port=389, timeout=3):
		self.username = username
		self.password = password
		self.domain = domain
		self.ldap_conn = ldap.initialize('ldap://%s:%d/' % (server, port))
		self.ldap_conn.set_option(ldap.OPT_NETWORK_TIMEOUT, timeout)
		self.ldap_conn.simple_bind_s('%s@%s' % (username, domain), password)
		
	def search_domain(self, ldap_filter):
		""" Args:
		ldap_filter: Filter in LDAP format (ugh!), eg, '(samAccountName=mmagin)'
		"""
		dn = ','.join('DC=%s' % s for s in self.domain.split('.'))
		return ADObjectList(self.ldap_conn.search_ext_s(dn, ldap.SCOPE_SUBTREE, ldap_filter), self)
		
	def search_confined(self, dn, ldap_filter, scope=ldap.SCOPE_SUBTREE):
		""" Args:
		dn: Distinguished Name of base container, eg, 'CN=Users,DC=cmvhodom,DC=cmv'
		ldap_filter: Filter in LDAP format (ugh!), eg, '(samAccountName=mmagin)'
		scope: Scope to search, options are:
				ldap.SCOPE_BASE to search only within the object defined by the dn
				ldap.SCOPE_ONELEVEL to search within the object and first-level children of the dn
				ldap.SCOPE_SUBTREE (default) to search within the object and all children of the dn
		"""
		return ADObjectList(self.ldap_conn.search_ext_s(dn, scope, ldap_filter), self)
		
	def get_single(self, ldap_filter):
		res = self.search_domain(ldap_filter)
		if len(res) > 1:
			raise AmbiguousResult
		else:
			try:
				return res[0]
			except IndexError:
				raise DoesNotExist
			
	def get_user(self, name):
		return self.get_single('(&(samAccountName=%s)(objectClass=user))' % name)
			
	def get_group(self, name):
		return self.get_single('(&(samAccountName=%s)(objectClass=group))' % name)
			
class ADObjectList(list):
	""" Takes a List of LDAP Result objects and converts to a list of ADObjects."""
	def __init__(self, ldap_result, ad_conn):
		self.ad_conn = ad_conn
		for r in ldap_result:
			try:
				self.append(ADObject(r, ad_conn))
			except AttributeError:
				continue

class ADObject:
	""" Converts an LDAP Result Object to a format easier to use."""
	def __init__(self, ldap_object, ad_conn):
		self.ad_conn = ad_conn
		for k in ldap_object[1].keys():
			setattr(self, k, ldap_object[1][k])
	
	def get_name(self):
		if hasattr(self, 'sAMAccountName'):
			return ','.join(self.sAMAccountName)
		else:
			return ''
			
	def get_dn(self):
		if hasattr(self, 'distinguishedName'):
			return ','.join(self.distinguishedName)
		else:
			return ''
			
	def get_primary_group(self):
		""" Gets the primary group for a user. We need this because the primary group 
		is not included in the "memberOf" attribute."""
		# Domain RID (relative identifier) is all of the user's SID (security identifier) 
		# except the last 4 bytes. The RID of the primary group is the 4 byte binary representation
		# of the primaryGroupID (int) attribute. Append the two to get the SID of the primary group.
		group_sid = self.objectSid[0][:-4] + pack('i', int(self.primaryGroupID[0]))
		# Convert it to hex so we can work with it. Binary data FTL.
		hex_sid = hexlify(group_sid)
		# Create an octet string from the hex representation. It's just the hex with a \ every 2 chars.
		octet_string = ''.join(["\%s" % hex_sid[i:i+2] for i in range(0, len(hex_sid), 2)])
		# Search the domain using the octet string and return the result. WIN!
		grp = self.ad_conn.get_single(r'(objectSid=%s)' % octet_string)
		return grp.name[0]
			
	def get_groups(self):
		groups = []
		try:
			groups.append(self.get_primary_group())
			groups = groups + [dn.split(',')[0].split("=")[1] for dn in self.memberOf]
		except AttributeError:
			pass
		return groups