mirror of
https://github.com/alexgetmancom/miband-bot.git
synced 2026-09-05 18:16:18 +03:00
Alex
452 lines
15 KiB
Python
452 lines
15 KiB
Python
"""小米账号认证管理器。
|
||
|
||
负责编排登录流程、token 持久化。具体登录实现委托给
|
||
``password``、``qr``、``passtoken`` 等子模块。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import tempfile
|
||
from collections.abc import Awaitable, Callable
|
||
from pathlib import Path
|
||
from typing import Self, TypeVar
|
||
|
||
from loguru import logger
|
||
|
||
from mi_fitness.exceptions import (
|
||
AuthError,
|
||
CaptchaRequiredError,
|
||
DeviceUntrustedError,
|
||
TokenExpiredError,
|
||
)
|
||
from mi_fitness.http import RetryAsyncClient
|
||
from mi_fitness.models import AuthToken
|
||
|
||
from . import passtoken as _pt
|
||
from . import password as _pwd
|
||
from . import qr as _qr
|
||
from . import sts as _sts
|
||
from ._helpers import create_login_http
|
||
|
||
_MAX_CAPTCHA_RETRIES = 3
|
||
_CaptchaStepT = TypeVar("_CaptchaStepT")
|
||
|
||
|
||
def _write_secret_text(path: Path, text: str) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent)
|
||
tmp_path = Path(tmp_name)
|
||
try:
|
||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||
fh.write(text)
|
||
fh.flush()
|
||
os.fsync(fh.fileno())
|
||
os.chmod(tmp_path, 0o600)
|
||
os.replace(tmp_path, path)
|
||
os.chmod(path, 0o600)
|
||
finally:
|
||
try:
|
||
tmp_path.unlink()
|
||
except FileNotFoundError:
|
||
pass
|
||
|
||
|
||
class XiaomiAuth:
|
||
"""小米账号认证管理器。
|
||
|
||
负责登录流程、token 持久化。通过 serviceLogin 获取 ssecurity
|
||
和 serviceToken，后续 API 请求使用 RC4 加密。
|
||
|
||
Attributes:
|
||
username: 小米账号（手机号或邮箱）。
|
||
token: 当前认证凭证。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
username: str = "",
|
||
password: str = "",
|
||
*,
|
||
device_id: str = "",
|
||
):
|
||
"""
|
||
Args:
|
||
username: 小米账号。
|
||
password: 密码（仅登录时需要，不会被存储）。
|
||
device_id: 设备标识符。留空则自动生成随机值。新设备首次登录
|
||
会触发短信验证码验证，可通过 ``login()`` 的
|
||
``verification_code_handler`` 回调自动处理。
|
||
"""
|
||
self.username = username
|
||
self._password = password
|
||
self.token = AuthToken()
|
||
if device_id:
|
||
self.token.device_id = device_id
|
||
self._http: RetryAsyncClient | None = None
|
||
self._ticket_token: str = ""
|
||
self._token_path: Path | None = None
|
||
|
||
@classmethod
|
||
def from_token(cls, path: Path | str) -> Self:
|
||
"""从文件加载已有 token，一步完成初始化。
|
||
|
||
Args:
|
||
path: token 文件路径。
|
||
|
||
Returns:
|
||
已加载 token 的认证管理器。
|
||
|
||
Raises:
|
||
AuthError: 文件不存在或格式错误。
|
||
"""
|
||
instance = cls()
|
||
instance.load_token(path)
|
||
return instance
|
||
|
||
def _ensure_http(self) -> RetryAsyncClient:
|
||
"""确保 HTTP 客户端已初始化（惰性创建）。"""
|
||
if self._http is None:
|
||
self._http = create_login_http()
|
||
return self._http
|
||
|
||
def _ensure_device_cookie(self) -> RetryAsyncClient:
|
||
"""确保 deviceId 已生成并写入登录 cookie。"""
|
||
http = self._ensure_http()
|
||
if not self.token.device_id:
|
||
self.token.device_id = f"an_{os.urandom(16).hex()}"
|
||
http.cookies.set("deviceId", self.token.device_id)
|
||
return http
|
||
|
||
async def _run_with_captcha_retries(
|
||
self,
|
||
http: RetryAsyncClient,
|
||
action: Callable[[str], Awaitable[_CaptchaStepT]],
|
||
*,
|
||
captcha_handler: Callable[[bytes], Awaitable[str]] | None = None,
|
||
) -> _CaptchaStepT:
|
||
"""统一处理图形验证码重试。"""
|
||
captcha_code = ""
|
||
for _ in range(_MAX_CAPTCHA_RETRIES):
|
||
try:
|
||
return await action(captcha_code)
|
||
except CaptchaRequiredError as e:
|
||
if captcha_handler is None:
|
||
raise
|
||
image = await _pwd.fetch_captcha_image(http, e.captcha_url)
|
||
captcha_code = await captcha_handler(image)
|
||
raise AuthError(f"图形验证码验证失败：已连续重试 {_MAX_CAPTCHA_RETRIES} 次")
|
||
|
||
# region 公共方法
|
||
async def login(
|
||
self,
|
||
*,
|
||
verification_code_handler: Callable[[str], Awaitable[str]] | None = None,
|
||
captcha_handler: Callable[[bytes], Awaitable[str]] | None = None,
|
||
) -> AuthToken:
|
||
"""执行完整登录流程。
|
||
|
||
Args:
|
||
verification_code_handler: 短信验证码回调。接收脱敏手机号
|
||
（如 ``"191******54"``），返回用户输入的 6 位验证码。
|
||
新设备首次登录需要短信验证时自动调用。
|
||
若未提供且需要验证，将抛出 ``DeviceUntrustedError``。
|
||
captcha_handler: 图形验证码回调。接收验证码图片字节
|
||
（PNG/JPEG），返回用户识别的验证码文本。
|
||
当登录流程触发图形验证码风控时自动调用。
|
||
若未提供且需要验证码，将抛出 ``CaptchaRequiredError``。
|
||
|
||
Returns:
|
||
登录成功后的 AuthToken。
|
||
|
||
Raises:
|
||
AuthError: 登录失败（密码错误等）。
|
||
DeviceUntrustedError: 需要短信验证但未提供回调。
|
||
CaptchaRequiredError: 需要图形验证码但未提供回调。
|
||
"""
|
||
if not self.username or not self._password:
|
||
raise AuthError("用户名和密码不能为空")
|
||
|
||
http = self._ensure_device_cookie()
|
||
|
||
logger.info("开始小米账号登录: {}", self.username)
|
||
|
||
sign, callback = await _pwd.get_login_page(http)
|
||
|
||
try:
|
||
await _pwd.submit_login(
|
||
http,
|
||
self.token,
|
||
self.username,
|
||
self._password,
|
||
sign,
|
||
callback,
|
||
)
|
||
except DeviceUntrustedError:
|
||
if verification_code_handler is None:
|
||
raise
|
||
phone = await self.send_verification_code(
|
||
captcha_handler=captcha_handler,
|
||
)
|
||
code = await verification_code_handler(phone)
|
||
await self.login_with_verification_code(code)
|
||
return self.token
|
||
|
||
await _sts.sts_exchange(http, self.token)
|
||
|
||
self._password = ""
|
||
logger.info("登录成功, user_id={}", self.token.user_id)
|
||
return self.token
|
||
|
||
def save_token(self, path: Path | str) -> None:
|
||
"""将 token 保存到 JSON 文件，便于下次免登录恢复。
|
||
|
||
Args:
|
||
path: 保存路径。
|
||
"""
|
||
path = Path(path)
|
||
_write_secret_text(path, self.token.model_dump_json(indent=2) + "\n")
|
||
self._token_path = path
|
||
logger.info("Token 已保存至 {}", path)
|
||
|
||
def load_token(self, path: Path | str) -> AuthToken:
|
||
"""从文件加载 token。
|
||
|
||
Args:
|
||
path: token 文件路径。
|
||
|
||
Returns:
|
||
加载的 AuthToken。
|
||
|
||
Raises:
|
||
AuthError: 文件不存在或格式错误。
|
||
"""
|
||
path = Path(path)
|
||
if not path.exists():
|
||
raise AuthError(f"Token 文件不存在: {path}")
|
||
try:
|
||
data = path.read_text(encoding="utf-8")
|
||
self.token = AuthToken.model_validate_json(data)
|
||
self._token_path = path
|
||
logger.info("Token 已从 {} 加载, user_id={}", path, self.token.user_id)
|
||
return self.token
|
||
except Exception as e:
|
||
raise AuthError(f"Token 文件解析失败: {e}") from e
|
||
|
||
@property
|
||
def is_authenticated(self) -> bool:
|
||
"""检查是否已登录。"""
|
||
return bool(self.token.service_token and self.token.ssecurity)
|
||
|
||
@property
|
||
def can_refresh(self) -> bool:
|
||
"""当前 token 是否具备自动刷新条件。"""
|
||
return bool(self.token.pass_token and self.token.user_id)
|
||
|
||
async def close(self) -> None:
|
||
"""关闭 HTTP 客户端（如有）。"""
|
||
if self._http is not None:
|
||
await self._http.aclose()
|
||
self._http = None
|
||
|
||
async def refresh(self) -> AuthToken:
|
||
"""用已有 passToken 刷新 serviceToken / ssecurity。
|
||
|
||
Returns:
|
||
刷新后的 AuthToken。
|
||
|
||
Raises:
|
||
TokenExpiredError: 当前 token 无法刷新，或刷新失败。
|
||
"""
|
||
if not self.can_refresh:
|
||
raise TokenExpiredError("Token 已过期，且缺少 passToken 或 user_id，无法自动刷新")
|
||
|
||
logger.info("开始刷新登录凭证, user_id={}", self.token.user_id)
|
||
try:
|
||
token = await self.login_passtoken(
|
||
pass_token=self.token.pass_token,
|
||
user_id=self.token.user_id,
|
||
device_id=self.token.device_id,
|
||
)
|
||
except AuthError as e:
|
||
raise TokenExpiredError(f"Token 已过期，自动刷新失败: {e}") from e
|
||
|
||
if self._token_path is not None:
|
||
self.save_token(self._token_path)
|
||
|
||
logger.info("登录凭证刷新成功, user_id={}", token.user_id)
|
||
return token
|
||
|
||
async def send_verification_code(
|
||
self,
|
||
*,
|
||
captcha_handler: Callable[[bytes], Awaitable[str]] | None = None,
|
||
) -> str:
|
||
"""发送短信验证码到用户手机。
|
||
|
||
在 ``login()`` 抛出 ``DeviceUntrustedError`` 后调用此方法
|
||
手动发起短信验证流程。
|
||
|
||
Args:
|
||
captcha_handler: 图形验证码回调。接收验证码图片字节，
|
||
返回用户识别的验证码文本。未提供时触发验证码将直接抛出
|
||
``CaptchaRequiredError``。
|
||
|
||
Returns:
|
||
脱敏手机号（如 ``"191******54"``）。
|
||
|
||
Raises:
|
||
AuthError: 获取手机信息或发送验证码失败。
|
||
CaptchaRequiredError: 需要图形验证码但未提供回调。
|
||
"""
|
||
http = self._ensure_http()
|
||
await _pwd.ensure_ticket_login_ready(http)
|
||
|
||
await self._run_with_captcha_retries(
|
||
http,
|
||
lambda captcha_code: _pwd.send_ticket(
|
||
http,
|
||
self.username,
|
||
captcha_code=captcha_code,
|
||
),
|
||
captcha_handler=captcha_handler,
|
||
)
|
||
phone, ticket_token = await self._run_with_captcha_retries(
|
||
http,
|
||
lambda captcha_code: _pwd.get_phone_info(
|
||
http,
|
||
self.username,
|
||
captcha_code=captcha_code,
|
||
),
|
||
captcha_handler=captcha_handler,
|
||
)
|
||
|
||
self._ticket_token = ticket_token
|
||
logger.info("验证码已发送至 {}", phone)
|
||
return phone
|
||
|
||
async def login_with_verification_code(self, code: str) -> AuthToken:
|
||
"""使用短信验证码完成登录。
|
||
|
||
在 ``send_verification_code()`` 之后调用，提交用户收到的验证码。
|
||
|
||
Args:
|
||
code: 6 位短信验证码。
|
||
|
||
Returns:
|
||
登录成功后的 AuthToken。
|
||
|
||
Raises:
|
||
AuthError: 验证码错误或登录失败。
|
||
"""
|
||
if not self._ticket_token:
|
||
raise AuthError("请先调用 send_verification_code() 发送验证码")
|
||
|
||
http = self._ensure_http()
|
||
http.cookies.set("ticketToken", self._ticket_token)
|
||
|
||
sign, callback = await _pwd.get_login_page(http, login_sign="ticket")
|
||
await _pwd.submit_ticket_auth(http, self.token, code, sign, callback)
|
||
await _sts.sts_exchange(http, self.token)
|
||
|
||
self._password = ""
|
||
self._ticket_token = ""
|
||
logger.info("短信验证码登录成功, user_id={}", self.token.user_id)
|
||
return self.token
|
||
|
||
def __repr__(self) -> str:
|
||
status = "已认证" if self.is_authenticated else "未认证"
|
||
uid = self.token.user_id or "N/A"
|
||
return f"XiaomiAuth(user={self.username or uid!r}, {status})"
|
||
|
||
async def login_qr(
|
||
self,
|
||
*,
|
||
qr_callback: Callable[[str, str], Awaitable[None]] | None = None,
|
||
poll_interval: float = 2.0,
|
||
max_wait: float = 300.0,
|
||
) -> AuthToken:
|
||
"""二维码扫码登录（无需密码，绕过验证码风控）。
|
||
|
||
用户用小米账号 APP 扫描二维码完成登录，SDK 通过长轮询
|
||
检测扫码结果并自动提取凭证。
|
||
|
||
Args:
|
||
qr_callback: 二维码展示回调。接收 ``(qr_image_url, login_url)``，
|
||
其中 ``qr_image_url`` 是二维码图片 URL（可下载显示），
|
||
``login_url`` 是备选的浏览器登录链接。
|
||
默认将信息打印到控制台。
|
||
poll_interval: 长轮询间隔（秒）。
|
||
max_wait: 扫码超时时间（秒）。
|
||
|
||
Returns:
|
||
登录成功后的 AuthToken。
|
||
|
||
Raises:
|
||
AuthError: 获取二维码失败或扫码超时。
|
||
"""
|
||
http = self._ensure_device_cookie()
|
||
|
||
await _qr.login_qr(
|
||
http,
|
||
self.token,
|
||
qr_callback=qr_callback,
|
||
poll_interval=poll_interval,
|
||
max_wait=max_wait,
|
||
)
|
||
|
||
await _sts.sts_exchange(http, self.token)
|
||
|
||
logger.info("二维码登录成功, user_id={}", self.token.user_id)
|
||
return self.token
|
||
|
||
async def login_passtoken(
|
||
self,
|
||
*,
|
||
pass_token: str = "",
|
||
user_id: str = "",
|
||
device_id: str = "",
|
||
) -> AuthToken:
|
||
"""使用 passToken 换取完整登录凭证（无需密码）。
|
||
|
||
passToken 可通过 ``migate.get_passtoken()`` 或浏览器登录小米账号
|
||
后从 Cookie 中提取获取。此方法用 passToken 调用 ``serviceLogin``
|
||
换取 ``ssecurity`` 和 ``serviceToken``。
|
||
|
||
Args:
|
||
pass_token: 小米账号 passToken。
|
||
user_id: 小米账号 userId。
|
||
device_id: 设备标识符（可选）。
|
||
|
||
Returns:
|
||
登录成功后的 AuthToken。
|
||
|
||
Raises:
|
||
AuthError: passToken 无效或换取凭证失败。
|
||
"""
|
||
http = self._ensure_http()
|
||
|
||
await _pt.login_passtoken(
|
||
http,
|
||
self.token,
|
||
pass_token=pass_token,
|
||
user_id=user_id,
|
||
device_id=device_id,
|
||
)
|
||
|
||
await _sts.sts_exchange(http, self.token)
|
||
|
||
logger.info("passToken 登录成功, user_id={}", self.token.user_id)
|
||
return self.token
|
||
|
||
# endregion
|
||
|
||
# region 上下文管理器
|
||
async def __aenter__(self) -> Self:
|
||
return self
|
||
|
||
async def __aexit__(self, *args: object) -> None:
|
||
await self.close()
|
||
|
||
# endregion
|