57 lines
1.4 KiB
Python
57 lines
1.4 KiB
Python
|
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()
|