From 55691ff2bf0d47fd5052b6c8e01ce496988aa730 Mon Sep 17 00:00:00 2001 From: knotteye Date: Sat, 17 Apr 2021 23:34:34 -0500 Subject: [PATCH] Add totp support --- plchat.py | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++---- pleroma.py | 13 +++++++++- 2 files changed, 82 insertions(+), 6 deletions(-) diff --git a/plchat.py b/plchat.py index 27f4930..87c6603 100644 --- a/plchat.py +++ b/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) diff --git a/pleroma.py b/pleroma.py index a779688..455d3c7 100644 --- a/pleroma.py +++ b/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')