253 lines
10 KiB
Python
253 lines
10 KiB
Python
#!/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
|