diff --git a/__pycache__/__init__.cpython-313.pyc b/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..733404c Binary files /dev/null and b/__pycache__/__init__.cpython-313.pyc differ diff --git a/__pycache__/__init__.cpython-314.pyc b/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000..bc1a66b Binary files /dev/null and b/__pycache__/__init__.cpython-314.pyc differ diff --git a/__pycache__/discord_image_node.cpython-313.pyc b/__pycache__/discord_image_node.cpython-313.pyc new file mode 100644 index 0000000..740dd90 Binary files /dev/null and b/__pycache__/discord_image_node.cpython-313.pyc differ diff --git a/__pycache__/discord_image_node.cpython-314.pyc b/__pycache__/discord_image_node.cpython-314.pyc new file mode 100644 index 0000000..3f357c2 Binary files /dev/null and b/__pycache__/discord_image_node.cpython-314.pyc differ diff --git a/__pycache__/discord_video_node.cpython-313.pyc b/__pycache__/discord_video_node.cpython-313.pyc new file mode 100644 index 0000000..9ac4175 Binary files /dev/null and b/__pycache__/discord_video_node.cpython-313.pyc differ diff --git a/__pycache__/discord_video_node.cpython-314.pyc b/__pycache__/discord_video_node.cpython-314.pyc new file mode 100644 index 0000000..cc36910 Binary files /dev/null and b/__pycache__/discord_video_node.cpython-314.pyc differ diff --git a/discord_image_node.py b/discord_image_node.py index bff0062..105e08c 100644 --- a/discord_image_node.py +++ b/discord_image_node.py @@ -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}" \ No newline at end of file + return True \ No newline at end of file diff --git a/discord_video_node.py b/discord_video_node.py index 459e9cb..73b3d4d 100644 --- a/discord_video_node.py +++ b/discord_video_node.py @@ -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): """ diff --git a/images/discord_formatting.png b/images/discord_formatting.png index d1544c2..7806860 100644 Binary files a/images/discord_formatting.png and b/images/discord_formatting.png differ diff --git a/images/discord_webhook_step1.png b/images/discord_webhook_step1.png index a3fb170..c2c2dae 100644 Binary files a/images/discord_webhook_step1.png and b/images/discord_webhook_step1.png differ diff --git a/images/discord_webhook_step2.png b/images/discord_webhook_step2.png index 3123938..77f6265 100644 Binary files a/images/discord_webhook_step2.png and b/images/discord_webhook_step2.png differ diff --git a/images/discord_webhook_step2_2.png b/images/discord_webhook_step2_2.png index 499ccbd..ce20f25 100644 Binary files a/images/discord_webhook_step2_2.png and b/images/discord_webhook_step2_2.png differ diff --git a/images/discord_webhook_step4.png b/images/discord_webhook_step4.png index 56c4bfa..76ed8e9 100644 Binary files a/images/discord_webhook_step4.png and b/images/discord_webhook_step4.png differ diff --git a/images/nodes_example_3.png b/images/nodes_example_3.png index 86c0734..548a250 100644 Binary files a/images/nodes_example_3.png and b/images/nodes_example_3.png differ diff --git a/pyproject.toml b/pyproject.toml index 812ee13..e5a83bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..80da0e7 --- /dev/null +++ b/utils/__init__.py @@ -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', +] diff --git a/utils/__pycache__/__init__.cpython-314.pyc b/utils/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000..834cb7f Binary files /dev/null and b/utils/__pycache__/__init__.cpython-314.pyc differ diff --git a/utils/__pycache__/discord_api.cpython-314.pyc b/utils/__pycache__/discord_api.cpython-314.pyc new file mode 100644 index 0000000..06e5629 Binary files /dev/null and b/utils/__pycache__/discord_api.cpython-314.pyc differ diff --git a/utils/__pycache__/github_integration.cpython-314.pyc b/utils/__pycache__/github_integration.cpython-314.pyc new file mode 100644 index 0000000..20fe7e9 Binary files /dev/null and b/utils/__pycache__/github_integration.cpython-314.pyc differ diff --git a/utils/__pycache__/prompt_extractor.cpython-314.pyc b/utils/__pycache__/prompt_extractor.cpython-314.pyc new file mode 100644 index 0000000..0b88602 Binary files /dev/null and b/utils/__pycache__/prompt_extractor.cpython-314.pyc differ diff --git a/utils/__pycache__/sanitizer.cpython-314.pyc b/utils/__pycache__/sanitizer.cpython-314.pyc new file mode 100644 index 0000000..32596a3 Binary files /dev/null and b/utils/__pycache__/sanitizer.cpython-314.pyc differ diff --git a/utils/discord_api.py b/utils/discord_api.py new file mode 100644 index 0000000..24dc813 --- /dev/null +++ b/utils/discord_api.py @@ -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" diff --git a/utils/github_integration.py b/utils/github_integration.py new file mode 100644 index 0000000..879216e --- /dev/null +++ b/utils/github_integration.py @@ -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}" diff --git a/utils/prompt_extractor.py b/utils/prompt_extractor.py new file mode 100644 index 0000000..71ffb6c --- /dev/null +++ b/utils/prompt_extractor.py @@ -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 diff --git a/utils/sanitizer.py b/utils/sanitizer.py new file mode 100644 index 0000000..ea2d0d0 --- /dev/null +++ b/utils/sanitizer.py @@ -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