From 512057ab61c8811df19f646711ce985e1872254b Mon Sep 17 00:00:00 2001 From: lucasheld Date: Fri, 5 Aug 2022 14:33:28 +0200 Subject: [PATCH] implement 2FA login --- tests/requirements.txt | 1 + tests/test_2fa.py | 56 ++++++++++++++++++++++++++++++++++++++++++ uptime_kuma_api/api.py | 10 ++++---- 3 files changed, 62 insertions(+), 5 deletions(-) create mode 100644 tests/requirements.txt create mode 100644 tests/test_2fa.py diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000..15e97f7 --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1 @@ +pyotp==2.6.0 diff --git a/tests/test_2fa.py b/tests/test_2fa.py new file mode 100644 index 0000000..d9d64e0 --- /dev/null +++ b/tests/test_2fa.py @@ -0,0 +1,56 @@ +import unittest +from urllib import parse + +import pyotp + +from uptime_kuma_test_case import UptimeKumaTestCase + + +def parse_secret(uri): + query = parse.urlsplit(uri).query + params = dict(parse.parse_qsl(query)) + return params["secret"] + + +def generate_token(secret): + totp = pyotp.TOTP(secret) + return totp.now() + + +class Test2FA(UptimeKumaTestCase): + def test_2fa(self): + # check 2fa is disabled + r = self.api.twofa_status() + self.assertEqual(r["status"], False) + + # prepare 2fa + r = self.api.prepare_2fa(self.password) + uri = r["uri"] + self.assertTrue(uri.startswith("otpauth://totp/")) + secret = parse_secret(uri) + + # verify token + token = generate_token(secret) + r = self.api.verify_token(token, self.password) + self.assertEqual(r["valid"], True) + + # save 2fa + r = self.api.save_2fa(self.password) + self.assertEqual(r["msg"], "2FA Enabled.") + + # check 2fa is enabled + r = self.api.twofa_status() + self.assertEqual(r["status"], True) + + # relogin using the totp token + self.api.logout() + token = generate_token(secret) + self.api.login(self.username, self.password, token) + + # disable 2fa + r = self.api.disable_2fa(self.password) + self.assertEqual(r["msg"], "2FA Disabled.") + + +if __name__ == '__main__': + unittest.main() diff --git a/uptime_kuma_api/api.py b/uptime_kuma_api/api.py index f18cd46..160e5aa 100644 --- a/uptime_kuma_api/api.py +++ b/uptime_kuma_api/api.py @@ -801,6 +801,9 @@ class UptimeKumaApi(object): def prepare_2fa(self, password: str): return self._call('prepare2FA', password) + def verify_token(self, token: str, password: str): + return self._call('verifyToken', (token, password)) + def save_2fa(self, password: str): return self._call('save2FA', password) @@ -809,19 +812,16 @@ class UptimeKumaApi(object): # login - def login(self, username: str, password: str): + def login(self, username: str, password: str, twofa_token: str = ""): return self._call('login', { "username": username, "password": password, - "token": "" + "token": twofa_token }) def login_by_token(self, token: str): return self._call('loginByToken', token) - def verify_token(self, token: str, password: str): - return self._call('verifyToken', (token, password)) - def logout(self): return self._call('logout')