From 39ad54f3d90e35757eccd003fc81d68f91ff7482 Mon Sep 17 00:00:00 2001 From: Edward <73746306+WangEdward@users.noreply.github.com> Date: Wed, 14 Feb 2024 05:30:41 +0000 Subject: [PATCH] =?UTF-8?q?feat=20=E6=96=B0=E5=A2=9E=202fa=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/helper/twofa.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 app/helper/twofa.py diff --git a/app/helper/twofa.py b/app/helper/twofa.py new file mode 100644 index 00000000..4c24dab1 --- /dev/null +++ b/app/helper/twofa.py @@ -0,0 +1,39 @@ +import base64 +import hashlib +import hmac +import struct +import sys +import time +from app.log import logger + + +class TwoFactorAuth: + def __init__(self, code_or_secret: str): + if len(code_or_secret) < 16: + self.code = code_or_secret + self.secret = None + else: + self.code = None + self.secret = code_or_secret + + def __calc(self, secret_key: str) -> str: + try: + input_time = int(time.time()) // 30 + key = base64.b32decode(secret_key) + msg = struct.pack(">Q", input_time) + google_code = hmac.new(key, msg, hashlib.sha1).digest() + o = ( + google_code[19] & 15 + if sys.version_info > (2, 7) + else ord(str(google_code[19])) & 15 + ) + google_code = str( + (struct.unpack(">I", google_code[o : o + 4])[0] & 0x7FFFFFFF) % 1000000 + ) + return f"0{google_code}" if len(google_code) == 5 else google_code + except Exception as e: + logger.error(f"计算动态验证码失败:{str(e)}") + return None + + def get_code(self) -> str: + return self.code or self.__calc(self.secret)