URL and Environment Variable Validation

This commit is contained in:
Mattabyte
2025-01-28 14:48:56 +10:00
parent 1958abc9ab
commit 3fd0284802
+34 -4
View File
@@ -2,10 +2,24 @@ import requests
import json
import server
import os
import urllib.parse
from typing import Optional
from .util import ComfyAnyType
class SecureApiCall:
ALLOWED_SCHEMES = {'https', 'http'}
@staticmethod
def validate_url(url: str) -> Optional[str]:
try:
parsed = urllib.parse.urlparse(url)
if parsed.scheme not in SecureApiCall.ALLOWED_SCHEMES:
return f"URL scheme must be one of: {', '.join(SecureApiCall.ALLOWED_SCHEMES)}"
return None
except Exception:
return "Invalid URL format"
@classmethod
def INPUT_TYPES(cls):
return {
@@ -26,24 +40,40 @@ class SecureApiCall:
CATEGORY = "notifications"
def hook(self, any, api_url, api_auth, data, full_comfyui_info, timeout, verify_ssl):
# Validate URL before proceeding
url_error = self.validate_url(api_url)
if url_error:
raise ValueError(f"Invalid API URL: {url_error}")
# Get the payload from the data string
payload = json.loads(data)
try:
payload = json.loads(data)
except json.JSONDecodeError:
raise ValueError("Invalid JSON data provided")
# Add ComfyUI execution information to the payload
current_queue = server.PromptServer.instance.prompt_queue.get_current_queue()
# The response is a massive object (containing all nodes and workflow info) -- we only want the prompt_id for tracking, but if full_comfyui_info is true, we want to include the entire object
prompt_id = current_queue[0][0][1]
payload.update({"comfyui_execution_info": {"prompt_id": prompt_id}})
if include_full_comfyui_info:
if full_comfyui_info:
payload.update({"comfyui_info": {"full_comfyui_info": current_queue}})
# Resolve any environment variables in the api_url, api_auth fields -- Only allow COMFYUI_SECUREAPICALL_ prefixed variables
if api_url.startswith("$ENV."):
env_var = f"COMFYUI_SECUREAPICALL_{api_url.removeprefix('$ENV.')}"
# Only allow alphanumeric and underscore characters in env var names
env_name = api_url.removeprefix("$ENV.")
if not env_name.replace("_", "").isalnum():
raise ValueError("Environment variable names must only contain letters, numbers, and underscores")
env_var = f"COMFYUI_SECUREAPICALL_{env_name}"
api_url = os.getenv(env_var) or 'NOT_SET'
if api_auth.startswith("$ENV."):
env_var = f"COMFYUI_SECUREAPICALL_{api_auth.removeprefix('$ENV.')}"
env_name = api_auth.removeprefix("$ENV.")
if not env_name.replace("_", "").isalnum():
raise ValueError("Environment variable names must only contain letters, numbers, and underscores")
env_var = f"COMFYUI_SECUREAPICALL_{env_name}"
api_auth = os.getenv(env_var) or 'NOT_SET'
# Send the payload to the API