mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-30 03:48:03 +08:00
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
"""
|
|
解析站点登录输入中的验证码或共享密钥，并按需生成 TOTP 验证码。
|
|
"""
|
|
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import struct
|
|
import sys
|
|
import time
|
|
|
|
from app.runtime.log import logger
|
|
|
|
|
|
class TwoFactorAuth:
|
|
"""
|
|
解析已有验证码或根据共享密钥生成 TOTP 验证码。
|
|
"""
|
|
|
|
def __init__(self, code_or_secret: str | None) -> None:
|
|
"""
|
|
按长度区分用户验证码和 TOTP 共享密钥。
|
|
"""
|
|
if code_or_secret and len(code_or_secret) >= 16:
|
|
self.code: str | None = None
|
|
self.secret: str | None = code_or_secret
|
|
else:
|
|
self.code = code_or_secret
|
|
self.secret = None
|
|
|
|
@staticmethod
|
|
def __calc(secret_key: str | None) -> str:
|
|
"""
|
|
按 30 秒时间窗计算六位 TOTP 验证码。
|
|
"""
|
|
if not secret_key:
|
|
return ""
|
|
try:
|
|
input_time = int(time.time()) // 30
|
|
key = base64.b32decode(secret_key)
|
|
msg = struct.pack(">Q", input_time)
|
|
digest = hmac.new(key, msg, hashlib.sha1).digest()
|
|
o = (
|
|
digest[19] & 15
|
|
if sys.version_info > (2, 7)
|
|
else ord(str(digest[19])) & 15
|
|
)
|
|
google_code = str(
|
|
(struct.unpack(">I", digest[o: o + 4])[0] & 0x7FFFFFFF) % 1000000
|
|
)
|
|
return google_code.zfill(6)
|
|
except Exception as e:
|
|
logger.error(f"计算动态验证码失败：{str(e)}")
|
|
return ""
|
|
|
|
def get_code(self) -> str:
|
|
"""
|
|
返回显式验证码，或从共享密钥实时计算。
|
|
"""
|
|
return self.code or self.__calc(self.secret)
|