Merge branch 'master' of git://github.com/macdiesel/SleekXMPP into hacks

6 files changed, 100 insertions, 32 deletions

diff --git a/sleekxmpp/__init__.py b/sleekxmpp/__init__.py
index 74eb290a..3379e45d 100644
--- a/sleekxmpp/__init__.py
+++ b/sleekxmpp/__init__.py
@@ -27,6 +27,10 @@ import sys
 import random
 import copy
 from . import plugins
+from xml.etree.cElementTree import tostring
+from xml.etree.cElementTree import Element
+from cStringIO import StringIO
+
 #from . import stanza
 srvsupport = True
 try:
@@ -71,8 +75,14 @@ class ClientXMPP(basexmpp, XMLStream):
 		self.sessionstarted = False
 		self.bound = False
 		self.bindfail = False
-		self.registerHandler(Callback('Stream Features', MatchXPath('{http://etherx.jabber.org/streams}features'), self._handleStreamFeatures, thread=True))
-		self.registerHandler(Callback('Roster Update', MatchXPath('{%s}iq/{jabber:iq:roster}query' % self.default_ns), self._handleRoster, thread=True))
+		self.digest_auth_started = False
+		XMLStream.registerHandler(self, Callback('Stream Features', MatchXPath('{http://etherx.jabber.org/streams}features'), self._handleStreamFeatures, thread=True))
+		XMLStream.registerHandler(self, Callback('Roster Update', MatchXPath('{%s}iq/{jabber:iq:roster}query' % self.default_ns), self._handleRoster, thread=True))
+		#SASL Auth handlers
+		basexmpp.add_handler(self, "<challenge xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_sasl_digest_md5_auth, instream=True)
+		basexmpp.add_handler(self, "<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'/>", self.handler_sasl_digest_md5_auth_fail, instream=True)
+		basexmpp.add_handler(self, "<success xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_auth_success, instream=True)
+		basexmpp.add_handler(self, "<failure xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_auth_fail, instream=True)
 		#self.registerHandler(Callback('Roster Update', MatchXMLMask("<presence xmlns='%s' type='subscribe' />" % self.default_ns), self._handlePresenceSubscribe, thread=True))
 		self.registerFeature("<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls' />", self.handler_starttls, True)
 		self.registerFeature("<mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_sasl_auth, True)
@@ -107,8 +117,7 @@ class ClientXMPP(basexmpp, XMLStream):
 			else:
 				logging.debug("Since no address is supplied, attempting SRV lookup.")
 				try:
-					answers = dns.resolver.query("_xmpp-client._tcp.%s" % self.domain, 
-					        dns.rdatatype.SRV )
+					answers = dns.resolver.query("_xmpp-client._tcp.%s" % self.server, dns.rdatatype.SRV)
 				except dns.resolver.NXDOMAIN:
 					logging.debug("No appropriate SRV record found.  Using JID server name.")
 				else:
@@ -192,7 +201,7 @@ class ClientXMPP(basexmpp, XMLStream):
 			_stanza = "<proceed xmlns='urn:ietf:params:xml:ns:xmpp-tls' />"
 			if not self.event_handlers.get(_stanza,None): # don't add handler > once
 				self.add_handler( _stanza, self.handler_tls_start, instream=True )
-			self.sendXML(xml)
+			self.sendPriorityRaw(self.tostring(xml))
 			return True
 		else:
 			logging.warning("The module tlslite is required in to some servers, and has not been found.")
@@ -207,17 +216,17 @@ class ClientXMPP(basexmpp, XMLStream):
 		if '{urn:ietf:params:xml:ns:xmpp-tls}starttls' in self.features:
 			return False
 		logging.debug("Starting SASL Auth")
-		self.add_handler("<success xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_auth_success, instream=True)
-		self.add_handler("<failure xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />", self.handler_auth_fail, instream=True)
 		sasl_mechs = xml.findall('{urn:ietf:params:xml:ns:xmpp-sasl}mechanism')
 		if len(sasl_mechs):
 			for sasl_mech in sasl_mechs:
 				self.features.append("sasl:%s" % sasl_mech.text)
-			if 'sasl:PLAIN' in self.features:
+			if 'sasl:DIGEST-MD5' in self.features:
+				self.sendPriorityRaw("""<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' mechanism='DIGEST-MD5'/>""")
+			elif 'sasl:PLAIN' in self.features:
 				if sys.version_info < (3,0):
-					self.send("""<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' mechanism='PLAIN'>%s</auth>""" % base64.b64encode(b'\x00' + bytes(self.username) + b'\x00' + bytes(self.password)).decode('utf-8'))
+					self.sendPriorityRaw("""<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' mechanism='PLAIN'>%s</auth>""" % base64.b64encode(b'\x00' + bytes(self.username) + b'\x00' + bytes(self.password)).decode('utf-8'))
 				else:
-					self.send("""<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' mechanism='PLAIN'>%s</auth>""" % base64.b64encode(b'\x00' + bytes(self.username, 'utf-8') + b'\x00' + bytes(self.password, 'utf-8')).decode('utf-8'))
+					self.sendPriorityRaw("""<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' mechanism='PLAIN'>%s</auth>""" % base64.b64encode(b'\x00' + bytes(self.username, 'utf-8') + b'\x00' + bytes(self.password, 'utf-8')).decode('utf-8'))
 			else:
 				logging.error("No appropriate login method.")
 				self.disconnect()
@@ -225,6 +234,40 @@ class ClientXMPP(basexmpp, XMLStream):
 				#	self._auth_digestmd5()
 		return True
 	
+	def handler_sasl_digest_md5_auth(self, xml):
+		logging.debug(tostring(xml))
+		logging.debug(xml)
+		logging.debug(type(xml).__name__)
+		
+		if self.digest_auth_started == False:
+			challenge = [item.split('=', 1) for item in base64.b64decode(xml.text).replace("\"", "").split(',', 6) ]
+			challenge = dict(challenge)
+			logging.debug(challenge)
+			
+			#Realm, nonce, qop should all be present
+			if not challenge['realm'] or not challenge['qop'] or not challenge['nonce']:
+				logging.error("Error during digest-md5 authentication. Challenge missing critical information. Challenge: %s" %base64.b64decode(xml.text))
+				self.disconnect()
+				self.event("failed_auth")
+				return
+			#TODO: charset can be either UTF-8 or if not present use ISO 8859-1 defaulting for UTF-8 for now
+			#Compute the cnonce - a unique hex string only used in this request
+			cnonce = ""
+			for i in range(7):
+				cnonce+=hex(int(random.random()*65536*4096))[2:]
+			cnonce = base64.encodestring(cnonce)[0:-1]
+			a1 = b"%s:%s:%s" %(md5("%s:%s:%s" % (self.username, self.domain, self.password)), challenge["nonce"].encode("UTF-8"), cnonce.encode("UTF-8") )
+			a2 = "AUTHENTICATE:xmpp/%s" %self.domain
+			responseHash = md5digest("%s:%s:00000001:%s:auth:%s" %(md5digest(a1), challenge["nonce"], cnonce, md5digest(a2) ) )
+			response = '''charset=utf-8,username="%s",realm="%s",nonce="%s",nc=00000001,cnonce="%s",digest-uri="%s",response=%s,qop=%s,'''  %(self.username, self.domain, challenge["nonce"], cnonce, "xmpp/%s" % self.domain, responseHash, challenge["qop"])
+			self.sendPriorityRaw("""<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>%s</response>""" %base64.encodestring(response)[:-1])
+		else:
+			logging.warn("handler_sasl_digest_md5_auth called while digest_auth_started is True (has already begun)")
+	
+	def handler_sasl_digest_md5_auth_fail(self, xml):
+		self.digest_auth_started = False
+		self.handler_auth_fail(xml)
+	
 	def handler_auth_success(self, xml):
 		logging.debug("Authentication successful.")
 		self.authenticated = True
@@ -233,6 +276,7 @@ class ClientXMPP(basexmpp, XMLStream):
 
 	def handler_auth_fail(self, xml):
 		logging.warning("Authentication failed.")
+		logging.debug(tostring(xml, 'utf-8'))
 		self.disconnect()
 		self.event("failed_auth")
 	
@@ -273,3 +317,21 @@ class ClientXMPP(basexmpp, XMLStream):
 			if iq['type'] == 'set':
 				self.send(self.Iq().setValues({'type': 'result', 'id': iq['id']}).enable('roster'))
 		self.event("roster_update", iq)
+
+def md5(data):
+	try:
+		import hashlib
+		md5 = hashlib.md5(data)
+	except ImportError:
+		import md5
+		md5 = md5.new(data)
+	return md5.digest()
+
+def md5digest(data):
+	try:
+		import hashlib
+		md5 = hashlib.md5(data)
+	except ImportError:
+		import md5
+		md5 = md5.new(data)
+	return md5.hexdigest()
diff --git a/sleekxmpp/basexmpp.py b/sleekxmpp/basexmpp.py
index ad737360..25eeedb0 100644
--- a/sleekxmpp/basexmpp.py
+++ b/sleekxmpp/basexmpp.py
@@ -111,6 +111,7 @@ class basexmpp(object):
 			logging.debug("Loaded Plugin %s%s" % (xep, self.plugin[plugin].description))
 		except:
 			logging.exception("Unable to load plugin: %s", plugin )
+
             
 	def register_plugins(self):
 		"""Initiates all plugins in the plugins/__init__.__all__"""
diff --git a/sleekxmpp/stanza/error.py b/sleekxmpp/stanza/error.py
index ee46722a..3346ceb2 100644
--- a/sleekxmpp/stanza/error.py
+++ b/sleekxmpp/stanza/error.py
@@ -1,9 +1,9 @@
 """
-    SleekXMPP: The Sleek XMPP Library
-    Copyright (C) 2010  Nathanael C. Fritz
-    This file is part of SleekXMPP.
+	SleekXMPP: The Sleek XMPP Library
+	Copyright (C) 2010  Nathanael C. Fritz
+	This file is part of SleekXMPP.
 
-    See the file license.txt for copying permission.
+	See the file license.txt for copying permission.
 """
 from .. xmlstream.stanzabase import ElementBase, ET
 
diff --git a/sleekxmpp/stanza/iq.py b/sleekxmpp/stanza/iq.py
index ded7515f..26f09268 100644
--- a/sleekxmpp/stanza/iq.py
+++ b/sleekxmpp/stanza/iq.py
@@ -1,9 +1,9 @@
 """
-    SleekXMPP: The Sleek XMPP Library
-    Copyright (C) 2010  Nathanael C. Fritz
-    This file is part of SleekXMPP.
+	SleekXMPP: The Sleek XMPP Library
+	Copyright (C) 2010  Nathanael C. Fritz
+	This file is part of SleekXMPP.
 
-    See the file license.txt for copying permission.
+	See the file license.txt for copying permission.
 """
 from .. xmlstream.stanzabase import StanzaBase
 from xml.etree import cElementTree as ET
@@ -67,11 +67,11 @@ class Iq(RootStanza):
 				self.xml.remove(child)
 		return self
 	
-	def send(self, block=True, timeout=10):
+	def send(self, block=True, timeout=10, priority=False):
 		if block and self['type'] in ('get', 'set'):
 			waitfor = Waiter('IqWait_%s' % self['id'], MatcherId(self['id']))
 			self.stream.registerHandler(waitfor)
-			StanzaBase.send(self)
+			StanzaBase.send(self, priority)
 			return waitfor.wait(timeout)
 		else:
-			return StanzaBase.send(self)
+			return StanzaBase.send(self, priority)
diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 64020c8f..34513807 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -1,9 +1,9 @@
 """
-    SleekXMPP: The Sleek XMPP Library
-    Copyright (C) 2010  Nathanael C. Fritz
-    This file is part of SleekXMPP.
+	SleekXMPP: The Sleek XMPP Library
+	Copyright (C) 2010  Nathanael C. Fritz
+	This file is part of SleekXMPP.
 
-    See the file license.txt for copying permission.
+	See the file license.txt for copying permission.
 """
 from xml.etree import cElementTree as ET
 import logging
@@ -383,6 +383,7 @@ class StanzaBase(ElementBase):
 	def exception(self, e):
 		logging.error(traceback.format_tb(e))
 	
-	def send(self):
-		self.stream.sendRaw(self.__str__())
-
+	def send(self, priority=False):
+		if priority: self.stream.sendPriorityRaw(self.__str__()) 
+		else: self.stream.sendRaw(self.__str__())
+		
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index d5238f80..842dfee2 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -81,7 +81,7 @@ class XMLStream(object):
 		self.stream_footer = "</stream>"
 
 		self.eventqueue = queue.Queue()
-		self.sendqueue = queue.Queue()
+		self.sendqueue = queue.PriorityQueue()
 		self.scheduler = scheduler.Scheduler(self.eventqueue)
 
 		self.namespace_map = {}
@@ -220,7 +220,7 @@ class XMLStream(object):
 		while self.run:
 			if not self.state.ensure('connected',wait=2): continue
 			try:
-				self.sendRaw(self.stream_header)
+				self.sendPriorityRaw(self.stream_header)
 				while self.run and self.__readXML(): pass
 			except socket.timeout:
 				logging.debug('socket rcv timeout')
@@ -281,7 +281,7 @@ class XMLStream(object):
 			
 			data = None
 			try:
-				data = self.sendqueue.get(True,5)
+				data = self.sendqueue.get(True,5)[1]
 				logging.debug("SEND: %s" % data)
 				self.socket.sendall(data.encode('utf-8'))
 			except queue.Empty:
@@ -302,7 +302,11 @@ class XMLStream(object):
 					self.disconnect(reconnect=True)
 	
 	def sendRaw(self, data):
-		self.sendqueue.put(data)
+		self.sendqueue.put((1, data))
+		return True
+	
+	def sendPriorityRaw(self, data):
+		self.sendqueue.put((0, data))
 		return True
 	
 	def disconnect(self, reconnect=False):