Refactor code into modular utilities and bump version to 1.1.0

This commit is contained in:
AEmotionStudio
2026-01-10 16:24:58 -08:00
parent 981df1c7b3
commit a9235af50e
25 changed files with 988 additions and 945 deletions
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+5 -615
View File
@@ -1,5 +1,6 @@
"""ComfyUI node for sending images to Discord and saving them locally."""
import os
import datetime
import json
import time
import numpy as np
@@ -14,463 +15,17 @@ import requests
from io import BytesIO
from uuid import uuid4
from typing import Any, Union, List, Optional
from pathlib import Path
import base64
class SaveImage:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
self.type = "output"
self.prefix_append = ""
self.compress_level = 4
# Import shared utilities
from utils import sanitize_json_for_export, update_github_cdn_urls, extract_prompts_from_workflow
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", {"tooltip": "The images to save."}),
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save."})
},
"hidden": {
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
},
}
RETURN_TYPES = ()
FUNCTION = "save_images"
OUTPUT_NODE = True
CATEGORY = "image"
DESCRIPTION = "Saves the input images to your ComfyUI output directory."
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
filename_prefix += self.prefix_append
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
results = list()
for (batch_number, image) in enumerate(images):
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = None
if not args.disable_metadata:
metadata = PngInfo()
if prompt is not None:
# Final sanitization check before embedding
sanitized_prompt = sanitize_json_for_export(prompt)
metadata.add_text("prompt", json.dumps(sanitized_prompt))
if extra_pnginfo is not None:
# Final sanitization check before embedding
sanitized_extra_pnginfo = sanitize_json_for_export(extra_pnginfo)
for x in sanitized_extra_pnginfo:
if x == "workflow":
# Extra sanitization for workflow data
workflow_data = sanitize_json_for_export(sanitized_extra_pnginfo[x])
metadata.add_text(x, json.dumps(workflow_data))
else:
metadata.add_text(x, json.dumps(sanitized_extra_pnginfo[x]))
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{counter:05}_.png"
img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
results.append({
"filename": file,
"subfolder": subfolder,
"type": self.type
})
counter += 1
return { "ui": { "images": results } }
# Helper function to convert tensor to OpenCV format
def tensor_to_cv(tensor: torch.Tensor) -> np.ndarray:
"""Convert a PyTorch tensor to an OpenCV-compatible numpy array."""
return np.clip(tensor.squeeze().cpu().numpy() * 255, 0, 255).astype(np.uint8)
# Helper function to sanitize JSON data by removing webhook information
def sanitize_json_for_export(json_data):
"""
Remove sensitive webhook data and GitHub tokens from JSON data to protect user security.
Parameters:
json_data: The JSON data object (dict) or string to sanitize
Returns:
The sanitized JSON data with webhook information and GitHub tokens removed
"""
if json_data is None:
return None
# Convert string to dict if necessary
if isinstance(json_data, str):
try:
data = json.loads(json_data)
except json.JSONDecodeError:
# If not valid JSON, return as is
return json_data
else:
data = json.loads(json.dumps(json_data))
if not isinstance(data, dict):
return json_data
# Create a deep copy to avoid modifying the original
sanitized_data = json.loads(json.dumps(data))
# Sanitize the workflow data (remove webhook URLs and GitHub tokens)
if "nodes" in sanitized_data:
# Handle the case where nodes is a list (not a dictionary)
if isinstance(sanitized_data["nodes"], list):
for node in sanitized_data["nodes"]:
if isinstance(node, dict):
# Check for webhook_url and github_token in inputs
if "inputs" in node and isinstance(node["inputs"], dict):
inputs = node["inputs"]
if "webhook_url" in inputs:
inputs["webhook_url"] = ""
print("Removed webhook URL from workflow JSON for security")
if "github_token" in inputs:
inputs["github_token"] = ""
print("Removed GitHub token from workflow JSON for security")
# If inputs has any fields that are dictionaries, check them too
for input_name, input_value in inputs.items():
if isinstance(input_value, dict):
if "webhook_url" in input_value:
input_value["webhook_url"] = ""
print(f"Removed nested webhook URL from {input_name} in workflow JSON")
if "github_token" in input_value:
input_value["github_token"] = ""
print(f"Removed nested GitHub token from {input_name} in workflow JSON")
# Check for webhook URL or GitHub token in widgets_values array
if "widgets_values" in node and isinstance(node["widgets_values"], list):
# First method: check for Discord/GitHub node by type
is_sensitive_node = False
if "type" in node and isinstance(node["type"], str):
# Check for any node that might contain Discord, webhook, or GitHub in its name
if ("Discord" in node["type"] or "discord" in node["type"] or
"webhook" in node["type"].lower() or "github" in node["type"].lower()):
is_sensitive_node = True
# Thoroughly scan all widget values for sensitive data regardless of node type
for i, value in enumerate(node["widgets_values"]):
# Check for sensitive data in any string value
if isinstance(value, str):
# Look for Discord webhook URLs
if "discord.com/api/webhooks" in value:
node["widgets_values"][i] = ""
print(f"Removed webhook URL from widgets_values[{i}] for security")
# Check for other webhook patterns
elif value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
node["widgets_values"][i] = ""
print(f"Removed potential webhook URL from widgets_values[{i}] for security")
# Check for GitHub tokens
elif (value.startswith("ghp_") or # GitHub personal access token
value.startswith("github_pat_") or # GitHub personal access token
value.startswith("gho_") or # GitHub OAuth token
value.startswith("ghs_") or # GitHub service token
value.startswith("ghu_")): # GitHub user-to-server token
node["widgets_values"][i] = ""
print(f"Removed GitHub token from widgets_values[{i}] for security")
# Check for generic tokens that could be GitHub tokens
elif len(value) >= 40 and "github" in node.get("type", "").lower() and any(c.isalnum() for c in value):
node["widgets_values"][i] = ""
print(f"Removed potential GitHub token from widgets_values[{i}] for security")
# Handle the case where nodes is a dictionary (node_id -> node)
elif isinstance(sanitized_data["nodes"], dict):
for node_id, node in sanitized_data["nodes"].items():
if isinstance(node, dict):
# Check for webhook_url and github_token in inputs
if "inputs" in node and isinstance(node["inputs"], dict):
inputs = node["inputs"]
if "webhook_url" in inputs:
inputs["webhook_url"] = ""
print("Removed webhook URL from workflow JSON for security")
if "github_token" in inputs:
inputs["github_token"] = ""
print("Removed GitHub token from workflow JSON for security")
# If inputs has any fields that are dictionaries, check them too
for input_name, input_value in inputs.items():
if isinstance(input_value, dict):
if "webhook_url" in input_value:
input_value["webhook_url"] = ""
print(f"Removed nested webhook URL from {input_name} in workflow JSON")
if "github_token" in input_value:
input_value["github_token"] = ""
print(f"Removed nested GitHub token from {input_name} in workflow JSON")
# Check for webhook URL or GitHub token in widgets_values array
if "widgets_values" in node and isinstance(node["widgets_values"], list):
# First method: check for Discord/GitHub node by type
is_sensitive_node = False
if "type" in node and isinstance(node["type"], str):
# Check for any node that might contain Discord, webhook, or GitHub in its name
if ("Discord" in node["type"] or "discord" in node["type"] or
"webhook" in node["type"].lower() or "github" in node["type"].lower()):
is_sensitive_node = True
# Thoroughly scan all widget values for sensitive data regardless of node type
for i, value in enumerate(node["widgets_values"]):
# Check for sensitive data in any string value
if isinstance(value, str):
# Look for Discord webhook URLs
if "discord.com/api/webhooks" in value:
node["widgets_values"][i] = ""
print(f"Removed webhook URL from widgets_values[{i}] for security")
# Check for other webhook patterns
elif value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
node["widgets_values"][i] = ""
print(f"Removed potential webhook URL from widgets_values[{i}] for security")
# Check for GitHub tokens
elif (value.startswith("ghp_") or # GitHub personal access token
value.startswith("github_pat_") or # GitHub personal access token
value.startswith("gho_") or # GitHub OAuth token
value.startswith("ghs_") or # GitHub service token
value.startswith("ghu_")): # GitHub user-to-server token
node["widgets_values"][i] = ""
print(f"Removed GitHub token from widgets_values[{i}] for security")
# Check for generic tokens that could be GitHub tokens
elif len(value) >= 40 and "github" in node.get("type", "").lower() and any(c.isalnum() for c in value):
node["widgets_values"][i] = ""
print(f"Removed potential GitHub token from widgets_values[{i}] for security")
# Also handle extra_pnginfo format where there might be direct sensitive data
if "webhook_url" in sanitized_data:
sanitized_data["webhook_url"] = ""
print("Removed top-level webhook URL from JSON data")
if "github_token" in sanitized_data:
sanitized_data["github_token"] = ""
print("Removed top-level GitHub token from JSON data")
# Recursively check for sensitive data in nested dictionaries
def check_nested_dict(d):
if isinstance(d, dict):
# Check for direct sensitive keys
if "webhook_url" in d:
d["webhook_url"] = ""
print("Removed nested webhook URL from JSON structure")
if "github_token" in d:
d["github_token"] = ""
print("Removed nested GitHub token from JSON structure")
# Recursively check all key-value pairs
for k, v in d.items():
if isinstance(v, (dict, list)):
check_nested_dict(v)
# Check if any string value looks like a GitHub token
elif isinstance(v, str):
if (v.startswith("ghp_") or v.startswith("github_pat_") or
v.startswith("gho_") or v.startswith("ghs_") or v.startswith("ghu_")):
d[k] = ""
print(f"Removed a GitHub token from field '{k}' in JSON structure")
elif isinstance(d, list):
for i, item in enumerate(d):
if isinstance(item, (dict, list)):
check_nested_dict(item)
# Check for webhook URLs or GitHub tokens in string items of lists
elif isinstance(item, str):
if "discord.com/api/webhooks" in item:
# Can't modify the string directly, but this will at least print a warning
print("Warning: Found webhook URL in a list item that cannot be directly sanitized")
elif (item.startswith("ghp_") or item.startswith("github_pat_") or
item.startswith("gho_") or item.startswith("ghs_") or item.startswith("ghu_")):
print("Warning: Found GitHub token in a list item that cannot be directly sanitized")
# Apply the recursive check
check_nested_dict(sanitized_data)
return sanitized_data
# Helper function to extract prompts from workflow data
def extract_prompts_from_workflow(workflow_data):
"""
Extract positive and negative prompts from workflow data.
Parameters:
workflow_data: The workflow data dictionary or object
Returns:
A tuple of (positive_prompt, negative_prompt) or (None, None) if not found
"""
print("extract_prompts_from_workflow called with workflow data")
if workflow_data is None:
print("extract_prompts_from_workflow: workflow_data is None")
return None, None
# Convert string to dict if necessary
if isinstance(workflow_data, str):
try:
data = json.loads(workflow_data)
print("extract_prompts_from_workflow: converted string to JSON")
except json.JSONDecodeError:
print("extract_prompts_from_workflow: failed to decode JSON string")
return None, None
else:
data = workflow_data
if not isinstance(data, dict):
print(f"extract_prompts_from_workflow: data is not a dict, but {type(data)}")
return None, None
positive_prompt = None
negative_prompt = None
# Basic approach: find CLIPTextEncode nodes and extract their text
if "nodes" in data:
nodes = data["nodes"]
print(f"extract_prompts_from_workflow: Found nodes key with {len(nodes)} items")
# Handle list-based nodes structure
if isinstance(nodes, list):
clip_text_encode_nodes = []
# Collect all CLIP nodes
for node in nodes:
if isinstance(node, dict) and "type" in node and node["type"] == "CLIPTextEncode":
if "widgets_values" in node and isinstance(node["widgets_values"], list) and len(node["widgets_values"]) > 0:
clip_text_encode_nodes.append(node)
print(f"extract_prompts_from_workflow: Found {len(clip_text_encode_nodes)} CLIP nodes in list structure")
# If we have exactly 2 CLIP nodes, determine which is which
if len(clip_text_encode_nodes) == 2:
# Try to determine by examining the text content first
# Negative prompts often contain terms like "bad quality", "deformed", etc.
negative_indicators = ["bad quality", "deformed", "blurry", "low quality", "worst quality",
"ugly", "disfigured", "low res", "deformed", "poorly drawn", "mutation"]
for node in clip_text_encode_nodes:
prompt_text = node["widgets_values"][0].lower()
# Check if this prompt contains negative indicators
matches = sum(1 for indicator in negative_indicators if indicator in prompt_text)
if matches >= 3: # If we have multiple matches, likely negative
negative_prompt = node["widgets_values"][0]
else:
# Assume it's positive if not strongly negative
positive_prompt = node["widgets_values"][0]
# If we couldn't determine based on content, try with connections if links exist
if (positive_prompt is None or negative_prompt is None) and "links" in data:
links = data["links"]
samplers = []
# Find all KSampler nodes first
for node in nodes:
if isinstance(node, dict) and "type" in node and "KSampler" in node["type"]:
samplers.append(node)
# If we found samplers, try to trace connections
if samplers and isinstance(links, list):
for link in links:
if len(link) >= 4:
from_node_id = link[0]
to_node_id = link[2]
to_slot = link[3]
# Check if this link connects a CLIP node to a sampler
for clip_node in clip_text_encode_nodes:
if clip_node.get("id") == from_node_id:
for sampler in samplers:
if sampler.get("id") == to_node_id:
# Check which input slot this connects to
if "inputs" in sampler and isinstance(sampler["inputs"], dict):
input_keys = list(sampler["inputs"].keys())
if to_slot < len(input_keys):
input_name = input_keys[to_slot]
if "positive" in input_name.lower():
positive_prompt = clip_node["widgets_values"][0]
elif "negative" in input_name.lower():
negative_prompt = clip_node["widgets_values"][0]
# If we still couldn't determine, go with our defaults
if positive_prompt is None and negative_prompt is None:
# Simply swap them (assuming first is negative, second is positive)
negative_prompt = clip_text_encode_nodes[0]["widgets_values"][0]
positive_prompt = clip_text_encode_nodes[1]["widgets_values"][0]
elif positive_prompt is None: # Only negative was found
# Find the other one
for node in clip_text_encode_nodes:
if node["widgets_values"][0] != negative_prompt:
positive_prompt = node["widgets_values"][0]
break
elif negative_prompt is None: # Only positive was found
# Find the other one
for node in clip_text_encode_nodes:
if node["widgets_values"][0] != positive_prompt:
negative_prompt = node["widgets_values"][0]
break
# If we have only one CLIP node, assume it's positive
elif len(clip_text_encode_nodes) == 1:
positive_prompt = clip_text_encode_nodes[0]["widgets_values"][0]
# Handle dict-based nodes structure
elif isinstance(nodes, dict):
clip_text_encode_nodes = []
# Find all CLIPTextEncode nodes
for node_id, node in nodes.items():
if isinstance(node, dict) and "type" in node and node["type"] == "CLIPTextEncode":
if "widgets_values" in node and isinstance(node["widgets_values"], list) and len(node["widgets_values"]) > 0:
# Store ID for reference
node_with_id = node.copy()
node_with_id["id"] = node_id
clip_text_encode_nodes.append(node_with_id)
# Same process as above, but for dict-based structure
if len(clip_text_encode_nodes) == 2:
# Try to determine by examining the text content first
negative_indicators = ["bad quality", "deformed", "blurry", "low quality", "worst quality",
"ugly", "disfigured", "low res", "deformed", "poorly drawn", "mutation"]
for node in clip_text_encode_nodes:
prompt_text = node["widgets_values"][0].lower()
# Check if this prompt contains negative indicators
matches = sum(1 for indicator in negative_indicators if indicator in prompt_text)
if matches >= 3: # If we have multiple matches, likely negative
negative_prompt = node["widgets_values"][0]
else:
# Assume it's positive if not strongly negative
positive_prompt = node["widgets_values"][0]
# If we couldn't determine by content
if positive_prompt is None and negative_prompt is None:
# Default if we can't determine
negative_prompt = clip_text_encode_nodes[0]["widgets_values"][0]
positive_prompt = clip_text_encode_nodes[1]["widgets_values"][0]
elif positive_prompt is None: # Only negative was found
# Find the other one
for node in clip_text_encode_nodes:
if node["widgets_values"][0] != negative_prompt:
positive_prompt = node["widgets_values"][0]
break
elif negative_prompt is None: # Only positive was found
# Find the other one
for node in clip_text_encode_nodes:
if node["widgets_values"][0] != positive_prompt:
negative_prompt = node["widgets_values"][0]
break
# If we have only one CLIP node, assume it's positive
elif len(clip_text_encode_nodes) == 1:
positive_prompt = clip_text_encode_nodes[0]["widgets_values"][0]
# If we've completed all our detection logic and still can't find a negative prompt
# but we have a positive prompt, assume there's no negative prompt
if positive_prompt is not None and negative_prompt is None:
negative_prompt = ""
print(f"extract_prompts_from_workflow returning: positive={positive_prompt is not None}, negative={negative_prompt is not None}")
# Return the prompts in the correct order
return positive_prompt, negative_prompt
class DiscordSendSaveImage:
"""
@@ -1424,169 +979,4 @@ class DiscordSendSaveImage:
include_prompts_in_message=False, include_format_in_message=False, group_batched_images=True,
send_workflow_json=False, save_cdn_urls=False, github_cdn_update=False, github_repo="",
github_token="", github_file_path="cdn_urls.md", prompt=None, extra_pnginfo=None):
return True
# Add function to send CDN URLs to GitHub repository
def update_github_cdn_urls(github_repo, github_token, file_path, cdn_urls, commit_message=None):
"""
Update a file in a GitHub repository with Discord CDN URLs.
Parameters:
github_repo: The GitHub repository (format: username/repo)
github_token: The GitHub personal access token for authentication
file_path: The path to the file within the repository to update
cdn_urls: List of (filename, url) tuples containing Discord CDN URLs
commit_message: Optional commit message, defaults to a standard message
Returns:
Tuple of (success, message) where success is a boolean and message is a status message
"""
print(f"update_github_cdn_urls called with repo: {github_repo}, file_path: {file_path}, URLs count: {len(cdn_urls)}")
# Check required parameters
if not github_repo:
print("Error: GitHub repository name is empty")
return False, "Missing GitHub repository name"
if not github_token:
print("Error: GitHub token is empty")
return False, "Missing GitHub personal access token"
if not file_path:
print("Error: GitHub file path is empty")
return False, "Missing file path in repository"
if not cdn_urls:
print("Error: No CDN URLs provided to update")
return False, "No CDN URLs to update"
# Ensure repository format is valid
if "/" not in github_repo:
print(f"Error: Invalid GitHub repository format: {github_repo}. Expected format: username/repo")
return False, f"Invalid GitHub repository format: {github_repo}. Expected format: username/repo"
# Setup API endpoint for the file
api_url = f"https://api.github.com/repos/{github_repo}/contents/{file_path}"
# Create headers with token but don't log the actual token
headers = {
"Authorization": f"token {github_token}",
"Accept": "application/vnd.github.v3+json"
}
# Log attempt without exposing the token
print(f"Attempting to access GitHub API at: {api_url} with authentication")
try:
# Check if file exists and get its SHA if it does
file_sha = None
try:
print("Checking if file exists on GitHub...")
response = requests.get(api_url, headers=headers)
print(f"GitHub API check response: Status {response.status_code}")
if response.status_code == 200:
file_data = response.json()
file_sha = file_data.get("sha")
print(f"File exists, got SHA: {file_sha[:7]}...")
# Get current content if file exists
current_content = ""
if file_data.get("content"):
current_content = base64.b64decode(file_data["content"]).decode("utf-8")
print(f"Retrieved existing file content ({len(current_content)} bytes)")
elif response.status_code == 404:
print("File doesn't exist yet, will create a new file")
else:
print(f"Unexpected response checking GitHub file: {response.status_code}")
print(f"Response body: {response.text[:200]}...")
return False, f"Error checking GitHub file: {response.status_code} - {response.text}"
except Exception as e:
# Continue with file creation if checking failed
print(f"Warning: Failed to check file existence: {str(e)}")
# Prepare the file content with the CDN URLs
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
# Start with new content template
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
# If we have existing content, try to merge it
if file_sha and 'current_content' in locals() and current_content:
print("Merging with existing content...")
# Extract existing URLs
existing_urls = {}
for line in current_content.splitlines():
if ": https://" in line and "cdn.discordapp.com" in line:
parts = line.split(": ", 1)
if len(parts) == 2:
name_part = parts[0]
if ". " in name_part: # Remove numbering if present
name_part = name_part.split(". ", 1)[1]
existing_urls[name_part] = parts[1]
print(f"Found {len(existing_urls)} existing URLs in the file")
# Add new URLs (don't duplicate filenames)
for filename, url in cdn_urls:
existing_urls[filename] = url
# Format all URLs
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
for i, (filename, url) in enumerate(existing_urls.items(), 1):
new_content += f"{i}. {filename}: {url}\n"
print(f"Final content has {len(existing_urls)} URLs")
else:
# Just add the new URLs
print("Creating new content with just the new URLs")
for i, (filename, url) in enumerate(cdn_urls, 1):
new_content += f"{i}. {filename}: {url}\n"
print(f"New content has {len(cdn_urls)} URLs")
# Set default commit message if not provided
if not commit_message:
commit_message = f"Update Discord CDN URLs - {timestamp}"
# Prepare the request data
data = {
"message": commit_message,
"content": base64.b64encode(new_content.encode("utf-8")).decode("utf-8"),
}
# Add SHA if file exists (for updating instead of creating)
if file_sha:
data["sha"] = file_sha
print(f"Adding SHA to request for updating existing file")
else:
print("Creating new file (no SHA included)")
# Make the request to create/update the file
print(f"Sending PUT request to GitHub API...")
response = requests.put(api_url, headers=headers, json=data)
print(f"GitHub API response: Status {response.status_code}")
if response.status_code in [200, 201]:
print(f"GitHub API success response: {response.text[:200]}...")
return True, f"Successfully updated GitHub file with {len(cdn_urls)} Discord CDN URLs"
else:
print(f"GitHub API error response: {response.text[:200]}...")
return False, f"Error updating GitHub file: {response.status_code} - {response.text}"
except Exception as e:
import traceback
print(f"Exception during GitHub update: {str(e)}")
# Scrub any potential token from error messages before logging
error_message = str(e)
if github_token and github_token in error_message:
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
# Get traceback but ensure it doesn't contain the token
tb = traceback.format_exc()
if github_token and github_token in tb:
tb = tb.replace(github_token, "[REDACTED_TOKEN]")
print(f"Traceback: {tb}")
return False, f"Exception during GitHub update: {error_message}"
return True
+4 -329
View File
@@ -1,5 +1,6 @@
"""ComfyUI node for sending videos to Discord and saving them locally."""
import os
import datetime
import json
import time
import numpy as np
@@ -20,173 +21,9 @@ import subprocess
import itertools
import functools
import server
import base64
# Helper function to send CDN URLs to GitHub repository
def update_github_cdn_urls(github_repo, github_token, file_path, cdn_urls, commit_message=None):
"""
Update a file in a GitHub repository with Discord CDN URLs.
Parameters:
github_repo: The GitHub repository (format: username/repo)
github_token: The GitHub personal access token for authentication
file_path: The path to the file within the repository to update
cdn_urls: List of (filename, url) tuples containing Discord CDN URLs
commit_message: Optional commit message, defaults to a standard message
Returns:
Tuple of (success, message) where success is a boolean and message is a status message
"""
print(f"update_github_cdn_urls called with repo: {github_repo}, file_path: {file_path}, URLs count: {len(cdn_urls)}")
# Check required parameters
if not github_repo:
print("Error: GitHub repository name is empty")
return False, "Missing GitHub repository name"
if not github_token:
print("Error: GitHub token is empty")
return False, "Missing GitHub personal access token"
if not file_path:
print("Error: GitHub file path is empty")
return False, "Missing file path in repository"
if not cdn_urls:
print("Error: No CDN URLs provided to update")
return False, "No CDN URLs to update"
# Ensure repository format is valid
if "/" not in github_repo:
print(f"Error: Invalid GitHub repository format: {github_repo}. Expected format: username/repo")
return False, f"Invalid GitHub repository format: {github_repo}. Expected format: username/repo"
# Setup API endpoint for the file
api_url = f"https://api.github.com/repos/{github_repo}/contents/{file_path}"
# Create headers with token but don't log the actual token
headers = {
"Authorization": f"token {github_token}",
"Accept": "application/vnd.github.v3+json"
}
# Log attempt without exposing the token
print(f"Attempting to access GitHub API at: {api_url} with authentication")
try:
# Check if file exists and get its SHA if it does
file_sha = None
try:
print("Checking if file exists on GitHub...")
response = requests.get(api_url, headers=headers)
print(f"GitHub API check response: Status {response.status_code}")
if response.status_code == 200:
file_data = response.json()
file_sha = file_data.get("sha")
print(f"File exists, got SHA: {file_sha[:7]}...")
# Get current content if file exists
current_content = ""
if file_data.get("content"):
current_content = base64.b64decode(file_data["content"]).decode("utf-8")
print(f"Retrieved existing file content ({len(current_content)} bytes)")
elif response.status_code == 404:
print("File doesn't exist yet, will create a new file")
else:
print(f"Unexpected response checking GitHub file: {response.status_code}")
print(f"Response body: {response.text[:200]}...")
return False, f"Error checking GitHub file: {response.status_code} - {response.text}"
except Exception as e:
# Continue with file creation if checking failed
print(f"Warning: Failed to check file existence: {str(e)}")
# Prepare the file content with the CDN URLs
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
# Start with new content template
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
# If we have existing content, try to merge it
if file_sha and 'current_content' in locals() and current_content:
print("Merging with existing content...")
# Extract existing URLs
existing_urls = {}
for line in current_content.splitlines():
if ": https://" in line and "cdn.discordapp.com" in line:
parts = line.split(": ", 1)
if len(parts) == 2:
name_part = parts[0]
if ". " in name_part: # Remove numbering if present
name_part = name_part.split(". ", 1)[1]
existing_urls[name_part] = parts[1]
print(f"Found {len(existing_urls)} existing URLs in the file")
# Add new URLs (don't duplicate filenames)
for filename, url in cdn_urls:
existing_urls[filename] = url
# Format all URLs
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
for i, (filename, url) in enumerate(existing_urls.items(), 1):
new_content += f"{i}. {filename}: {url}\n"
print(f"Final content has {len(existing_urls)} URLs")
else:
# Just add the new URLs
print("Creating new content with just the new URLs")
for i, (filename, url) in enumerate(cdn_urls, 1):
new_content += f"{i}. {filename}: {url}\n"
print(f"New content has {len(cdn_urls)} URLs")
# Set default commit message if not provided
if not commit_message:
commit_message = f"Update Discord CDN URLs - {timestamp}"
# Prepare the request data
data = {
"message": commit_message,
"content": base64.b64encode(new_content.encode("utf-8")).decode("utf-8"),
}
# Add SHA if file exists (for updating instead of creating)
if file_sha:
data["sha"] = file_sha
print(f"Adding SHA to request for updating existing file")
else:
print("Creating new file (no SHA included)")
# Make the request to create/update the file
print(f"Sending PUT request to GitHub API...")
response = requests.put(api_url, headers=headers, json=data)
print(f"GitHub API response: Status {response.status_code}")
if response.status_code in [200, 201]:
print(f"GitHub API success response: {response.text[:200]}...")
return True, f"Successfully updated GitHub file with {len(cdn_urls)} Discord CDN URLs"
else:
print(f"GitHub API error response: {response.text[:200]}...")
return False, f"Error updating GitHub file: {response.status_code} - {response.text}"
except Exception as e:
import traceback
print(f"Exception during GitHub update: {str(e)}")
# Scrub any potential token from error messages before logging
error_message = str(e)
if github_token and github_token in error_message:
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
# Get traceback but ensure it doesn't contain the token
tb = traceback.format_exc()
if github_token and github_token in tb:
tb = tb.replace(github_token, "[REDACTED_TOKEN]")
print(f"Traceback: {tb}")
return False, f"Exception during GitHub update: {error_message}"
# Import shared utilities
from utils import sanitize_json_for_export, update_github_cdn_urls
# Define cached decorator for local use
def cached(max_size=None):
"""
@@ -253,168 +90,6 @@ if ffmpeg_path is None:
except Exception as e:
print(f"Error during ffmpeg detection: {str(e)}")
# Improved sanitization function for webhook URLs and GitHub tokens
def sanitize_json_for_export(json_data):
"""
Enhanced sanitization to remove webhook URLs and GitHub tokens for security in exported data.
Thoroughly scans and removes any Discord webhook URLs, GitHub tokens, or similar sensitive data
from JSON structures, making them safe for sharing and exporting.
Parameters:
json_data: The JSON data object (dict, list, or string) to sanitize
Returns:
The sanitized JSON data with webhook information and GitHub tokens removed
"""
if json_data is None:
return None
# Convert string to dict if necessary
if isinstance(json_data, str):
try:
data = json.loads(json_data)
result = sanitize_json_for_export(data)
return json.dumps(result)
except json.JSONDecodeError:
# Check if it's a webhook URL in string format
if ("discord.com/api/webhooks" in json_data or
(json_data.startswith("http") and ("webhook" in json_data.lower() or "discord" in json_data.lower()))):
return ""
# Check if it's a GitHub token in string format
elif (json_data.startswith("ghp_") or # GitHub personal access token
json_data.startswith("github_pat_") or # GitHub personal access token
json_data.startswith("gho_") or # GitHub OAuth token
json_data.startswith("ghs_") or # GitHub service token
json_data.startswith("ghu_")): # GitHub user-to-server token
print("Sanitized GitHub token from JSON string")
return ""
return json_data
if isinstance(json_data, dict):
result = {}
for key, value in json_data.items():
# Handle sensitive keys
sanitized_key = key
if isinstance(key, str):
if "webhook" in key.lower():
sanitized_value = ""
print("Sanitized webhook URL from JSON key")
elif key == "github_token" or "github" in key.lower() and "token" in key.lower():
sanitized_value = ""
print("Sanitized GitHub token from JSON key")
else:
sanitized_value = sanitize_json_for_export(value)
else:
sanitized_value = sanitize_json_for_export(value)
result[sanitized_key] = sanitized_value
# Special handling for ComfyUI workflow structure
if "nodes" in result:
nodes = result["nodes"]
if isinstance(nodes, dict):
for node_id, node in nodes.items():
if isinstance(node, dict):
# Check inputs for sensitive data
if "inputs" in node and isinstance(node["inputs"], dict):
for input_key, input_val in node["inputs"].items():
if input_key == "webhook_url":
node["inputs"][input_key] = ""
print("Sanitized webhook URL from node inputs")
elif input_key == "github_token":
node["inputs"][input_key] = ""
print("Sanitized GitHub token from node inputs")
# Check widgets_values for sensitive data
if "widgets_values" in node and isinstance(node["widgets_values"], list):
# Check if this might be a Discord or GitHub node by its type
is_sensitive_node = False
if "type" in node and isinstance(node["type"], str):
if ("Discord" in node["type"] or "discord" in node["type"] or
"webhook" in node["type"].lower() or "github" in node["type"].lower()):
is_sensitive_node = True
# Check all widget values for sensitive data
for i, value in enumerate(node["widgets_values"]):
if isinstance(value, str):
# Check for Discord webhook URLs
if "discord.com/api/webhooks" in value:
node["widgets_values"][i] = ""
print(f"Sanitized webhook URL from widgets_values[{i}]")
elif value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
node["widgets_values"][i] = ""
print(f"Sanitized potential webhook URL from widgets_values[{i}]")
# Check for GitHub tokens
elif (value.startswith("ghp_") or # GitHub personal access token
value.startswith("github_pat_") or # GitHub personal access token
value.startswith("gho_") or # GitHub OAuth token
value.startswith("ghs_") or # GitHub service token
value.startswith("ghu_")): # GitHub user-to-server token
node["widgets_values"][i] = ""
print(f"Sanitized GitHub token from widgets_values[{i}]")
# Check for generic tokens in GitHub-related nodes
elif len(value) >= 40 and "github" in node.get("type", "").lower() and any(c.isalnum() for c in value):
node["widgets_values"][i] = ""
print(f"Sanitized potential GitHub token from widgets_values[{i}]")
elif isinstance(nodes, list):
for node in nodes:
if isinstance(node, dict):
# Check inputs for sensitive data
if "inputs" in node and isinstance(node["inputs"], dict):
for input_key, input_val in node["inputs"].items():
if input_key == "webhook_url":
node["inputs"][input_key] = ""
print("Sanitized webhook URL from node inputs")
elif input_key == "github_token":
node["inputs"][input_key] = ""
print("Sanitized GitHub token from node inputs")
# Also check widgets_values for sensitive data
if "widgets_values" in node and isinstance(node["widgets_values"], list):
# Check if this might be a Discord or GitHub node by its type
is_sensitive_node = False
if "type" in node and isinstance(node["type"], str):
if ("Discord" in node["type"] or "discord" in node["type"] or
"webhook" in node["type"].lower() or "github" in node["type"].lower()):
is_sensitive_node = True
# Check all widget values for sensitive data
for i, value in enumerate(node["widgets_values"]):
if isinstance(value, str):
# Check for Discord webhook URLs
if "discord.com/api/webhooks" in value:
node["widgets_values"][i] = ""
print(f"Sanitized webhook URL from widgets_values[{i}]")
elif value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
node["widgets_values"][i] = ""
print(f"Sanitized potential webhook URL from widgets_values[{i}]")
# Check for GitHub tokens
elif (value.startswith("ghp_") or # GitHub personal access token
value.startswith("github_pat_") or # GitHub personal access token
value.startswith("gho_") or # GitHub OAuth token
value.startswith("ghs_") or # GitHub service token
value.startswith("ghu_")): # GitHub user-to-server token
node["widgets_values"][i] = ""
print(f"Sanitized GitHub token from widgets_values[{i}]")
# Check for generic tokens in GitHub-related nodes
elif len(value) >= 40 and "github" in node.get("type", "").lower() and any(c.isalnum() for c in value):
node["widgets_values"][i] = ""
print(f"Sanitized potential GitHub token from widgets_values[{i}]")
# Check for top-level sensitive fields
if "webhook_url" in result:
result["webhook_url"] = ""
print("Sanitized top-level webhook URL")
if "github_token" in result:
result["github_token"] = ""
print("Sanitized top-level GitHub token")
return result
elif isinstance(json_data, list):
return [sanitize_json_for_export(item) for item in json_data]
# Function to validate video files for Discord compatibility
def validate_video_for_discord(file_path):
"""
Binary file not shown.

Before

Width:  |  Height:  |  Size: 417 KiB

After

Width:  |  Height:  |  Size: 131 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 33 KiB

After

Width:  |  Height:  |  Size: 130 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 13 KiB

After

Width:  |  Height:  |  Size: 130 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 14 KiB

After

Width:  |  Height:  |  Size: 130 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 115 KiB

After

Width:  |  Height:  |  Size: 131 B

Binary file not shown.

Before

Width:  |  Height:  |  Size: 174 KiB

After

Width:  |  Height:  |  Size: 131 B

+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui-discordsend"
description = "A ComfyUI extension that enables seamless sharing of AI-generated images and videos directly to Discord."
version = "1.0.0"
version = "1.1.0"
license = {file = "LICENSE"}
dependencies = ["requests>=2.25.0"]
+18
View File
@@ -0,0 +1,18 @@
"""
ComfyUI-DiscordSend Utility Package
Shared utilities for Discord integration, sanitization, and GitHub CDN operations.
"""
from .sanitizer import sanitize_json_for_export
from .github_integration import update_github_cdn_urls
from .prompt_extractor import extract_prompts_from_workflow
from .discord_api import DiscordWebhookClient, validate_webhook_url
__all__ = [
'sanitize_json_for_export',
'update_github_cdn_urls',
'extract_prompts_from_workflow',
'DiscordWebhookClient',
'validate_webhook_url',
]
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+314
View File
@@ -0,0 +1,314 @@
"""
Discord API Utilities for ComfyUI-DiscordSend
Provides a client for interacting with Discord webhooks and validation utilities.
"""
import os
import re
import time
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
import requests
# Discord webhook URL patterns
WEBHOOK_URL_PATTERNS = [
r"https?://(?:www\.)?discord(?:app)?\.com/api/webhooks/\d+/[\w-]+",
r"https?://(?:www\.)?discordapp\.com/api/webhooks/\d+/[\w-]+",
]
def validate_webhook_url(url: str) -> Tuple[bool, str]:
"""
Validate a Discord webhook URL.
Args:
url: The webhook URL to validate
Returns:
Tuple of (is_valid, message)
"""
if not url:
return False, "Webhook URL is empty"
if not url.startswith("http"):
return False, "Webhook URL must start with http:// or https://"
# Check against known patterns
for pattern in WEBHOOK_URL_PATTERNS:
if re.match(pattern, url, re.IGNORECASE):
return True, "Valid Discord webhook URL"
# More lenient check
if "discord" in url.lower() and "webhook" in url.lower():
return True, "Appears to be a Discord webhook URL"
return False, "URL does not appear to be a valid Discord webhook URL"
def sanitize_webhook_for_logging(url: str) -> str:
"""
Sanitize a webhook URL for safe logging (hide the token portion).
Args:
url: The webhook URL
Returns:
Sanitized URL safe for logging
"""
if not url:
return ""
# Pattern: https://discord.com/api/webhooks/{id}/{token}
match = re.match(r"(https?://[^/]+/api/webhooks/\d+/)(.+)", url)
if match:
return f"{match.group(1)}[REDACTED]"
return "[REDACTED_WEBHOOK_URL]"
class DiscordWebhookClient:
"""
Client for sending messages and files to Discord via webhooks.
Features:
- Automatic retry with exponential backoff
- Rate limit handling
- File size validation
- Error handling with sanitized logging
"""
# Discord limits
MAX_FILE_SIZE = 25 * 1024 * 1024 # 25MB for regular users
MAX_FILES_PER_MESSAGE = 10
MAX_MESSAGE_LENGTH = 2000
def __init__(self, webhook_url: str, max_retries: int = 3):
"""
Initialize the Discord webhook client.
Args:
webhook_url: The Discord webhook URL
max_retries: Maximum number of retry attempts for failed requests
"""
self.webhook_url = webhook_url
self.max_retries = max_retries
self._validated = False
def validate(self) -> Tuple[bool, str]:
"""Validate the webhook URL."""
is_valid, message = validate_webhook_url(self.webhook_url)
self._validated = is_valid
return is_valid, message
def send_message(
self,
content: str = "",
files: Optional[List[Tuple[str, bytes, str]]] = None,
embeds: Optional[List[Dict]] = None
) -> Tuple[bool, Dict[str, Any]]:
"""
Send a message to Discord.
Args:
content: Text message content
files: List of (filename, file_bytes, content_type) tuples
embeds: List of embed dictionaries
Returns:
Tuple of (success, response_data)
"""
if not self.webhook_url:
return False, {"error": "No webhook URL configured"}
# Validate URL if not already done
if not self._validated:
is_valid, message = self.validate()
if not is_valid:
return False, {"error": message}
# Truncate message if too long
if content and len(content) > self.MAX_MESSAGE_LENGTH:
content = content[:self.MAX_MESSAGE_LENGTH - 3] + "..."
# Prepare request
data = {}
if content:
data["content"] = content
if embeds:
data["embeds"] = embeds
# Prepare files
request_files = None
if files:
request_files = []
for i, (filename, file_bytes, content_type) in enumerate(files):
# Validate file size
if len(file_bytes) > self.MAX_FILE_SIZE:
continue # Skip oversized files
request_files.append(
(f"file{i}", (filename, BytesIO(file_bytes), content_type))
)
if len(request_files) > self.MAX_FILES_PER_MESSAGE:
request_files = request_files[:self.MAX_FILES_PER_MESSAGE]
# Send with retry logic
return self._send_with_retry(data, request_files)
def send_file(
self,
file_path: str,
message: str = "",
additional_files: Optional[List[str]] = None
) -> Tuple[bool, Dict[str, Any]]:
"""
Send a file from disk to Discord.
Args:
file_path: Path to the file to send
message: Optional message to accompany the file
additional_files: Optional list of additional file paths
Returns:
Tuple of (success, response_data)
"""
files = []
# Add main file
if os.path.exists(file_path):
file_size = os.path.getsize(file_path)
if file_size <= self.MAX_FILE_SIZE:
with open(file_path, "rb") as f:
filename = os.path.basename(file_path)
content_type = self._get_content_type(filename)
files.append((filename, f.read(), content_type))
# Add additional files
if additional_files:
for path in additional_files:
if os.path.exists(path):
file_size = os.path.getsize(path)
if file_size <= self.MAX_FILE_SIZE:
with open(path, "rb") as f:
filename = os.path.basename(path)
content_type = self._get_content_type(filename)
files.append((filename, f.read(), content_type))
if not files:
return False, {"error": "No valid files to send"}
return self.send_message(content=message, files=files)
def _send_with_retry(
self,
data: Dict,
files: Optional[List] = None
) -> Tuple[bool, Dict[str, Any]]:
"""Send request with retry logic."""
last_error = None
for attempt in range(self.max_retries):
try:
if files:
response = requests.post(
self.webhook_url,
data={"payload_json": str(data)} if data else None,
files=files,
timeout=60
)
else:
response = requests.post(
self.webhook_url,
json=data,
timeout=30
)
# Handle rate limiting
if response.status_code == 429:
retry_after = response.json().get("retry_after", 1)
time.sleep(retry_after)
continue
# Success
if response.status_code in [200, 204]:
try:
return True, response.json() if response.content else {}
except:
return True, {}
# Client errors (don't retry)
if 400 <= response.status_code < 500:
return False, {
"error": f"Discord API error: {response.status_code}",
"details": response.text[:500]
}
# Server errors (retry)
last_error = f"Discord API returned {response.status_code}"
except requests.exceptions.Timeout:
last_error = "Request timed out"
except requests.exceptions.RequestException as e:
last_error = str(e)
# Exponential backoff
if attempt < self.max_retries - 1:
time.sleep(2 ** attempt)
return False, {"error": last_error or "Unknown error"}
def _get_content_type(self, filename: str) -> str:
"""Get content type based on file extension."""
ext = os.path.splitext(filename)[1].lower()
content_types = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".mp4": "video/mp4",
".webm": "video/webm",
".mov": "video/quicktime",
".json": "application/json",
".txt": "text/plain",
}
return content_types.get(ext, "application/octet-stream")
def validate_file_for_discord(file_path: str) -> Tuple[bool, str]:
"""
Validate that a file is compatible with Discord uploads.
Args:
file_path: Path to the file
Returns:
Tuple of (is_valid, message)
"""
if not os.path.exists(file_path):
return False, f"File does not exist: {file_path}"
file_size = os.path.getsize(file_path)
if file_size == 0:
return False, "File is empty"
if file_size < 1024:
return False, f"File is suspiciously small: {file_size} bytes"
max_size = 25 * 1024 * 1024 # 25MB
if file_size > max_size:
return False, f"File exceeds Discord's size limit: {file_size} bytes (max {max_size} bytes)"
# Check extension
ext = os.path.splitext(file_path)[1].lower().lstrip('.')
supported_formats = ['png', 'jpg', 'jpeg', 'gif', 'webp', 'mp4', 'webm', 'mov']
if ext in supported_formats:
return True, f"Valid {ext.upper()} file for Discord"
else:
return False, f"Format {ext} may not be fully supported by Discord"
+137
View File
@@ -0,0 +1,137 @@
"""
GitHub Integration Utilities for ComfyUI-DiscordSend
Handles updating GitHub repositories with Discord CDN URLs.
"""
import base64
import time
from typing import List, Optional, Tuple
import requests
def update_github_cdn_urls(
github_repo: str,
github_token: str,
file_path: str,
cdn_urls: List[Tuple[str, str]],
commit_message: Optional[str] = None
) -> Tuple[bool, str]:
"""
Update a file in a GitHub repository with Discord CDN URLs.
Args:
github_repo: The GitHub repository (format: username/repo)
github_token: The GitHub personal access token for authentication
file_path: The path to the file within the repository to update
cdn_urls: List of (filename, url) tuples containing Discord CDN URLs
commit_message: Optional commit message, defaults to a standard message
Returns:
Tuple of (success, message) where success is a boolean and message is a status message
"""
# Validate required parameters
if not github_repo:
return False, "Missing GitHub repository name"
if not github_token:
return False, "Missing GitHub personal access token"
if not file_path:
return False, "Missing file path in repository"
if not cdn_urls:
return False, "No CDN URLs to update"
# Ensure repository format is valid
if "/" not in github_repo:
return False, f"Invalid GitHub repository format: {github_repo}. Expected format: username/repo"
# Setup API endpoint
api_url = f"https://api.github.com/repos/{github_repo}/contents/{file_path}"
headers = {
"Authorization": f"token {github_token}",
"Accept": "application/vnd.github.v3+json"
}
try:
# Check if file exists and get its SHA
file_sha = None
current_content = ""
response = requests.get(api_url, headers=headers, timeout=30)
if response.status_code == 200:
file_data = response.json()
file_sha = file_data.get("sha")
# Get current content
if file_data.get("content"):
current_content = base64.b64decode(file_data["content"]).decode("utf-8")
elif response.status_code == 404:
pass # File doesn't exist, will create new
else:
return False, f"Error checking GitHub file: {response.status_code} - {response.text}"
# Prepare file content
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
# Merge with existing content if present
existing_urls = {}
if current_content:
for line in current_content.splitlines():
if ": https://" in line and "cdn.discordapp.com" in line:
parts = line.split(": ", 1)
if len(parts) == 2:
name_part = parts[0]
# Remove numbering if present
if ". " in name_part:
name_part = name_part.split(". ", 1)[1]
existing_urls[name_part] = parts[1]
# Add new URLs (overwrites duplicates)
for filename, url in cdn_urls:
existing_urls[filename] = url
# Format content
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
for i, (filename, url) in enumerate(existing_urls.items(), 1):
new_content += f"{i}. {filename}: {url}\n"
# Set commit message
if not commit_message:
commit_message = f"Update Discord CDN URLs - {timestamp}"
# Prepare request data
data = {
"message": commit_message,
"content": base64.b64encode(new_content.encode("utf-8")).decode("utf-8"),
}
if file_sha:
data["sha"] = file_sha
# Create/update file
response = requests.put(api_url, headers=headers, json=data, timeout=30)
if response.status_code in [200, 201]:
return True, f"Successfully updated GitHub file with {len(cdn_urls)} Discord CDN URLs"
else:
return False, f"Error updating GitHub file: {response.status_code} - {response.text}"
except requests.exceptions.Timeout:
return False, "GitHub API request timed out"
except requests.exceptions.RequestException as e:
# Scrub token from error messages
error_message = str(e)
if github_token and github_token in error_message:
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
return False, f"GitHub API request failed: {error_message}"
except Exception as e:
error_message = str(e)
if github_token and github_token in error_message:
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
return False, f"Exception during GitHub update: {error_message}"
+240
View File
@@ -0,0 +1,240 @@
"""
Prompt Extraction Utilities for ComfyUI-DiscordSend
Extracts positive and negative prompts from ComfyUI workflow data.
"""
import json
from typing import Any, Dict, List, Optional, Tuple, Union
# Common negative prompt indicators
NEGATIVE_INDICATORS = [
"bad quality", "deformed", "blurry", "low quality", "worst quality",
"ugly", "disfigured", "low res", "poorly drawn", "mutation",
"extra limbs", "bad anatomy", "watermark", "text", "signature"
]
def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Optional[str]]:
"""
Extract positive and negative prompts from workflow data.
Analyzes ComfyUI workflow structure to find CLIPTextEncode nodes and
determine which contains the positive vs negative prompt.
Args:
workflow_data: The workflow data dictionary or JSON string
Returns:
A tuple of (positive_prompt, negative_prompt) or (None, None) if not found
"""
if workflow_data is None:
return None, None
# Convert string to dict if necessary
if isinstance(workflow_data, str):
try:
data = json.loads(workflow_data)
except json.JSONDecodeError:
return None, None
else:
data = workflow_data
if not isinstance(data, dict):
return None, None
positive_prompt = None
negative_prompt = None
# Find CLIPTextEncode nodes
if "nodes" not in data:
return None, None
nodes = data["nodes"]
clip_nodes = _find_clip_text_encode_nodes(nodes)
if not clip_nodes:
return None, None
# Determine positive/negative based on content and structure
if len(clip_nodes) == 1:
# Single CLIP node - assume it's the positive prompt
positive_prompt = _get_prompt_text(clip_nodes[0])
elif len(clip_nodes) >= 2:
# Multiple CLIP nodes - need to determine which is which
positive_prompt, negative_prompt = _classify_prompts(clip_nodes, data)
# Ensure we return empty string for negative if we have positive but not negative
if positive_prompt is not None and negative_prompt is None:
negative_prompt = ""
return positive_prompt, negative_prompt
def _find_clip_text_encode_nodes(nodes: Union[List, Dict]) -> List[Dict]:
"""Find all CLIPTextEncode nodes in the workflow."""
clip_nodes = []
if isinstance(nodes, list):
for node in nodes:
if _is_clip_text_encode(node):
clip_nodes.append(node)
elif isinstance(nodes, dict):
for node_id, node in nodes.items():
if _is_clip_text_encode(node):
node_copy = dict(node)
node_copy["id"] = node_id
clip_nodes.append(node_copy)
return clip_nodes
def _is_clip_text_encode(node: Any) -> bool:
"""Check if a node is a CLIPTextEncode node with valid text."""
if not isinstance(node, dict):
return False
if node.get("type") != "CLIPTextEncode":
return False
widgets = node.get("widgets_values", [])
if not isinstance(widgets, list) or len(widgets) == 0:
return False
# Must have a string value at position 0
if not isinstance(widgets[0], str):
return False
return True
def _get_prompt_text(node: Dict) -> Optional[str]:
"""Extract the prompt text from a CLIP node."""
widgets = node.get("widgets_values", [])
if widgets and isinstance(widgets[0], str):
return widgets[0]
return None
def _classify_prompts(clip_nodes: List[Dict], workflow_data: Dict) -> Tuple[Optional[str], Optional[str]]:
"""
Classify which CLIP nodes contain positive vs negative prompts.
Uses multiple heuristics:
1. Content analysis (negative prompts often contain quality-related terms)
2. Connection analysis (traces connections to sampler nodes)
"""
positive_prompt = None
negative_prompt = None
# First pass: content-based classification
for node in clip_nodes:
prompt_text = _get_prompt_text(node)
if prompt_text is None:
continue
prompt_lower = prompt_text.lower()
# Count negative indicators
negative_matches = sum(1 for indicator in NEGATIVE_INDICATORS if indicator in prompt_lower)
if negative_matches >= 3:
# Likely a negative prompt
if negative_prompt is None:
negative_prompt = prompt_text
else:
# Assume positive
if positive_prompt is None:
positive_prompt = prompt_text
# Second pass: connection-based classification (if needed)
if positive_prompt is None or negative_prompt is None:
positive_prompt, negative_prompt = _classify_by_connections(
clip_nodes, workflow_data, positive_prompt, negative_prompt
)
# Fallback: if we still can't determine, use first two nodes
if positive_prompt is None and negative_prompt is None and len(clip_nodes) >= 2:
# Convention: assume first is positive, second is negative
positive_prompt = _get_prompt_text(clip_nodes[0])
negative_prompt = _get_prompt_text(clip_nodes[1])
elif positive_prompt is None and negative_prompt is not None:
# Find the other prompt
for node in clip_nodes:
text = _get_prompt_text(node)
if text != negative_prompt:
positive_prompt = text
break
elif negative_prompt is None and positive_prompt is not None:
# Find the other prompt
for node in clip_nodes:
text = _get_prompt_text(node)
if text != positive_prompt:
negative_prompt = text
break
return positive_prompt, negative_prompt
def _classify_by_connections(
clip_nodes: List[Dict],
workflow_data: Dict,
existing_positive: Optional[str],
existing_negative: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
"""Try to classify prompts by analyzing node connections to samplers."""
positive = existing_positive
negative = existing_negative
links = workflow_data.get("links", [])
if not isinstance(links, list):
return positive, negative
nodes = workflow_data.get("nodes", [])
# Find sampler nodes
samplers = []
if isinstance(nodes, list):
for node in nodes:
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
samplers.append(node)
elif isinstance(nodes, dict):
for node_id, node in nodes.items():
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
node_copy = dict(node)
node_copy["id"] = node_id
samplers.append(node_copy)
if not samplers:
return positive, negative
# Trace connections from CLIP nodes to samplers
for link in links:
if len(link) < 4:
continue
from_node_id = link[0]
to_node_id = link[2]
to_slot = link[3]
# Find matching CLIP node and sampler
for clip_node in clip_nodes:
clip_id = clip_node.get("id")
if clip_id == from_node_id:
for sampler in samplers:
if sampler.get("id") == to_node_id:
# Determine slot type
inputs = sampler.get("inputs", {})
if isinstance(inputs, dict):
input_keys = list(inputs.keys())
if to_slot < len(input_keys):
input_name = input_keys[to_slot].lower()
prompt_text = _get_prompt_text(clip_node)
if "positive" in input_name and positive is None:
positive = prompt_text
elif "negative" in input_name and negative is None:
negative = prompt_text
return positive, negative
+269
View File
@@ -0,0 +1,269 @@
"""
JSON Sanitization Utilities for ComfyUI-DiscordSend
Provides comprehensive sanitization of sensitive data (webhook URLs, GitHub tokens)
from JSON data structures before export or sharing.
"""
import json
import re
from typing import Any, Dict, List, Optional, Union
# Patterns for detecting sensitive data
WEBHOOK_PATTERNS = [
r"discord\.com/api/webhooks",
r"discordapp\.com/api/webhooks",
]
GITHUB_TOKEN_PREFIXES = [
"ghp_", # GitHub personal access token
"github_pat_", # GitHub personal access token (new format)
"gho_", # GitHub OAuth token
"ghs_", # GitHub service token
"ghu_", # GitHub user-to-server token
]
def is_webhook_url(value: str) -> bool:
"""Check if a string appears to be a Discord webhook URL."""
if not isinstance(value, str):
return False
for pattern in WEBHOOK_PATTERNS:
if re.search(pattern, value, re.IGNORECASE):
return True
# Also check for generic webhook patterns in URLs
if value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
return True
return False
def is_github_token(value: str) -> bool:
"""Check if a string appears to be a GitHub token."""
if not isinstance(value, str):
return False
for prefix in GITHUB_TOKEN_PREFIXES:
if value.startswith(prefix):
return True
return False
def is_potential_token(value: str, context_type: str = "") -> bool:
"""Check if a string might be a token based on length and context."""
if not isinstance(value, str):
return False
# Long alphanumeric strings in GitHub-related contexts
if len(value) >= 40 and "github" in context_type.lower():
if any(c.isalnum() for c in value):
return True
return False
def sanitize_string(value: str, context_type: str = "") -> str:
"""
Sanitize a single string value.
Args:
value: The string to sanitize
context_type: Optional context (e.g., node type) to help identify tokens
Returns:
Empty string if sensitive data detected, original value otherwise
"""
if is_webhook_url(value):
return ""
if is_github_token(value):
return ""
if is_potential_token(value, context_type):
return ""
return value
def sanitize_widget_values(widgets: List, node_type: str = "") -> List:
"""
Sanitize a list of widget values from a ComfyUI node.
Args:
widgets: List of widget values
node_type: The type of node for context
Returns:
Sanitized list with sensitive values replaced with empty strings
"""
result = []
for value in widgets:
if isinstance(value, str):
result.append(sanitize_string(value, node_type))
elif isinstance(value, dict):
result.append(sanitize_dict(value))
elif isinstance(value, list):
result.append(sanitize_widget_values(value, node_type))
else:
result.append(value)
return result
def sanitize_node_inputs(inputs: Dict, node_type: str = "") -> Dict:
"""
Sanitize the inputs dictionary of a ComfyUI node.
Args:
inputs: Dictionary of node inputs
node_type: The type of node for context
Returns:
Sanitized dictionary
"""
result = {}
for key, value in inputs.items():
if key in ("webhook_url", "github_token"):
result[key] = ""
elif isinstance(value, str):
result[key] = sanitize_string(value, node_type)
elif isinstance(value, dict):
result[key] = sanitize_node_inputs(value, node_type)
elif isinstance(value, list):
result[key] = sanitize_widget_values(value, node_type)
else:
result[key] = value
return result
def sanitize_node(node: Dict) -> Dict:
"""
Sanitize a single ComfyUI node.
Args:
node: The node dictionary
Returns:
Sanitized node dictionary
"""
if not isinstance(node, dict):
return node
result = dict(node)
node_type = result.get("type", "")
# Sanitize inputs
if "inputs" in result and isinstance(result["inputs"], dict):
result["inputs"] = sanitize_node_inputs(result["inputs"], node_type)
# Sanitize widget values
if "widgets_values" in result and isinstance(result["widgets_values"], list):
result["widgets_values"] = sanitize_widget_values(result["widgets_values"], node_type)
return result
def sanitize_dict(data: Dict) -> Dict:
"""
Recursively sanitize a dictionary.
Args:
data: Dictionary to sanitize
Returns:
Sanitized dictionary
"""
result = {}
for key, value in data.items():
# Handle known sensitive keys
if key in ("webhook_url", "github_token"):
result[key] = ""
continue
# Handle nested structures
if isinstance(value, dict):
result[key] = sanitize_dict(value)
elif isinstance(value, list):
result[key] = sanitize_list(value)
elif isinstance(value, str):
result[key] = sanitize_string(value)
else:
result[key] = value
# Special handling for ComfyUI workflow structure
if "nodes" in result:
nodes = result["nodes"]
if isinstance(nodes, list):
result["nodes"] = [sanitize_node(n) for n in nodes]
elif isinstance(nodes, dict):
result["nodes"] = {k: sanitize_node(v) for k, v in nodes.items()}
return result
def sanitize_list(data: List) -> List:
"""
Recursively sanitize a list.
Args:
data: List to sanitize
Returns:
Sanitized list
"""
result = []
for item in data:
if isinstance(item, dict):
result.append(sanitize_dict(item))
elif isinstance(item, list):
result.append(sanitize_list(item))
elif isinstance(item, str):
result.append(sanitize_string(item))
else:
result.append(item)
return result
def sanitize_json_for_export(json_data: Any) -> Any:
"""
Remove sensitive webhook data and GitHub tokens from JSON data to protect user security.
This is the main entry point for sanitization. It handles:
- Discord webhook URLs (multiple formats)
- GitHub personal access tokens (all variants)
- Nested dictionaries and lists
- ComfyUI workflow structures
Args:
json_data: The JSON data object (dict, list, or string) to sanitize
Returns:
The sanitized JSON data with sensitive information removed
"""
if json_data is None:
return None
# Handle string input (may be JSON string or plain string)
if isinstance(json_data, str):
try:
data = json.loads(json_data)
result = sanitize_json_for_export(data)
return json.dumps(result)
except json.JSONDecodeError:
# Not valid JSON - check if it's sensitive data directly
return sanitize_string(json_data)
# Handle dictionary
if isinstance(json_data, dict):
return sanitize_dict(json_data)
# Handle list
if isinstance(json_data, list):
return sanitize_list(json_data)
# Return other types as-is
return json_data