#!/usr/bin/env python # -*- coding: utf-8 -*- """ Substance Painter 注册表管理器 用于在启动前配置库，关闭后清理库 """ import winreg import re import os class SPRegistryManager: """SP 注册表管理器""" REGISTRY_KEY = r"SOFTWARE\Adobe\Adobe Substance 3D Painter\Shelf\pathInfos" SHELF_KEY = r"SOFTWARE\Adobe\Adobe Substance 3D Painter\Shelf" @staticmethod def normalize_library_name(name): """规范化库名称（小写、下划线、连字符）""" safe_name = name.lower().replace(' ', '_') safe_name = re.sub(r'[^a-z0-9_-]', '', safe_name) return safe_name @staticmethod def get_all_libraries(): """获取所有库配置 Returns: list: [(id, name, path, disabled), ...] """ try: reg_conn = winreg.ConnectRegistry(None, winreg.HKEY_CURRENT_USER) key = winreg.OpenKey(reg_conn, SPRegistryManager.REGISTRY_KEY, 0, winreg.KEY_READ) libraries = [] sub_key_count = winreg.QueryInfoKey(key)[0] for i in range(sub_key_count): sub_key_name = winreg.EnumKey(key, i) sub_key = winreg.OpenKey(reg_conn, f"{SPRegistryManager.REGISTRY_KEY}\\{sub_key_name}", 0, winreg.KEY_READ) try: name = winreg.QueryValueEx(sub_key, "name")[0] path = winreg.QueryValueEx(sub_key, "path")[0] disabled = winreg.QueryValueEx(sub_key, "disabled")[0] libraries.append((int(sub_key_name), name, path, disabled)) except: pass finally: winreg.CloseKey(sub_key) winreg.CloseKey(key) return libraries except Exception as e: print(f"[Registry] Error reading libraries: {e}") return [] @staticmethod def remove_project_libraries(): """删除所有项目库（保留系统库） 系统库：your_assets, starter_assets, system_fonts, user_fonts """ system_libs = ['your_assets', 'starter_assets', 'system_fonts', 'user_fonts'] try: libraries = SPRegistryManager.get_all_libraries() to_remove = [] for lib_id, name, path, disabled in libraries: if name not in system_libs: to_remove.append(lib_id) print(f"[Registry] Will remove: {name} (ID: {lib_id})") if not to_remove: print(f"[Registry] No project libraries to remove") return True # 删除项目库 reg_conn = winreg.ConnectRegistry(None, winreg.HKEY_CURRENT_USER) for lib_id in to_remove: try: winreg.DeleteKey(winreg.HKEY_CURRENT_USER, f"{SPRegistryManager.REGISTRY_KEY}\\{lib_id}") print(f"[Registry] Removed library ID: {lib_id}") except Exception as e: print(f"[Registry] Error removing {lib_id}: {e}") # 重新编号剩余的库（确保连续） SPRegistryManager._reindex_libraries() return True except Exception as e: print(f"[Registry] Error removing libraries: {e}") return False @staticmethod def _reindex_libraries(): """重新编号库，确保 ID 连续""" try: libraries = SPRegistryManager.get_all_libraries() libraries.sort(key=lambda x: x[0]) # 按 ID 排序 reg_conn = winreg.ConnectRegistry(None, winreg.HKEY_CURRENT_USER) # 如果已经连续，不需要重新编号 expected_ids = list(range(1, len(libraries) + 1)) actual_ids = [lib[0] for lib in libraries] if expected_ids == actual_ids: print(f"[Registry] Library IDs are already continuous") return print(f"[Registry] Reindexing libraries...") # 先将所有库移到临时 ID temp_offset = 1000 for lib_id, name, path, disabled in libraries: old_key_path = f"{SPRegistryManager.REGISTRY_KEY}\\{lib_id}" new_key_path = f"{SPRegistryManager.REGISTRY_KEY}\\{lib_id + temp_offset}" # 读取旧值 old_key = winreg.OpenKey(reg_conn, old_key_path, 0, winreg.KEY_READ) name_val = winreg.QueryValueEx(old_key, "name")[0] path_val = winreg.QueryValueEx(old_key, "path")[0] disabled_val = winreg.QueryValueEx(old_key, "disabled")[0] winreg.CloseKey(old_key) # 创建新键 new_key = winreg.CreateKey(winreg.HKEY_CURRENT_USER, new_key_path) winreg.SetValueEx(new_key, "name", 0, winreg.REG_SZ, name_val) winreg.SetValueEx(new_key, "path", 0, winreg.REG_SZ, path_val) winreg.SetValueEx(new_key, "disabled", 0, winreg.REG_SZ, disabled_val) winreg.CloseKey(new_key) # 删除旧键 winreg.DeleteKey(winreg.HKEY_CURRENT_USER, old_key_path) # 再将临时 ID 移到正确的连续 ID for new_id, (old_id, name, path, disabled) in enumerate(libraries, start=1): temp_id = old_id + temp_offset temp_key_path = f"{SPRegistryManager.REGISTRY_KEY}\\{temp_id}" final_key_path = f"{SPRegistryManager.REGISTRY_KEY}\\{new_id}" # 读取临时键 temp_key = winreg.OpenKey(reg_conn, temp_key_path, 0, winreg.KEY_READ) name_val = winreg.QueryValueEx(temp_key, "name")[0] path_val = winreg.QueryValueEx(temp_key, "path")[0] disabled_val = winreg.QueryValueEx(temp_key, "disabled")[0] winreg.CloseKey(temp_key) # 创建最终键 final_key = winreg.CreateKey(winreg.HKEY_CURRENT_USER, final_key_path) winreg.SetValueEx(final_key, "name", 0, winreg.REG_SZ, name_val) winreg.SetValueEx(final_key, "path", 0, winreg.REG_SZ, path_val) winreg.SetValueEx(final_key, "disabled", 0, winreg.REG_SZ, disabled_val) winreg.CloseKey(final_key) # 删除临时键 winreg.DeleteKey(winreg.HKEY_CURRENT_USER, temp_key_path) # 更新 size key = winreg.OpenKeyEx(reg_conn, SPRegistryManager.REGISTRY_KEY, 0, winreg.KEY_SET_VALUE) winreg.SetValueEx(key, "size", 0, winreg.REG_DWORD, len(libraries)) winreg.CloseKey(key) print(f"[Registry] Reindexing complete, new size: {len(libraries)}") except Exception as e: print(f"[Registry] Error reindexing: {e}") import traceback traceback.print_exc() @staticmethod def add_project_library(library_name, library_path, set_as_default=True): """添加项目库 Args: library_name: 库名称 library_path: 库路径 set_as_default: 是否设置为默认 Returns: bool: 是否成功 """ try: safe_name = SPRegistryManager.normalize_library_name(library_name) normalized_path = library_path.replace('\\', '/') print(f"[Registry] Adding library: {safe_name}") print(f"[Registry] Path: {normalized_path}") # 检查是否已存在 libraries = SPRegistryManager.get_all_libraries() for lib_id, name, path, disabled in libraries: if name == safe_name: print(f"[Registry] Library already exists: {name}") return True # 找到下一个 ID if libraries: next_id = max(lib[0] for lib in libraries) + 1 else: next_id = 1 # 创建新库 reg_conn = winreg.ConnectRegistry(None, winreg.HKEY_CURRENT_USER) new_key_path = f"{SPRegistryManager.REGISTRY_KEY}\\{next_id}" new_key = winreg.CreateKey(winreg.HKEY_CURRENT_USER, new_key_path) winreg.SetValueEx(new_key, "name", 0, winreg.REG_SZ, safe_name) winreg.SetValueEx(new_key, "path", 0, winreg.REG_SZ, normalized_path) winreg.SetValueEx(new_key, "disabled", 0, winreg.REG_SZ, "false") winreg.CloseKey(new_key) # 更新 size key = winreg.OpenKeyEx(reg_conn, SPRegistryManager.REGISTRY_KEY, 0, winreg.KEY_SET_VALUE) winreg.SetValueEx(key, "size", 0, winreg.REG_DWORD, len(libraries) + 1) winreg.CloseKey(key) print(f"[Registry] Library added with ID: {next_id}") # 设置为默认库 if set_as_default: shelf_key = winreg.OpenKeyEx(reg_conn, SPRegistryManager.SHELF_KEY, 0, winreg.KEY_SET_VALUE) winreg.SetValueEx(shelf_key, "writableShelf", 0, winreg.REG_SZ, safe_name) winreg.CloseKey(shelf_key) print(f"[Registry] Set as default library") return True except Exception as e: print(f"[Registry] Error adding library: {e}") import traceback traceback.print_exc() return False @staticmethod def needs_restart(): """检查是否需要重启 SP 以应用更改 Returns: bool: 是否需要重启 """ # 如果 SP 正在运行，注册表更改需要重启才能生效 import psutil for proc in psutil.process_iter(['name']): if 'Adobe Substance 3D Painter' in proc.info['name']: return True return False