From aeb43dff55175c512ca541f92bc151d09d18d610 Mon Sep 17 00:00:00 2001 From: moyamo Date: Thu, 7 Jan 2016 19:31:05 +0200 Subject: [PATCH] Add ability to request password after receiving sms code --- Spectrum2/backend.py | 6 +++++ registersession.py | 62 +++++++++++++++++++++++++++++++++++++++----- whatsappbackend.py | 16 +++++++++++- yowsupwrapper.py | 18 ++++++++++++- 4 files changed, 94 insertions(+), 8 deletions(-) diff --git a/Spectrum2/backend.py b/Spectrum2/backend.py index 5a91a27..b2e5891 100644 --- a/Spectrum2/backend.py +++ b/Spectrum2/backend.py @@ -239,6 +239,12 @@ class SpectrumBackend: message = WRAP(c.SerializeToString(), protocol_pb2.WrapperMessage.TYPE_BACKEND_CONFIG); self.send(message) + def handleQuery(self, command): + c = protocol_pb2.BackendConfig() + c.config = command + message = WRAP(c.SerializeToString(), protocol_pb2.WrapperMessage.TYPE_QUERY); + self.send(message) + def handleLoginPayload(self, data): payload = protocol_pb2.Login() if (payload.ParseFromString(data) == False): diff --git a/registersession.py b/registersession.py index b07b675..963eeb2 100644 --- a/registersession.py +++ b/registersession.py @@ -3,6 +3,7 @@ from Spectrum2 import protocol_pb2 from yowsupwrapper import YowsupApp import logging import threadutils +import sys class RegisterSession(YowsupApp): """ @@ -14,6 +15,7 @@ class RegisterSession(YowsupApp): self.user = user self.number = legacyName self.backend = backend + self.countryCode = '' self.logger = logging.getLogger(self.__class__.__name__) self.state = self.WANT_CC @@ -39,25 +41,73 @@ class RegisterSession(YowsupApp): else: self.backend.handleMessage(self.user, 'bot', 'Requesting sms code') self.logger.debug('Requesting SMS code for %s', self.user) - self._requestSMSCodeNonBlock(country_code) + self.countryCode = country_code + self._requestSMSCodeNonBlock() elif buddy == 'bot' and self.state == self.WANT_SMS: - self.backend.handleMessage(self.user, 'bot', 'Not implemented') + code = message.strip() + if self._checkSMSFormat(code): + self._requestPassword(code) + else: + self.backend.handleMessage(self.user, + 'bot', 'Invalid code. Must be of the form XXX-XXX.') else: self.logger.warn('Unauthorised user (%s) attempting to send messages', self.user) self.backend.handleMessage(self.user, buddy, 'You are not logged in yet. You can only send messages to bot.') - def _requestSMSCodeNonBlock(self, country_code): - number = self.number[len(country_code):] - threadFunc = lambda: self.requestSMSCode(country_code, number) + def _checkSMSFormat(self, sms): + splitting = sms.split('-') + if len(splitting) != 2: + return False + a, b = splitting + if len(a) != 3 and len(b) != 3: + return False + try: + int(a) + int(b) + except ValueError: + return False + return True + + def _requestSMSCodeNonBlock(self): + number = self.number[len(self.countryCode):] + threadFunc = lambda: self.requestSMSCode(self.countryCode, number) threadutils.runInThread(threadFunc, self._confirmation) + self.backend.handleMessage(self.user, 'bot', 'SMS Code Sent') def _confirmation(self, result): - self.backend.handleMessage(self.user, 'bot', 'SMS Code Sent') self.state = self.WANT_SMS + resultStr = self._resultToString(result) + self.backend.handleMessage(self.user, 'bot', 'Response:') + self.backend.handleMessage(self.user, 'bot', resultStr) self.backend.handleMessage(self.user, 'bot', 'Please enter SMS Code') + def _requestPassword(self, smsCode): + cc = self.countryCode + number = self.number[len(cc):] + threadFunc = lambda: self.requestPassword(cc, number, smsCode) + threadutils.runInThread(threadFunc, self._gotPassword) + self.backend.handleMessage(self.user, 'bot', 'Getting Password') + + def _gotPassword(self, result): + resultStr = self._resultToString(result) + self.backend.handleMessage(self.user, 'bot', 'Response:') + self.backend.handleMessage(self.user, 'bot', resultStr) + self.backend.handleMessage(self.user, 'bot', 'Logging you in') + password = result['pw'] + self.backend.relogin(self.user, self.number, password, None) + + def _resultToString(self, result): + unistr = str if sys.version_info >= (3, 0) else unicode + out = [] + for k, v in result.items(): + if v is None: + continue + out.append("%s: %s" %(k, v.encode("utf-8") if type(v) is unistr else v)) + + return "\n".join(out) + # Dummy methods. Whatsapp backend might call these, but they should have no # effect def logout(self): diff --git a/whatsappbackend.py b/whatsappbackend.py index 068f1a8..46b2dc5 100644 --- a/whatsappbackend.py +++ b/whatsappbackend.py @@ -44,7 +44,7 @@ class WhatsAppBackend(SpectrumBackend): # RequestsHandlers def handleLoginRequest(self, user, legacyName, password, extra): self.logger.debug("handleLoginRequest(user=%s, legacyName=%s)", user, legacyName) - # Key word means we should register + # Key word means we should register a new password if password == 'register': if user not in self.sessions: self.sessions[user] = RegisterSession(self, user, legacyName, extra) @@ -117,6 +117,20 @@ class WhatsAppBackend(SpectrumBackend): self.logger.debug("handleVCardRequest(user=%s, buddy=%s, ID=%s)", user, buddy, ID) self.sessions[user].requestVCard(buddy, ID) + def relogin(self, user, legacyName, password, extra): + """ + Used to re-initialize the session object. Used when finished with + registration session and the user needs to login properly + """ + self.logger.debug("relogin(user=%s, legacyName=%s)", user, legacyName) + # Change password in spectrum database + self.handleQuery('register %s %s %s' % (user, legacyName, password)) + # Key word means we should register a new password + if password == 'register': # This shouldn't happen, but just in case + self.sessions[user] = RegisterSession(self, user, legacyName, extra) + else: + self.sessions[user] = Session(self, user, legacyName, extra) + self.sessions[user].login(password) # TODO def handleBuddyBlockToggled(self, user, buddy, blocked): diff --git a/yowsupwrapper.py b/yowsupwrapper.py index f6f0116..23b84de 100644 --- a/yowsupwrapper.py +++ b/yowsupwrapper.py @@ -45,7 +45,8 @@ from yowsup.layers.protocol_media.mediauploader import MediaUploader # Registration -from yowsup.registration.coderequest import WACodeRequest +from yowsup.registration import WACodeRequest +from yowsup.registration import WARegRequest from functools import partial @@ -388,6 +389,21 @@ class YowsupApp(object): request = WACodeRequest(countryCode, phoneNumber) return request.send() + def requestPassword(self, countryCode, phoneNumber, smsCode): + """ + Request a password. WARNING: this function is blocking + + Args: + countryCode: The country code of the phone you wish to register + phoneNumber: phoneNumber of the phone you wish to register without + the country code. + smsCode: The sms code that you asked for previously + """ + smsCode = smsCode.replace('-', '') + request = WARegRequest(countryCode, phoneNumber, smsCode) + return request.send() + + def onAuthSuccess(self, status, kind, creation, expiration, props, nonce, t): """