LongjoeAgent/utils/registration.py

import json
import base64
import hashlib
from datetime import datetime
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import os
import sys
from typing import Tuple, Optional, Dict


class RegistrationValidator:
    def __init__(self, secret_key: str):
        """
        初始化注册验证器

        Args:
            secret_key: 密钥（必须与生成时相同）
        """
        if len(secret_key) not in [16, 24, 32]:
            # 自动调整密钥长度
            if len(secret_key) < 16:
                secret_key = secret_key.ljust(16, "0")
            elif len(secret_key) < 24:
                secret_key = secret_key.ljust(24, "0")
            elif len(secret_key) < 32:
                secret_key = secret_key.ljust(32, "0")
            else:
                secret_key = secret_key[:32]

        self.secret_key = secret_key.encode("utf-8")

    def _get_cipher(self, iv=None):
        """获取AES加密对象"""
        if iv is None:
            # 使用密钥的前16字节作为固定IV
            iv = hashlib.md5(self.secret_key).digest()[:16]
        return AES.new(self.secret_key, AES.MODE_CBC, iv)

    def decrypt_registration(self, registration_code: str) -> Optional[Dict]:
        """
        解密注册码

        Returns:
            注册信息字典，解密失败返回None
        """
        try:
            # Base64解码
            encrypted = base64.urlsafe_b64decode(registration_code.encode("utf-8"))

            # AES解密
            cipher = self._get_cipher()
            decrypted = unpad(cipher.decrypt(encrypted), AES.block_size)

            # JSON解析
            reg_data = json.loads(decrypted.decode("utf-8"))

            return reg_data
        except Exception as e:
            print(f"解密注册码失败: {e}")
            return None

    def validate(self, device_id: str, registration_code: str) -> Tuple[bool, str]:
        """
        验证注册码

        Args:
            device_id: 当前设备特征码
            registration_code: 注册码

        Returns:
            (是否有效, 消息)
        """
        if not registration_code or registration_code.strip() == "":
            return False, "未提供注册码"

        # 解密注册码
        reg_info = self.decrypt_registration(registration_code)
        if not reg_info:
            return False, "注册码无效或已损坏"

        # 检查设备ID
        reg_device_id = reg_info.get("device_id", "")
        if reg_device_id != device_id:
            return (
                False,
                f"设备不匹配。注册设备: {reg_device_id}，当前设备: {device_id}",
            )

        # 检查有效期
        expire_date_str = reg_info.get("expire_date", "")
        if not expire_date_str:
            return False, "注册信息不完整（缺少有效期）"

        try:
            expire_date = datetime.strptime(expire_date_str, "%Y-%m-%d").date()
            current_date = datetime.now().date()

            if current_date > expire_date:
                days_overdue = (current_date - expire_date).days
                return (
                    False,
                    f"注册已过期 {days_overdue} 天！有效期至: {expire_date_str}",
                )

            # 计算剩余天数
            days_left = (expire_date - current_date).days

            if days_left <= 7:
                return (
                    True,
                    f"注册有效（{days_left}天后到期）有效期至: {expire_date_str}",
                )
            elif days_left <= 30:
                return True, f"注册有效（剩余{days_left}天）有效期至: {expire_date_str}"
            else:
                return True, f"注册有效，有效期至: {expire_date_str}"

        except ValueError as e:
            return False, f"日期格式错误: {str(e)}"


def check_registration_required() -> bool:
    """
    检查是否需要注册验证

    可以通过环境变量控制是否启用注册验证
    """
    return os.getenv("REGISTRATION_REQUIRED", "true").lower() == "true"


def get_registration_code() -> str:
    """
    获取注册码

    优先级：
    1. 环境变量 REGISTRATION_CODE(取消)
    2. 配置文件 .registration
    3. 当前目录的 registration.txt
    """
    # # 从环境变量获取
    # reg_code = os.getenv("REGISTRATION_CODE", "").strip()
    # if reg_code:
    #     return reg_code
    current_file_dir = os.path.dirname(os.path.abspath(__file__))  # utils目录
    project_root = os.path.dirname(current_file_dir)  # 项目根目录
    # 配置文件列表
    config_files = [
        os.path.join(project_root, ".registration"),
        os.path.join(project_root, "registration.txt"),
        os.path.join(project_root, "config", "registration.txt"),
        # 保留当前目录作为备用
        ".registration",
        "registration.txt",
    ]

    for config_file in config_files:
        if os.path.exists(config_file):
            try:
                with open(config_file, "r") as f:
                    for line in f:
                        if line.startswith("REGISTRATION_CODE="):
                            return line.strip().split("=", 1)[1]
                        elif (
                            not line.startswith("#")
                            and "=" not in line
                            and len(line.strip()) > 20
                        ):
                            # 可能是单独的注册码
                            return line.strip()
            except:
                continue

    return ""


def validate_registration(
    secret_key: str, device_id_func=None
) -> Tuple[bool, str, Optional[Dict]]:
    """
    验证注册的主函数

    Args:
        secret_key: 密钥
        device_id_func: 获取设备ID的函数，默认为导入device_id模块

    Returns:
        (是否通过, 消息, 注册信息)
    """
    # 如果不要求注册验证，直接通过
    if not check_registration_required():
        return True, "注册验证已禁用", None

    # 获取设备ID函数
    if device_id_func is None:
        try:
            from device_id import get_device_id

            device_id_func = get_device_id
        except ImportError:
            # 如果在同一文件中定义了get_device_id
            from device_id import get_device_id as device_id_func

    # 获取当前设备ID
    try:
        current_device_id = device_id_func()
    except Exception as e:
        return False, f"获取设备特征码失败: {e}", None

    # 获取注册码
    registration_code = get_registration_code()
    if not registration_code:
        return (
            False,
            "未找到注册码。本机特征码为: {current_device_id}，请申请注册并设置REGISTRATION_CODE环境变量或创建注册文件。",
            None,
        )

    # 创建验证器并验证
    validator = RegistrationValidator(secret_key)

    # 解密注册信息
    reg_info = validator.decrypt_registration(registration_code)
    if not reg_info:
        return False, "注册码无效或已损坏", None

    # 验证
    is_valid, message = validator.validate(current_device_id, registration_code)

    return is_valid, message, reg_info


# 使用示例
if __name__ == "__main__":
    # 测试验证
    secret_key = "ialwayslovelongjoe"
    from device_id import get_device_id

    is_valid, message, reg_info = validate_registration(secret_key, get_device_id)

    print(f"验证结果: {is_valid}")
    print(f"消息: {message}")
    if reg_info:
        print(f"注册信息: {reg_info}")