Refactor code into modular utilities and bump version to 1.1.0
@@ -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
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
|
Before Width: | Height: | Size: 417 KiB After Width: | Height: | Size: 131 B |
|
Before Width: | Height: | Size: 33 KiB After Width: | Height: | Size: 130 B |
|
Before Width: | Height: | Size: 13 KiB After Width: | Height: | Size: 130 B |
|
Before Width: | Height: | Size: 14 KiB After Width: | Height: | Size: 130 B |
|
Before Width: | Height: | Size: 115 KiB After Width: | Height: | Size: 131 B |
|
Before Width: | Height: | Size: 174 KiB After Width: | Height: | Size: 131 B |
@@ -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"]
|
||||
|
||||
|
||||
@@ -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',
|
||||
]
|
||||
@@ -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"
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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
|
||||