Add totp support

plchat.py

@@ -55,6 +55,7 @@ class App(QMainWindow):
         self.setGeometry(self.settings.value('left', type=int) or 10, self.settings.value('top', type=int) or 10, self.settings.value('width', type=int) or 640, self.settings.value('height', type=int) or 480)
 
         self._exit = False
+        self.totpReady = False
         self.Err = QErrorMessage()
 
         exitAction = QAction('&Exit', self)
@@ -277,6 +278,20 @@ class App(QMainWindow):
             self.settings.setValue('closed'+u+i, closedList)
             CallThread(self.accts[u+i].addChat, self.populateChats, result['id']).start()
 
+    def getTotp(self, acct, challenge_type):
+        if type(challenge_type) != list:
+            challenge_type = [challenge_type, 'recovery']
+        self._eventloop.call_soon_threadsafe(self.makeTotpCard, '@'+acct.username+'@'+acct.instance, challenge_type)
+        while not self.totpReady:
+            monkeypatch.sleep(0.2)
+        t = self.totpReady
+        self.totpReady = None
+        return t
+
+    def makeTotpCard(self, acct, challenge_type):
+        dialog = TotpCard(acct, challenge_type)
+        dialog.getInput()
+
     def contactDialog(self):
         dialog = ContactCard(self)
         dialog.show()
@@ -294,7 +309,7 @@ class App(QMainWindow):
 
     def initAcct(self, instance, username, password=None):
         if password:
-            acct = pleroma.Account(instance, username, password)
+            acct = pleroma.Account(instance, username, password, totpFunc=self.getTotp)
         else:
             token = keyring.get_password('plchat', instance+username+'access_token')
             refresh_token = keyring.get_password('plchat', instance+username+'refresh_token')
@@ -304,7 +319,8 @@ class App(QMainWindow):
                 token=token,
                 refresh_token=refresh_token,
                 clientID=clientID,
-                clientSecret=clientSecret
+                clientSecret=clientSecret,
+                totpFunc=self.getTotp
             )
         RegisterThread(acct, self.doneRegister).start()
 
@@ -352,7 +368,8 @@ class App(QMainWindow):
         else:
             self.setWindowTitle('PlChat')
         self.tabs.clear()
-        CallThread(self.accts[u+i].listChats, self.populateChats).start()
+        if u and i:
+            CallThread(self.accts[u+i].listChats, self.populateChats).start()
         
     def populateChats(self, chatList):
         if type(chatList) == dict:
@@ -641,9 +658,9 @@ class MessageArea(QWidget):
         CallThread(ex.accts[u+i].getMessages, self._update, self.chatID).start()
 
     def markRead(self):
-        if not self.last_read_id:
-            return
         u, i = ex.getCurrentAcc()
+        if not self.last_read_id or not u or not i:
+            return
         acc = ex.accts[u+i]
         acc.markChatRead(self.chatID, self.last_read_id)
 
@@ -1067,6 +1084,54 @@ class DetachDialog(QDialog):
         self.raise_()
         self.activateWindow()
 
+class TotpCard(QDialog):
+    def __init__(self, acct, challenge_types, parent=None,):
+        super().__init__(parent=parent)
+        self.setWindowTitle("MFA Request")
+        self.result = None
+        self.waiting = False
+
+        QBtn = QDialogButtonBox.Ok
+
+        layout = QVBoxLayout()
+
+        self.buttonBox = QDialogButtonBox(QBtn)
+        self.buttonBox.accepted.connect(self.accept)
+
+        self.message = QLabel("2FA Required for "+acct)
+        tmp = QHBoxLayout()
+        tmpw = QWidget()
+        tmpw.setLayout(tmp)
+        self.comboboxlabel = QLabel("2FA Code Type:")
+        self.combobox = QComboBox()
+
+        self.combobox.addItems(challenge_types)
+        self.combobox.setDuplicatesEnabled(False)
+        self.combobox.setEditable(False)
+
+        tmp.addWidget(self.comboboxlabel)
+        tmp.addWidget(self.combobox)
+        self.code = QLineEdit()
+        self.code.setPlaceholderText('2FA Code')
+
+        layout.addWidget(self.message)
+        layout.addWidget(tmpw)
+        layout.addWidget(self.code)
+        layout.addWidget(self.buttonBox)
+
+        self.setLayout(layout)
+        self.setFocusProxy(self.code)
+        self.accepted.connect(self._finished)
+
+    def _finished(self):
+        ex.totpReady = (self.code.text(), self.combobox.currentText())
+
+    def getInput(self, *args, **kwargs):
+        self.exec_()
+        self.show()
+        self.raise_()
+        self.activateWindow()
+
 class LoginDialog(QDialog):
     def __init__(self, parent=None):
         super().__init__(parent=parent)

pleroma.py

@@ -16,7 +16,7 @@
 import re, requests, os.path, websockets, json
 
 class Account():
-	def __init__(self, instance, username=' ', password='', clientID=None, clientSecret=None, token=None, refresh_token=None):
+	def __init__(self, instance, username=' ', password='', clientID=None, clientSecret=None, token=None, refresh_token=None, totpFunc=None):
 		scheme = re.compile('https?://')
 		self.instance = re.sub(scheme, '', instance)
 		if(self.instance[len(self.instance) - 1] == '/'):
@@ -39,6 +39,7 @@ class Account():
 		self._instanceInfo = None
 		self.flakeid = None
 		self.acct = None
+		self.totpFunc = totpFunc
 
 		self.chat_update = None
 
@@ -83,6 +84,16 @@ class Account():
 			'redirect_uris': "urn:ietf:wg:oauth:2.0:oob"
 		}
 		response = self.apiRequest('POST', '/oauth/token', request_data)
+		if 'error' in response and response['error'] == 'mfa_required':
+			if response['supported_challenge_types'] == 'totp' or 'totp' in response['supported_challenge_types']:
+				mfa_code, code_type = self.totpFunc(self, response['supported_challenge_types'])
+				response = self.apiRequest('POST', '/oauth/mfa/challenge', {
+					'client_id': self.clientID,
+					'client_secret': self.clientSecret,
+					'mfa_token': response['mfa_token'],
+					'challenge_type': code_type,
+					'code': mfa_code
+				})
 		self.token = response['access_token']
 		self.refresh_token = response['refresh_token']
 		r = self.apiRequest('GET', '/api/v1/accounts/verify_credentials')