nodes/ReSerendipity-ComfyUI-BatchPromptLoader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

security: confine user-supplied paths to ComfyUI root; localize node UI to English

Browse Source 
Addresses maintainer review feedback on Comfy-Org/ComfyUI-Manager#3183.

Path traversal:
- Add _confine_to_base(), which resolves the real path (following symlinks)
  and verifies it stays inside the ComfyUI root directory. Applied to
  folder_path, log_folder and output_folder.
- Previously an absolute path was trusted verbatim and relative paths could
  escape via '..', allowing a graph to read arbitrary files and to write or
  delete .json files anywhere the server could reach.
- Log clearing now only removes JSON files previously written by this node
  (validated by the 'prompt_file' key) instead of every .json in the folder.

UI localization:
- Node display name and all widget labels are now English.
This commit is contained in:
ReSerendipity
2026-09-02 16:46:18 +08:00
parent e75614c151
commit 98a61ad89e
2 changed files with 136 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+1 -1
View File
@@ -6,5 +6,5 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BatchPromptReaderWithClip": "批量提示词加载器 (Batch Prompt Loader)",
"BatchPromptReaderWithClip": "Batch Prompt Loader",
}
batch_prompt_loader.py
+135 -103
View File
@@ -3,6 +3,9 @@ ComfyUI Batch Prompt Loader - Enhanced Version
A custom node for loading and encoding prompts from TXT files with advanced features:
- Recursive scanning, wildcard filtering, skip exists, metadata output, preview display
- External log recording: records which prompt files have been processed, and auto-skips them next run
All user-supplied paths (folder_path / log_folder / output_folder) are confined to the
ComfyUI root directory to prevent path traversal.
"""
import os
import torch
@@ -15,8 +18,28 @@ _global_mode_state = {}
_log_cleared_this_session = False
def _confine_to_base(base_dir, user_path, label):
"""Resolve a user-supplied path and verify it stays inside base_dir.
Mitigates path traversal: absolute paths pointing outside the ComfyUI root and
relative paths escaping via '..' are rejected. Symlinks are resolved via realpath.
"""
base_real = os.path.normcase(os.path.realpath(base_dir))
if os.path.isabs(user_path):
candidate = user_path
else:
candidate = os.path.join(base_dir, user_path)
candidate_real = os.path.normcase(os.path.realpath(candidate))
if candidate_real != base_real and not candidate_real.startswith(base_real + os.sep):
raise Exception(
f"[BatchPromptLoader] '{label}' must resolve to a location inside the ComfyUI directory.\n"
f"Resolved path: {candidate_real}\nAllowed base: {base_real}"
)
return candidate_real
def _sanitize_log_name(rel_path):
"""将相对路径转为安全的日志文件名（替换路径分隔符和非法字符）"""
"""Convert a relative prompt path into a safe log filename (replacing separators and illegal chars)."""
name = rel_path.replace('\\', '_').replace('/', '_')
for ch in '<>:"/\\|?*':
name = name.replace(ch, '_')
@@ -24,7 +47,7 @@ def _sanitize_log_name(rel_path):
def _load_log(log_dir, source_dir):
"""读取日志文件夹中匹配当前来源目录的已处理文件集合"""
"""Read the set of already-processed prompt files matching the current source directory."""
processed = set()
if not os.path.isdir(log_dir):
return processed
@@ -40,12 +63,12 @@ def _load_log(log_dir, source_dir):
except Exception:
continue
except Exception as e:
print(f"[BatchPromptLoader] Warning: 读取日志失败: {e}")
print(f"[BatchPromptLoader] Warning: failed to read log: {e}")
return processed
def _write_log(log_dir, source_dir, prompt_file, index, total_count):
"""将已处理的提示词文件写入日志（每个文件一个 JSON，便于选择性删除）"""
"""Record a processed prompt file (one JSON per file, so entries can be deleted selectively)."""
try:
os.makedirs(log_dir, exist_ok=True)
log_name = _sanitize_log_name(prompt_file) + '.json'
@@ -59,7 +82,8 @@ def _write_log(log_dir, source_dir, prompt_file, index, total_count):
with open(os.path.join(log_dir, log_name), 'w', encoding='utf-8') as f:
json.dump(entry, f, ensure_ascii=False, indent=2)
except Exception as e:
print(f"[BatchPromptLoader] Warning: 写入日志失败: {e}")
print(f"[BatchPromptLoader] Warning: failed to write log: {e}")
class BatchPromptReaderWithClip:
@classmethod
@@ -70,7 +94,7 @@ class BatchPromptReaderWithClip:
"folder_path": ("STRING", {
"default": "input/batch_prompts",
"multiline": False,
"label": "文件夹路径"
"label": "Folder Path"
}),
"current_number": ("INT", {
"default": 0,
@@ -78,44 +102,44 @@ class BatchPromptReaderWithClip:
"max": 999999,
"step": 1,
"control_after_generate": ["fixed", "increment", "decrement", "random"],
"label": "当前编号"
"label": "Current Index"
}),
"recursive": ("BOOLEAN", {
"default": True,
"label": "递归扫描子文件夹"
"label": "Recursive Scan"
}),
"reverse_order": ("BOOLEAN", {
"default": False,
"label": "倒序排列文件"
"label": "Reverse Order"
}),
"file_pattern": ("STRING", {
"default": "*.txt",
"multiline": False,
"label": "文件名过滤 (通配符)"
"label": "Filename Pattern (wildcard)"
}),
"output_folder": ("STRING", {
"default": "output/",
"multiline": False,
"label": "输出目录 (用于跳过检查)"
"label": "Output Folder (skip check)"
}),
"enable_logging": ("BOOLEAN", {
"default": False,
"label": "启用日志记录 (跳过已处理)"
"label": "Enable Logging (skip processed)"
}),
"log_folder": ("STRING", {
"default": "user/default/batch_prompt_logs",
"multiline": False,
"label": "日志文件夹 (外置可清除)"
"label": "Log Folder (external, clearable)"
}),
"clear_log_on_start": ("BOOLEAN", {
"default": False,
"label": "会话启动时清空日志"
"label": "Clear Log On Session Start"
}),
},
"optional": {
"skip_exists": ("BOOLEAN", {
"default": False,
"label": "跳过已存在的图片"
"label": "Skip Existing Images"
}),
},
"hidden": {
@@ -123,112 +147,118 @@ class BatchPromptReaderWithClip:
"unique_id": "UNIQUE_ID"
}
}
RETURN_TYPES = ("CONDITIONING", "STRING", "INT", "INT")
RETURN_NAMES = ("conditioning", "filename", "index", "total_count")
FUNCTION = "read_and_encode"
OUTPUT_NODE = True
CATEGORY = "batch_tools"
def read_and_encode(self, clip, folder_path, current_number, recursive=True, reverse_order=False,
def read_and_encode(self, clip, folder_path, current_number, recursive=True, reverse_order=False,
file_pattern="*.txt", output_folder="output/", skip_exists=False,
enable_logging=False, log_folder="user/default/batch_prompt_logs", clear_log_on_start=False,
extra_pnginfo=None, unique_id=None):
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if os.path.isabs(folder_path):
target_dir = folder_path
else:
target_dir = os.path.join(base_dir, folder_path)
# 优化 3: 增强错误提示 - 先检查路径是否存在
# Confine user-supplied paths to the ComfyUI root (path traversal protection)
target_dir = _confine_to_base(base_dir, folder_path, "folder_path")
if not os.path.exists(target_dir):
raise Exception(f"[错误] 路径不存在：{target_dir}\n\n请检查 folder_path 是否正确。")
raise Exception(
f"[BatchPromptLoader] Folder does not exist: {target_dir}\n\n"
f"Please check the 'folder_path' value."
)
if not os.path.isdir(target_dir):
raise Exception(f"[错误] 路径不是文件夹：{target_dir}\n\n请确保这是一个有效的目录。")
raise Exception(
f"[BatchPromptLoader] Not a directory: {target_dir}\n\n"
f"Please make sure this is a valid folder."
)
try:
txt_files = []
# 解析文件模式
pattern_lower = file_pattern.lower().strip()
if not pattern_lower or pattern_lower == '*.txt':
search_pattern = '*.txt'
required_ext = '.txt'
else:
search_pattern = pattern_lower
# 提取扩展名（如 *.prompt.txt -> .prompt.txt）
if '*' in search_pattern:
ext_part = search_pattern.split('*', 1)[-1]
required_ext = ext_part if ext_part else '.txt'
else:
required_ext = '.txt'
if recursive:
# 递归扫描所有子文件夹
for root, dirs, files in os.walk(target_dir):
for f in files:
# 优化 2: 通配符过滤
if fnmatch.fnmatch(f.lower(), search_pattern.lower()):
if required_ext and not f.lower().endswith(required_ext):
continue
rel_path = os.path.relpath(os.path.join(root, f), target_dir)
txt_files.append(rel_path)
else:
# 只读取根目录
all_files = os.listdir(target_dir)
for f in all_files:
if fnmatch.fnmatch(f.lower(), search_pattern.lower()):
if required_ext and not f.lower().endswith(required_ext):
continue
txt_files.append(f)
# 排序
txt_files.sort(key=lambda x: x.lower())
# 倒序处理
if reverse_order:
txt_files.reverse()
except Exception as e:
raise Exception(f"[错误] 读取文件夹失败：{target_dir}\n\n错误信息：{str(e)}")
# 优化 3: 空文件夹检测
raise Exception(f"[BatchPromptLoader] Failed to read folder: {target_dir}\n\nError: {str(e)}")
if len(txt_files) == 0:
raise Exception(f"[错误] 文件夹中没有匹配的文件：{target_dir}\n\n搜索模式：{search_pattern}\n\n请添加对应的 TXT 文件或调整文件名过滤规则。")
# 优化 6: 外置日志记录 - 跳过已处理过的提示词文件
raise Exception(
f"[BatchPromptLoader] No matching files in folder: {target_dir}\n\n"
f"Search pattern: {search_pattern}\n\n"
f"Please add matching TXT files or adjust the filename pattern."
)
# External logging: skip already-processed prompt files
log_dir = None
if enable_logging:
if os.path.isabs(log_folder):
log_dir = log_folder
else:
log_dir = os.path.join(base_dir, log_folder)
# 会话启动后第一次运行时清空日志（配合 clear_log_on_start）
log_dir = _confine_to_base(base_dir, log_folder, "log_folder")
global _log_cleared_this_session
if clear_log_on_start and not _log_cleared_this_session:
try:
if os.path.isdir(log_dir):
for fn in os.listdir(log_dir):
if fn.endswith('.json'):
os.remove(os.path.join(log_dir, fn))
if not fn.endswith('.json'):
continue
# Only remove log files previously written by this node
log_path = os.path.join(log_dir, fn)
try:
with open(log_path, 'r', encoding='utf-8') as lf:
data = json.load(lf)
if isinstance(data, dict) and 'prompt_file' in data:
os.remove(log_path)
except Exception:
continue
_log_cleared_this_session = True
print("[BatchPromptLoader] 🧹 已清空日志文件夹（会话启动后第一次运行）")
print("[BatchPromptLoader] Cleared log folder (first run this session)")
except Exception as e:
print(f"[BatchPromptLoader] Warning: 清空日志失败: {e}")
print(f"[BatchPromptLoader] Warning: failed to clear log: {e}")
processed_files = _load_log(log_dir, os.path.normpath(target_dir))
if processed_files:
remaining = [f for f in txt_files if f not in processed_files]
print(f"[BatchPromptLoader] 📋 日志过滤：已跳过 {len(txt_files) - len(remaining)} 个已处理文件，剩余 {len(remaining)} 个待处理")
print(f"[BatchPromptLoader] Log filter: skipped {len(txt_files) - len(remaining)} "
f"processed file(s), {len(remaining)} remaining")
if not remaining:
print("[BatchPromptLoader] ✅ 所有提示词已处理完成（清除日志文件夹可重新开始）")
print("[BatchPromptLoader] All prompts already processed "
"(clear the log folder to start over)")
empty_cond = [[torch.zeros(1, 1, 2816), {"pooled_output": torch.zeros(1, 2816)}]]
return (empty_cond, "ALL_DONE", 0, 0)
txt_files = remaining
counter_file = os.path.join(base_dir, "user", "default", "batch_prompt_counter.json")
try:
if os.path.exists(counter_file) and os.path.getsize(counter_file) > 0:
@@ -236,91 +266,94 @@ class BatchPromptReaderWithClip:
counter_data = json.load(f)
else:
counter_data = {"last_index": 0}
except:
except Exception:
counter_data = {"last_index": 0}
mode = "fixed"
effective_current_number = current_number
if isinstance(current_number, (list, tuple)) and len(current_number) >= 2:
effective_current_number = current_number[0]
mode = current_number[1]
elif isinstance(current_number, dict) and 'mode' in current_number:
mode = current_number['mode']
effective_current_number = current_number.get('value', 0)
if mode == "increment":
old_index = counter_data.get("last_index", effective_current_number)
new_index = (old_index + 1) % len(txt_files)
effective_index = new_index
counter_data["last_index"] = new_index
print(f"[BatchPromptLoader] Increment: {old_index} → {new_index} ({txt_files[new_index][:30]}...)")
print(f"[BatchPromptLoader] Increment: {old_index} -> {new_index} "
f"({txt_files[new_index][:30]}...)")
elif mode == "decrement":
old_index = counter_data.get("last_index", effective_current_number)
new_index = (old_index - 1) % len(txt_files)
effective_index = new_index
counter_data["last_index"] = new_index
print(f"[BatchPromptLoader] Decrement: {old_index} → {new_index} ({txt_files[new_index][:30]}...)")
print(f"[BatchPromptLoader] Decrement: {old_index} -> {new_index} "
f"({txt_files[new_index][:30]}...)")
elif mode == "random":
effective_index = random_module.randint(0, len(txt_files) - 1)
counter_data["last_index"] = effective_index
print(f"[BatchPromptLoader] Random: {effective_index} ({txt_files[effective_index][:30]}...)")
print(f"[BatchPromptLoader] Random: {effective_index} "
f"({txt_files[effective_index][:30]}...)")
else:
effective_index = effective_current_number % len(txt_files)
counter_data["last_index"] = effective_index
print(f"[BatchPromptLoader] Fixed: {effective_index} ({txt_files[effective_index][:30]}...)")
print(f"[BatchPromptLoader] Fixed: {effective_index} "
f"({txt_files[effective_index][:30]}...)")
try:
os.makedirs(os.path.dirname(counter_file), exist_ok=True)
with open(counter_file, 'w', encoding='utf-8') as f:
json.dump(counter_data, f, ensure_ascii=False)
except Exception as e:
print(f"[BatchPromptLoader] Warning: Failed to save counter: {e}")
selected_file = txt_files[effective_index]
file_path = os.path.join(target_dir, selected_file)
# 优化 4: 跳过已存在的图片
# Skip images that already exist in the output folder
if skip_exists and output_folder.strip():
out_base = os.path.join(base_dir, output_folder.rstrip('/').rstrip('\\'))
out_base = _confine_to_base(base_dir, output_folder.rstrip('/').rstrip('\\'), "output_folder")
base_name = os.path.splitext(os.path.basename(selected_file))[0]
possible_extensions = ['.png', '.jpg', '.jpeg', '.webp']
file_exists = False
for ext in possible_extensions:
check_path = os.path.join(out_base, f"{base_name}{ext}")
if os.path.exists(check_path):
file_exists = True
print(f"[BatchPromptLoader] ⏭️ 跳过已存在：{base_name}{ext}")
print(f"[BatchPromptLoader] Skipped existing: {base_name}{ext}")
break
if file_exists:
print(f"[BatchPromptLoader] 已跳过：{selected_file}")
if enable_logging:
_write_log(log_dir, os.path.normpath(target_dir), selected_file, effective_index, len(txt_files))
_write_log(log_dir, os.path.normpath(target_dir), selected_file,
effective_index, len(txt_files))
empty_cond = [[torch.zeros(1, 1, 2816), {"pooled_output": torch.zeros(1, 2816)}]]
return (empty_cond, f"SKIPPED:{selected_file}", effective_index, len(txt_files))
# 优化 1: 显示预览信息
preview_msg = f"✓ [{effective_index + 1}/{len(txt_files)}] {selected_file}"
preview_msg = f"[{effective_index + 1}/{len(txt_files)}] {selected_file}"
print(f"[BatchPromptLoader] {preview_msg}")
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
positive = ""
negative = ""
section = None
for line in content.split('\n'):
line = line.strip()
if not line:
continue
lower_line = line.lower()
if lower_line.startswith('positive:'):
section = 'positive'
@@ -332,47 +365,46 @@ class BatchPromptReaderWithClip:
positive += (' ' if positive else '') + line
elif section == 'negative':
negative += (' ' if negative else '') + line
if not positive and content.strip():
positive = content.strip()
if not negative:
negative = ""
with torch.no_grad():
tokens_positive = clip.tokenize(positive)
tokens_negative = clip.tokenize(negative)
cond_positive, pooled_positive = clip.encode_from_tokens(tokens_positive, return_pooled=True)
cond_negative, pooled_negative = clip.encode_from_tokens(tokens_negative, return_pooled=True)
if cond_positive.dim() == 2:
cond_positive = cond_positive.unsqueeze(0)
if cond_negative.dim() == 2:
cond_negative = cond_negative.unsqueeze(0)
if pooled_positive is None:
pooled_positive = torch.zeros(1, 2816)
elif pooled_positive.dim() == 1:
pooled_positive = pooled_positive.unsqueeze(0)
if pooled_negative is None:
pooled_negative = torch.zeros(1, 2816)
elif pooled_negative.dim() == 1:
pooled_negative = pooled_negative.unsqueeze(0)
conditioning = [[cond_positive, {"pooled_output": pooled_positive}]]
# 优化 6: 记录到外置日志（成功编码后标记为已处理）
if enable_logging:
_write_log(log_dir, os.path.normpath(target_dir), selected_file, effective_index, len(txt_files))
# 优化 5: 输出元数据
_write_log(log_dir, os.path.normpath(target_dir), selected_file,
effective_index, len(txt_files))
return (
conditioning,
selected_file, # filename
effective_index, # index (从 0 开始)
effective_index, # index (0-based)
len(txt_files) # total_count
)
except Exception as e:
raise Exception(f"Failed to read file: {selected_file}\nError: {str(e)}")