rice_install_client.py

import os
import sys
import tempfile
import requests
from .rice_def import RiceRoundErrorDef
from .rice_url_config import RiceUrlConfig
from .auth_unit import AuthUnit
from .rice_prompt_info import RiceEnvConfig
from .utils import get_local_app_setting_path


class RiceInstallClient:
    def __init__(self):
        self.current_path = os.path.dirname(os.path.abspath(__file__))
        self.app_path = get_local_app_setting_path()
        (
            self.source_executable_filename,
            self.executable_filename,
        ) = self._get_platform_executables()

    def _get_platform_executables(self):
        if sys.platform == "win32":
            return "share_client_windows.exe", "share_client.exe"
        elif sys.platform == "darwin":
            return "share_client_mac", "share_client"
        elif sys.platform == "linux":
            return "share_client_linux", "share_client"
        else:
            raise OSError(f"Unsupported platform: {sys.platform}")

    def is_client_running(self):
        if not self.is_client_installed():
            return False
        lock_path = os.path.join(tempfile.gettempdir(), "rice_client.lock")
        if not os.path.exists(lock_path):
            return False
        try:
            import portalocker

            with open(lock_path, "w") as f:
                portalocker.lock(f, portalocker.LOCK_EX | portalocker.LOCK_NB)
                portalocker.unlock(f)
                return False
        except portalocker.LockException:
            return True

    def is_client_installed(self):
        if not self.app_path.exists():
            return False
        executable_path = self.app_path / self.executable_filename
        if not executable_path.exists():
            return False
        toml_path = self.app_path / "client.toml"
        if not toml_path.exists():
            return False
        return True

    def repair_client_toml(self, client_toml_path):
        if not client_toml_path.exists():
            return False
        try:
            import tomlkit
            from tomlkit import dumps

            env_config = RiceEnvConfig().read_env()
            if not env_config:
                print("Error: Failed to read environment config")
                return False
            with open(client_toml_path, "r", encoding="utf-8") as f:
                toml_data = tomlkit.load(f)
            comfyui_config = toml_data.get("ComfyUI")
            if not comfyui_config or not isinstance(comfyui_config, dict):
                return False
            comfyui_config["PythonPath"] = env_config["PythonPath"]
            comfyui_config["WorkingDirectory"] = env_config["WorkingDirectory"]
            comfyui_config["ComfyuiScriptName"] = env_config["ScriptName"]
            with open(client_toml_path, "w", encoding="utf-8") as f:
                f.write(dumps(toml_data))
            return True
        except Exception as e:
            print(f"Error repairing client.toml: {str(e)}")
            return False

    def _generate_toml_config(
        self, secret_token, comfyui_port=6607, local_server_port=6608
    ):
        try:
            import tomlkit
            from tomlkit import comment, table, dumps

            env_config = RiceEnvConfig().read_env()
            config = tomlkit.document()
            config.add(comment("日志级别设置"))
            config.add(comment("可选值: 'debug', 'info', 'warn', 'error'"))
            config["LogLevel"] = "info"
            config.add(comment("机器码，非常重要，用于登录鉴权"))
            config.add(comment("在官网可以获取自己的机器码，普通用户也可以由管理员授予"))
            config["SecretToken"] = secret_token
            config.add(comment("本地服务端口"))
            config.add(comment("用于本地服务端口，通常为 6608"))
            config["Port"] = local_server_port
            comfyui_table = table()
            comfyui_table.add(comment("ComfyUI 监听的端口"))
            comfyui_table.add(comment("端口号，默认为 6607"))
            comfyui_table["Port"] = comfyui_port
            comfyui_table.add(comment("Python 解释器路径"))
            comfyui_table.add(comment("这里填写你安装的 Python 解释器路径，确保 Python 环境已经配置好"))
            comfyui_table["PythonPath"] = str(env_config["PythonPath"])
            comfyui_table.add(comment("ComfyUI 脚本的文件名"))
            comfyui_table.add(comment("这里填写 ComfyUI 的启动脚本名，通常是 'main.py'"))
            comfyui_table["ComfyuiScriptName"] = env_config["ScriptName"]
            comfyui_table.add(comment("ComfyUI 工作目录"))
            comfyui_table.add(comment("这里填写 ComfyUI 所在的目录路径"))
            comfyui_table["WorkingDirectory"] = str(env_config["WorkingDirectory"])
            comfyui_table.add(comment("环境命令，用于激活相关环境"))
            comfyui_table.add(comment("例如可以填写 conda 环境的激活命令 conda activate comfyui"))
            comfyui_table["EnvCmd"] = ""
            comfyui_table.add(comment("启动时附加的命令行参数"))
            comfyui_table.add(comment("可根据需要添加，常用的如 '--disable-metadata'"))
            comfyui_table["AddCmd"] = env_config["AddCmd"]
            config["ComfyUI"] = comfyui_table
            return dumps(config)
        except Exception as e:
            print(f"Error generating TOML content: {str(e)}")
            raise e

    def install_client_toml(self, comfyui_port, local_server_port, secret_token):
        try:
            os.makedirs(self.app_path, exist_ok=True)
            toml_content = self._generate_toml_config(
                secret_token, comfyui_port, local_server_port
            )
            client_toml_path = self.app_path / "client.toml"
            with open(client_toml_path, "w", encoding="utf-8") as f:
                f.write(toml_content)
            return True
        except Exception as e:
            print(f"Error writing client.toml: {str(e)}")
            return False

    def export_toml(self, secret_token):
        return self._generate_toml_config(secret_token)

    def auto_fix_toml(self, comfyui_port=6607, local_server_port=8689):
        toml_path = self.app_path / "client.toml"
        if not toml_path.exists():
            secret_token, error_message, error_code = self.get_secret_token()
            if not secret_token:
                return (
                    error_code
                    if error_code != RiceRoundErrorDef.SUCCESS
                    else RiceRoundErrorDef.ERROR_SECRET_TOKEN,
                    error_message,
                )
            if not self.install_client_toml(
                comfyui_port, local_server_port, secret_token
            ):
                return RiceRoundErrorDef.ERROR_INSTALL_CLIENT_TOML, "安装client.toml失败"
        else:
            self.repair_client_toml(toml_path)

    def get_secret_token(self):
        token, error_message, error_code = AuthUnit().get_user_token()
        if not token:
            return None, error_message, error_code
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
        }
        try:
            response = requests.get(
                RiceUrlConfig().machine_bind_key_url, headers=headers, timeout=10
            )
            if response.status_code != 200:
                error_code = response.status_code
                return (
                    None,
                    "获取密钥失败",
                    RiceRoundErrorDef.calc_error_code(
                        RiceRoundErrorDef.ERROR_MACHINE_CODE_BASE, error_code
                    ),
                )
            response_data = response.json()
            if response_data.get("code") != 0:
                return None, "获取密钥失败: 响应码不为0", RiceRoundErrorDef.ERROR_SECRET_TOKEN
            secret_token = response_data.get("data", {}).get("key")
            if not secret_token:
                return None, "获取密钥失败: 密钥为空", RiceRoundErrorDef.ERROR_SECRET_TOKEN
            return secret_token, "", 0
        except Exception as e:
            return None, "获取密钥失败" + str(e), RiceRoundErrorDef.ERROR_SECRET_TOKEN