Fix wrong admin jids detection

Fix longpoll: now user won't fall out the Poll
Fix several captcha exceptions
Split get&set in mod_iq_register to a different functions
Fix several wrong argument exceptions in mod_prs_main
No code climate changes were applied.

4 files changed, 86 insertions, 92 deletions

diff --git a/gateway.py b/gateway.py
index 2e22df9..4d48bbf 100644
--- a/gateway.py
+++ b/gateway.py
@@ -83,8 +83,8 @@ import vkapi as api
 import utils
 
 # Compatibility with old config files
-if not ADMIN_JIDS:
-	ADMIN_JIDS = [evalJID]
+if not ADMIN_JIDS and evalJID:
+	ADMIN_JIDS.append(evalJID)
 
 # Setting variables
 # DefLang for language id, root for the translations directory
@@ -193,10 +193,10 @@ badChars = [x for x in xrange(32) if x not in (9, 10, 13)] + [57003, 65535]
 escape = re.compile("|".join(unichr(x) for x in badChars),
 	re.IGNORECASE | re.UNICODE | re.DOTALL).sub
 
+checkOnline = lambda friends, key: friends[key]["online"]
 sortMsg = lambda first, second: first.get("mid", 0) - second.get("mid", 0)
 require = lambda name: os.path.exists("extensions/%s.py" % name)
 isdef = lambda var: var in globals()
-findListByID = lambda id, list: [key for key in list if key["lid"] == id]
 
 
 class VK(object):
@@ -223,7 +223,7 @@ class VK(object):
 		Checks the api token
 		"""
 		try:
-			int(self.method("isAppUser", force=True))
+			int(self.engine.method("isAppUser"))
 		except (api.VkApiError, TypeError, AttributeError):
 			return False
 		return True
@@ -279,7 +279,7 @@ class VK(object):
 			nodecode: decode flag (make json.loads or not)
 			force: says that method will be executed even the captcha and not online
 		See library/vkapi.py for more information about exceptions
-		Returns method execition result
+		Returns method execution result
 		"""
 		args = args or {}
 		result = {}
@@ -291,7 +291,11 @@ class VK(object):
 				if force:
 					raise
 
-			except api.NetworkNotFound as e:
+			except api.CaptchaNeeded:
+				executeHandlers("evt04", (self, self.engine.captcha["img"]))
+				self.online = False
+
+			except api.NetworkNotFound:
 				self.online = False
 
 			except api.NotAllowed as e:
@@ -495,8 +499,7 @@ class User(object):
 		runDatabaseQuery("update users set rosterSet=? where jid=?",
 			(self.rosterSet, self.source), True)
 
-
-	def initialize(self, force=False, send=True, resource=None, raise_exc=False):
+	def initialize(self, force=False, send=True, resource=None):
 		"""
 		Initializes user after self.connect() is done:
 			1. Receives friends list and set 'em to self.friends
@@ -529,9 +532,9 @@ class User(object):
 
 	def sendInitPresence(self):
 		"""
-		Sends init presence (available) to user from all his online friends
+		Sends init presence (available) to the user from all their online friends
 		"""
-		if not self.settings.i_am_ghost:
+		if not self.settings.i_am_ghost and not self.vk.engine.captcha:
 			if not self.friends:
 				self.friends = self.vk.getFriends()
 			logger.debug("User: sending init presence (friends count: %s) (jid %s)",
@@ -543,12 +546,11 @@ class User(object):
 				if value["online"]:
 					sendPresence(self.source, vk2xmpp(uid), None,
 						value.get(key, "Unknown"), caps=True)
-		if not self.vk.engine.captcha:
 			sendPresence(self.source, TransportID, None, IDENTIFIER["name"], caps=True)
 
 	def sendOutPresence(self, destination, reason=None, all=False):
 		"""
-		Sends out presence (unavailable) to destination and set reason if exists
+		Sends out presence (unavailable) to destination. Defines a reason, if set.
 		Parameters:
 			destination: to whom send the stanzas
 			reason: offline status message
@@ -557,7 +559,7 @@ class User(object):
 		logger.debug("User: sending out presence to %s", self.source)
 		friends = self.friends.keys()
 		if not all and friends:
-			friends = filter(lambda key: self.friends[key]["online"], friends)
+			friends = filter(checkOnline, friends)
 
 		for uid in friends + [TransportID]:
 			sendPresence(destination, vk2xmpp(uid), "unavailable", reason=reason)
@@ -649,7 +651,7 @@ class User(object):
 			return 1
 
 		if "failed" in data:
-			logger.debug("longpoll: failed. Searching for new server(jid: %s)",
+			logger.debug("longpoll: failed. Searching for a new server (jid: %s)",
 				self.source)
 			return 0
 
diff --git a/library/longpoll.py b/library/longpoll.py
index b89ad50..377ab27 100644
--- a/library/longpoll.py
+++ b/library/longpoll.py
@@ -168,9 +168,9 @@ class Poll:
 		"""
 		Processes the select result (see above)
 		Handles answers from user.processPollResult()
-		Decides neiter if need to add user to poll or not
+		Decides if need to add user to poll or not
 		"""
-		result = utils.execute(user.processPollResult, (opener,)) or -1
+		result = utils.execute(user.processPollResult, (opener,))
 		if DEBUG_POLL:
 			logger.debug("longpoll: result=%d (jid: %s)" % (result, user.source))
 		if result == -1:
diff --git a/modules/mod_iq_register.py b/modules/mod_iq_register.py
index 7d76b55..d3acc85 100644
--- a/modules/mod_iq_register.py
+++ b/modules/mod_iq_register.py
@@ -4,6 +4,7 @@
 
 from __main__ import *
 from __main__ import _
+import forms
 
 
 @utils.threaded
@@ -18,30 +19,19 @@ def initializeUser(user, cl, iq, kwargs):
 		result = utils.buildIQError(iq, xmpp.ERR_NOT_AUTHORIZED, _(str(e) + " Try logging in by token."))
 	else:
 		if connect:
-			try:
-				user.initialize(resource=resource)
-			except api.CaptchaNeeded:
-				user.vk.captchaChallenge()
-			except Exception:
-				crashLog("user.init")
-				result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("Initialization failed."))
-			else:
-				executeHandlers("evt08", (source,))
+			user.initialize(resource=resource)
+			executeHandlers("evt08", (source,))
 		else:
 			logger.error("user connection failed (jid: %s)" % source)
 			result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("Incorrect password or access token!"))
 	sender(cl, result)
 
 
-import forms
-
+@utils.safe
 def register_handler(cl, iq):
-	jidTo = iq.getTo()
-	jidFrom = iq.getFrom()
-	source = jidFrom.getStripped()
-	destination = jidTo.getStripped()
+	source = iq.getFrom().getStripped()
+	destination = iq.getTo().getStripped()
 	iType = iq.getType()
-	queryChildren = iq.getQueryChildren()
 	result = iq.buildReply("result")
 	if USER_LIMIT:
 		count = calcStats()[0]
@@ -49,66 +39,63 @@ def register_handler(cl, iq):
 			sender(cl, utils.buildIQError(iq, xmpp.ERR_NOT_ALLOWED, _("Transport's admins limited registrations, sorry.")))
 			raise xmpp.NodeProcessed()
 
-	if destination == TransportID:
-		if iType == "get" and not queryChildren:
-			logger.debug("Send registration form to user (jid: %s)", source)
-			form = utils.buildDataForm(fields=forms.Forms.getComlicatedForm(), data=[_("Fill the fields below")])
-			result.setQueryPayload([form])
+	if destination == TransportID and iq.getQueryChildren():
+		phone, password, use_password, token, result = None, None, None, None, None
+		query = iq.getTag("query")
+		data = query.getTag("x", namespace=xmpp.NS_DATA)
+		if data:
+			form = xmpp.DataForm(node=data).asDict()
+			phone = str(form.get("phone", "")).lstrip("+")
+			password = str(form.get("password", ""))
+			use_password = utils.normalizeValue(form.get("use_password", ""))  # In case here comes some unknown crap
 
-		elif iType == "set" and queryChildren:
-			phone, password, use_password, token, result = False, False, False, False, False # Why result is here?
-			query = iq.getTag("query")
-			data = query.getTag("x", namespace=xmpp.NS_DATA)
-			if data:
-				form = xmpp.DataForm(node=data).asDict()
-				phone = str(form.get("phone", "")).lstrip("+")
-				password = str(form.get("password", ""))
-				use_password = utils.normalizeValue(form.get("use_password", ""))  # In case here comes some unknown crap
-
-				if not password:
-					result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("The token/password field can't be empty!"))
+			if not password:
+				result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("The token/password field can't be empty!"))
+			else:
+				if use_password:
+					logger.debug("user want to use a password (jid: %s)" % source)
+					if not phone or phone == "+":
+						result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("Phone is incorrect."))
 				else:
-					if use_password:
-						logger.debug("user want to use a password (jid: %s)" % source)
-						if not phone or phone == "+":
-							result = utils.buildIQError(iq, xmpp.ERR_BAD_REQUEST, _("Phone is incorrect."))
+					logger.debug("user won't use a password (jid: %s)" % source)
+					token = password
+					password = None
+					# If not using a password, then we need to check if there a link or token. It's possible that user's wrong and that's a password.
+					match = api.token_exp.search(token)
+					if match:
+						token = match.group(0)
+					elif phone:
+						password = token
 					else:
-						logger.debug("user won't use a password (jid: %s)" % source)
-						token = password
-						password = False
-				
-						## If not using a password, then we need to check if there a link or token. It's possible that user's wrong and that's a password.
-						_token = api.token_exp.search(token)
-						if _token:
-							token = _token.group(0)
-						elif phone:
-							password = token
-						else:
-							result = utils.buildIQError(iq, xmpp.ERR_NOT_AUTHORIZED, _("Fill the fields!"))
-
-					# If phone or password (token)
-					if token or (phone and password):
-						user = User(source)
-						initializeUser(user, cl, iq, {"username": phone, "password": password, "token": token})
-						result = None
+						result = utils.buildIQError(iq, xmpp.ERR_NOT_AUTHORIZED, _("Fill the fields!"))
 
-			elif query.getTag("remove"):
-				logger.debug("user %s want to remove me..." % source)
-				if source in Transport:
-					user = Transport[source]
-					removeUser(user, True, False)
-					result = iq.buildReply("result") # Is it required?
-					result.setPayload([], add=False)
-					executeHandlers("evt09", (source,))
-				else:
-					logger.debug("... but they don't know that they were removed already!")
+				# If phone or password (token)
+				if token or (phone and password):
+					user = User(source)
+					initializeUser(user, cl, iq, {"username": phone, "password": password, "token": token})
+					result = None
 
-		else:
-			result = utils.buildIQError(iq, 0, _("Feature not implemented."))
+		elif query.getTag("remove"):
+			logger.debug("user %s want to remove me..." % source)
+			if source in Transport:
+				user = Transport[source]
+				removeUser(user, True, False)
+				result = iq.buildReply("result")
+				result.setPayload([], add=False)
+				executeHandlers("evt09", (source,))
 	if result:
 		sender(cl, result)
 
 
+
+def sendRegisterForm(cl, iq):
+	logger.debug("Send registration form to user (jid: %s)", iq.getFrom().getStripped())
+	form = utils.buildDataForm(fields=forms.Forms.getComlicatedForm(), data=[_("Fill the fields below")])
+	result = iq.buildReply("result")
+	result.setQueryPayload([form])
+	sender(cl, result)
+
+
 MOD_TYPE = "iq"
 MOD_FEATURES = [xmpp.NS_DATA, xmpp.NS_REGISTER]
-MOD_HANDLERS = ((register_handler, "", xmpp.NS_REGISTER, False),)
+MOD_HANDLERS = ((register_handler, "set", xmpp.NS_REGISTER, False), (sendRegisterForm, "get", xmpp.NS_REGISTER, False))
diff --git a/modules/mod_prs_main.py b/modules/mod_prs_main.py
index a18ec42..8c6ac46 100644
--- a/modules/mod_prs_main.py
+++ b/modules/mod_prs_main.py
@@ -24,11 +24,16 @@ def initializeUser(source, resource, prs):
 		logger.debug("User has been found in database (jid: %s)" % source)
 		jid, phone = data
 		Transport[jid] = user = User(jid)
-		if user.connect():
+		try:
+			connect = user.connect()
+		except Exception:
+			sendMessage(Component, jid, TransportID, 
+				_("Auth failed! If this error repeated, "
+					"please register again. This incident will be reported."))
+			crashLog("user.connect")
+		else:
 			user.initialize(send=True, resource=resource)  # probably we need to know resource a bit earlier than this time
 			utils.runThread(executeHandlers, ("prs01", (source, prs)))
-		else:
-			sendMessage(Component, jid, TransportID, _("Auth failed! If this error repeated, please register again. This incident will be reported."))
 
 	if source in USERS_ON_INIT:
 		USERS_ON_INIT.remove(source)
@@ -54,11 +59,11 @@ def presence_handler(cl, prs):
 			if jidTo == TransportID and resource in user.resources:
 				user.resources.remove(resource)
 				if user.resources:
-					user.sendOutPresence(jidFrom)
+					user.sendOutPresence(source)
 			if not user.resources:
-				sendPresence(jidFrom, TransportID, "unavailable")
+				sendPresence(source, TransportID, "unavailable")
 				if transportSettings.send_unavailable:
-					user.sendOutPresence(jidFrom)
+					user.sendOutPresence(source)
 				user.vk.disconnect()
 				try:
 					del Transport[source]
@@ -75,9 +80,9 @@ def presence_handler(cl, prs):
 				id = vk2xmpp(destination)
 				if id in user.friends:
 					if user.friends[id]["online"]:
-						sendPresence(jidFrom, jidTo)
+						sendPresence(source, destination)
 			if destination == TransportID:
-				sendPresence(jidFrom, TransportID)
+				sendPresence(source, destination)
 	
 		elif pType == "unsubscribe":
 			if destination == TransportID: