AEmotionStudio-ComfyUI-DiscordSend/shared/workflow/sanitizer.py

"""
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
# Pre-compile regex for faster matching
WEBHOOK_REGEX = re.compile(
    r"(discord\.com/api/webhooks|discordapp\.com/api/webhooks)",
    re.IGNORECASE
)

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
    
    if WEBHOOK_REGEX.search(value):
        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
    
    return value.startswith(GITHUB_TOKEN_PREFIXES)


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: Any) -> Any:
    """
    Sanitize a single ComfyUI node.
    
    Args:
        node: The node dictionary or value
        
    Returns:
        Sanitized node dictionary or value
    """
    if not isinstance(node, dict):
        if isinstance(node, str):
            return sanitize_string(node)
        return node
    
    result = {}
    node_type = node.get("type", "")
    
    for key, value in node.items():
        # Handle known sensitive keys
        if key in ("webhook_url", "github_token"):
            result[key] = ""
            continue

        # Context-aware sanitization for inputs and widgets
        if key == "inputs" and isinstance(value, dict):
            result[key] = sanitize_node_inputs(value, node_type)
        elif key == "widgets_values" and isinstance(value, list):
            result[key] = sanitize_widget_values(value, node_type)
        else:
            # Generic sanitization for other fields
            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
    
    return result


def sanitize_dict(data: Dict) -> Dict:
    """
    Recursively sanitize a dictionary.
    
    Args:
        data: Dictionary to sanitize
        
    Returns:
        Sanitized dictionary
    """
    result = {}
    
    # Check if this is a workflow object with nodes
    # We want to process nodes specifically using sanitize_node to ensure correct context
    # and avoid double-processing (once as generic dict/list, once as nodes)
    is_workflow = "nodes" in data and isinstance(data["nodes"], (list, dict))

    for key, value in data.items():
        # Handle known sensitive keys
        if key in ("webhook_url", "github_token"):
            result[key] = ""
            continue
        
        # Special handling for "nodes" in workflow
        if is_workflow and key == "nodes":
            if isinstance(value, list):
                result[key] = [sanitize_node(n) for n in value]
            elif isinstance(value, dict):
                result[key] = {k: sanitize_node(v) for k, v in value.items()}
            else:
                # Fallback if nodes is neither list nor dict (unlikely)
                result[key] = value
            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
    
    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