RiceRound-ComfyUI_RiceRound/auth_unit.py

import os
import time
import requests
import configparser
from .rice_def import RiceRoundErrorDef
from .utils import get_local_app_setting_path, get_machine_id, generate_random_string
from .rice_url_config import RiceUrlConfig
from server import PromptServer


class AuthUnit:
    _instance = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super(AuthUnit, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "initialized"):
            self.machine_id = get_machine_id()
            local_app_path = get_local_app_setting_path()
            local_app_path.mkdir(parents=True, exist_ok=True)
            self.config_path = local_app_path / "config.ini"
            self.last_check_time = 0
            self.initialized = True
            self.user_id = 0

    def empty_token(self, need_clear=False):
        self.token = ""
        self.last_check_time = 0
        if need_clear:
            self.clear_user_token()

    def get_user_token(self):
        self.token = self.read_user_token()
        if (
            time.time() - self.last_check_time > 120
            and self.token
            and len(self.token) > 50
        ):
            try:
                headers = {
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.token}",
                }
                response = requests.get(
                    RiceUrlConfig().get_info_url, headers=headers, timeout=10
                )
                if response.status_code == 200:
                    user_info = response.json()
                    try:
                        self.user_id = int(user_info.get("user_id", 0) or 0)
                    except (ValueError, TypeError):
                        self.user_id = 0
                    self.last_check_time = time.time()
                    return self.token, "", RiceRoundErrorDef.SUCCESS
                else:
                    error_message = "登录结果错误"
                    error_code = RiceRoundErrorDef.UNKNOWN_ERROR
                    try:
                        response_data = response.json()
                        if "message" in response_data:
                            error_message = response_data["message"]
                    except ValueError:
                        pass
                    if response.status_code == 401:
                        error_message = "登录已过期，请重新登录"
                        error_code = RiceRoundErrorDef.HTTP_UNAUTHORIZED
                    elif response.status_code == 500:
                        error_message = "服务器内部错误，请稍后重试"
                        error_code = RiceRoundErrorDef.HTTP_INTERNAL_ERROR
                    elif response.status_code == 503:
                        error_message = "服务不可用，请稍后重试"
                        error_code = RiceRoundErrorDef.HTTP_SERVICE_UNAVAILABLE
                    self.empty_token(response.status_code == 401)
                    return None, error_message, error_code
            except requests.exceptions.Timeout:
                self.empty_token()
                return None, "请求超时，请检查网络连接", RiceRoundErrorDef.HTTP_TIMEOUT
            except requests.exceptions.ConnectionError:
                self.empty_token()
                return None, "网络连接失败，请检查网络", RiceRoundErrorDef.NETWORK_ERROR
            except requests.exceptions.RequestException as e:
                self.empty_token()
                return None, f"请求失败: {str(e)}", RiceRoundErrorDef.REQUEST_ERROR
        if self.token and len(self.token) > 50:
            return self.token, "", RiceRoundErrorDef.SUCCESS
        return None, "未读取到有效的token，请重新登录", RiceRoundErrorDef.NO_TOKEN_ERROR

    def get_user_info(self):
        _, error_message, error_code = self.get_user_token()
        if error_code == RiceRoundErrorDef.SUCCESS and self.user_id:
            return error_code, self.user_id
        else:
            return error_code, error_message

    def login_dialog(self, title=""):
        self.client_key = generate_random_string(8)
        PromptServer.instance.send_sync(
            "riceround_login_dialog",
            {
                "machine_id": self.machine_id,
                "client_key": self.client_key,
                "title": title,
            },
        )

    def read_user_token(self):
        if not os.path.exists(self.config_path):
            return ""
        try:
            config = configparser.ConfigParser()
            config.read(self.config_path, encoding="utf-8")
            return config.get("Auth", "user_token", fallback="")
        except Exception as e:
            print(f"Error reading token: {e}")
            return ""

    def set_user_token(self, user_token, client_key):
        if not client_key or self.client_key != client_key:
            return
        if not user_token:
            user_token = ""
            print("user_token is empty")
        self.save_user_token(user_token)

    def save_user_token(self, user_token):
        try:
            config = configparser.ConfigParser()
            try:
                if os.path.exists(self.config_path):
                    config.read(self.config_path, encoding="utf-8")
            except Exception as read_error:
                print(f"Warning: Error reading existing config: {read_error}")
            if "Auth" not in config:
                config.add_section("Auth")
            config["Auth"]["user_token"] = user_token
            with open(self.config_path, "w", encoding="utf-8") as f:
                config.write(f)
        except Exception as e:
            print(f"Error saving token: {e}")
            raise RuntimeError(f"Failed to save token: {e}")

    def set_long_token(self, long_token):
        if not long_token:
            return
        self.save_user_token(long_token)
        self.client_key = ""

    def clear_user_token(self):
        PromptServer.instance.send_sync(
            "riceround_clear_user_info", {"clear_key": "all"}
        )
        if os.path.exists(self.config_path):
            try:
                config = configparser.ConfigParser()
                config.read(self.config_path, encoding="utf-8")
                if "Auth" not in config:
                    return
                if "user_token" not in config["Auth"]:
                    return
                config["Auth"]["user_token"] = ""
                with open(self.config_path, "w", encoding="utf-8") as f:
                    config.write(f)
            except Exception as e:
                print(f"Error clearing token: {e}")
                raise RuntimeError(f"Failed to clear token: {e}")