feat: secure webhook storage (secrets or env), added .gitignore

This commit is contained in:
maxmad62
2025-10-08 20:24:05 +02:00
parent 1f7a02b8ba
commit f79793558a
2 changed files with 150 additions and 129 deletions
+10
View File
@@ -0,0 +1,10 @@
# Ignore compiled Python files
__pycache__/
*.pyc
# Ignore local test copies
nodes - Copie.py
# Ignore secrets
secrets/
discord_webhook_url.txt
+140 -129
View File
@@ -1,28 +1,13 @@
"""
Copyright (C) 2024 Dayuppy
Discord Webhook nodes (safe version: no URL stored in workflows/PNGs)
This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.
This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.
You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301
USA
@author: Dayuppy
@title: Discord Webhook
@nickname: DiscordWebhook
@description: A very simple Discord webhook integration node for ComfyUI that lets you post images and text.
Base: Dayuppy
Mod: +message input + env/secrets for webhook URL
"""
import os
import io
import json
import tempfile
import shutil
import numpy as np
@@ -30,18 +15,113 @@ import asyncio
from PIL import Image, ImageDraw
from discord_webhook import AsyncDiscordWebhook
# === Secrets handling ===
BASE_DIR = os.path.dirname(__file__)
SECRETS_DIR = os.path.join(BASE_DIR, "secrets")
WEBHOOK_FILE = os.path.join(SECRETS_DIR, "webhook.txt")
LEGACY_FILE = os.path.join(BASE_DIR, "discord_webhook_url.txt") # backward-compat, deprecated
def _read_file(path):
try:
with open(path, "r", encoding="utf-8") as f:
return f.read().strip()
except Exception:
return ""
def get_webhook_url():
# 1) env var (recommended)
env_url = os.getenv("DISCORD_WEBHOOK_URL", "").strip()
if env_url:
return env_url
# 2) local secret file
if os.path.exists(WEBHOOK_FILE):
return _read_file(WEBHOOK_FILE)
# 3) legacy fallback for old setups
if os.path.exists(LEGACY_FILE):
return _read_file(LEGACY_FILE)
return ""
# === Utilities ===
def create_default_image():
"""Create a simple TV test pattern image."""
image = Image.new("RGB", (128, 128), "black")
colors = ["white", "yellow", "cyan", "green", "magenta", "red", "blue", "black"]
bar_width = 128 // len(colors)
draw = ImageDraw.Draw(image)
for i, color in enumerate(colors):
draw.rectangle([i * bar_width, 0, (i + 1) * bar_width, 128], fill=color)
return image
def _tensor_or_nd_to_pil_list(image):
"""
Accepts:
- None -> default image
- numpy ndarray HxWxC or NxHxWxC
- torch tensor-like with .cpu().numpy()
- PIL.Image.Image
Returns: list of PIL images
"""
images_to_send = []
if image is None:
images_to_send.append(create_default_image())
return images_to_send
if isinstance(image, Image.Image):
images_to_send.append(image)
return images_to_send
if isinstance(image, np.ndarray):
if image.ndim == 4:
for i in range(image.shape[0]):
arr = np.clip(image[i] * 255, 0, 255).astype(np.uint8)
images_to_send.append(Image.fromarray(arr))
elif image.ndim == 3:
arr = np.clip(image * 255, 0, 255).astype(np.uint8)
images_to_send.append(Image.fromarray(arr))
else:
raise ValueError("Input numpy array must be 3D or 4D.")
return images_to_send
if hasattr(image, "cpu"): # torch-like
array = image.cpu().numpy()
if array.ndim == 4:
for i in range(array.shape[0]):
arr = np.clip(array[i] * 255, 0, 255).astype(np.uint8)
images_to_send.append(Image.fromarray(arr))
elif array.ndim == 3:
arr = np.clip(array * 255, 0, 255).astype(np.uint8)
images_to_send.append(Image.fromarray(arr))
else:
raise ValueError("Input tensor must be 3D or 4D.")
return images_to_send
# Fallback: try to treat as PIL
images_to_send.append(image)
return images_to_send
def _prepare_files_for_discord(images_to_send):
"""Save PIL images to temp PNGs, enforce <20MB, return list of dicts {data, name}."""
files = []
temp_dir = tempfile.mkdtemp()
try:
for idx, img in enumerate(images_to_send):
file_path = os.path.join(temp_dir, f"image_{idx}.png")
img.save(file_path, format="PNG", compress_level=1)
# if >20MB, resize and recompress
if os.path.getsize(file_path) > 20 * 1024 * 1024:
w, h = img.width, img.height
img_small = img.resize((max(1, w // 2), max(1, h // 2)))
img_small.save(file_path, format="PNG", compress_level=9)
with open(file_path, "rb") as f:
files.append({"data": f.read(), "name": f"image_{idx}.png"})
finally:
shutil.rmtree(temp_dir)
return files
# === Nodes ===
class DiscordSetWebhook:
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("DUMMY IMAGE",)
@@ -52,16 +132,15 @@ class DiscordSetWebhook:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"URL": ("STRING",)}}
def execute(self, URL):
if not URL.startswith("https://discord.com/api/webhooks/"):
raise ValueError("Invalid URL format. URL should start with 'https://discord.com/api/webhooks/' Please reference https://support.discord.com/hc/en-us/articles/228383668-Intro-to-Webhooks")
with open("discord_webhook_url.txt", "w") as f:
f.write(URL)
raise ValueError("Invalid URL. Must start with 'https://discord.com/api/webhooks/'.")
os.makedirs(SECRETS_DIR, exist_ok=True)
with open(WEBHOOK_FILE, "w", encoding="utf-8") as f:
f.write(URL.strip())
return (create_default_image(),)
class DiscordPostViaWebhook:
RETURN_TYPES = ("IMAGE",)
OUTPUT_NODE = True
@@ -73,7 +152,7 @@ class DiscordPostViaWebhook:
return {
"required": {
"image": ("IMAGE",)
},
},
"optional": {
"send_Message": ("BOOLEAN", {"default": False}),
"send_Image": ("BOOLEAN", {"default": True}),
@@ -82,136 +161,68 @@ class DiscordPostViaWebhook:
},
}
async def send_webhook(self, url, message, files=None):
"""Send the webhook with the given message and up to 4 files."""
async def _send_webhook(self, url, message, files=None):
webhook = AsyncDiscordWebhook(url=url, content=message[:2000], timeout=30.0)
if files:
for file in files:
webhook.add_file(file=file["data"], filename=file["name"])
await webhook.execute()
def process_image(self, image):
"""Process the image (or batch of images) and return them in a format suitable for Discord."""
images_to_send = []
if image is None:
image = create_default_image()
# Check if it's a batched image (4D array: [batch_size, height, width, channels])
if isinstance(image, np.ndarray):
if image.ndim == 4:
# Batch of images, process each image separately
for i in range(image.shape[0]):
img = image[i] # Don't squeeze unless we know the axis has size 1
img = Image.fromarray(np.clip(img * 255, 0, 255).astype(np.uint8))
images_to_send.append(img)
elif image.ndim == 3:
# Single image (3D array: [height, width, channels])
image = Image.fromarray(np.clip(image * 255, 0, 255).astype(np.uint8))
images_to_send.append(image)
else:
raise ValueError("Input image array must be 3D or 4D (batch of images).")
elif hasattr(image, "cpu"):
array = image.cpu().numpy()
# Handle batched images
if array.ndim == 4: # Batch of images
for i in range(array.shape[0]):
img_array = array[i] # Only access the i-th image, no squeeze needed
img = Image.fromarray(np.clip(img_array * 255, 0, 255).astype(np.uint8))
images_to_send.append(img)
elif array.ndim == 3:
# Single image
array = np.clip(array * 255, 0, 255).astype(np.uint8)
images_to_send.append(Image.fromarray(array))
else:
raise ValueError("Input tensor must be 3D or 4D (batch of images).")
else:
images_to_send.append(image)
# Save each image to a temporary file and collect file data
files = []
temp_dir = tempfile.mkdtemp()
for idx, img in enumerate(images_to_send):
file_path = os.path.join(temp_dir, f"image_{idx}.png")
img.save(file_path, format="PNG", compress_level=1)
# If the file size exceeds 20MB, resize and save again
if os.path.getsize(file_path) > 20 * 1024 * 1024:
img = img.resize((img.width // 2, img.height // 2))
img.save(file_path, format="PNG", compress_level=9)
with open(file_path, "rb") as f:
files.append({"data": f.read(), "name": f"image_{idx}.png"})
# Clean up the temporary directory
shutil.rmtree(temp_dir)
return files
def _process_image(self, image):
imgs = _tensor_or_nd_to_pil_list(image)
return _prepare_files_for_discord(imgs) if imgs else None
async def execute(self, image, send_Message=True, send_Image=True, message="", prepend_message=""):
with open("discord_webhook_url.txt", "r") as f:
webhook_url = f.read().strip()
webhook_url = get_webhook_url()
if not webhook_url:
raise ValueError("Webhook URL is empty.")
raise ValueError("Missing Discord webhook URL. Set env DISCORD_WEBHOOK_URL "
"or create secrets/webhook.txt")
final_message = ""
if send_Message:
message = f"{prepend_message}\n{message}"
files = self.process_image(image) if send_Image else None
final_message = (prepend_message or "").strip()
if final_message:
final_message += "\n"
final_message += (message or "")
files = self._process_image(image) if send_Image else None
if files:
# Split files into batches of 4 (Discord limit)
# Discord: max 10 files per message nowadays, mais on reste safe à 4
batches = [files[i:i + 4] for i in range(0, len(files), 4)]
# Send multiple webhooks if necessary
for batch in batches:
await self.send_webhook(webhook_url, message, batch)
await self._send_webhook(webhook_url, final_message, batch)
else:
# No images to send, just send the message
await self.send_webhook(webhook_url, message)
await self._send_webhook(webhook_url, final_message)
return (image,)
NODE_CLASS_MAPPINGS = {
"DiscordSetWebhook": DiscordSetWebhook,
"DiscordPostViaWebhook": DiscordPostViaWebhook
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DiscordSetWebhook": "Set Discord Webhook",
"DiscordPostViaWebhook": "Use Discord Webhook"
}
class DiscordPostViaWebhookWithInput(DiscordPostViaWebhook):
"""
Variante du node avec entrée connectable pour le message (prompt).
Basée sur DiscordPostViaWebhook pour ne rien casser.
Variant with connectable inputs for message/prepend (safe: no URL stored in workflow).
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"message": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
},
"optional": {
"send_Message": ("BOOLEAN", {"default": False}),
"send_Image": ("BOOLEAN", {"default": True}),
# 👇 Ajout important: on force une ENTRÉE (socket) pour le texte
"message": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
# Optionnel: on peut aussi forcer une entrée pour le prepend
"prepend_message": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
},
}
# Ajoute ces lignes à la fin des mappings existants
NODE_CLASS_MAPPINGS.update({
"DiscordPostViaWebhookWithInput": DiscordPostViaWebhookWithInput
})
NODE_CLASS_MAPPINGS = {
"DiscordSetWebhook": DiscordSetWebhook,
"DiscordPostViaWebhook": DiscordPostViaWebhook,
"DiscordPostViaWebhookWithInput": DiscordPostViaWebhookWithInput,
}
NODE_DISPLAY_NAME_MAPPINGS.update({
"DiscordPostViaWebhookWithInput": "Use Discord Webhook (+message input)"
})
NODE_DISPLAY_NAME_MAPPINGS = {
"DiscordSetWebhook": "Set Discord Webhook",
"DiscordPostViaWebhook": "Use Discord Webhook",
"DiscordPostViaWebhookWithInput": "Use Discord Webhook (+message input)",
}