isala404-comfy-workflow-api/server/uploads.py

"""Handle file uploads for webhook requests."""

import logging
import os
import shutil
import uuid
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# Track uploads per request for cleanup
_upload_registry: dict[str, list[str]] = {}


def get_upload_directory() -> str:
    """Get the directory for webhook uploads."""
    try:
        import folder_paths
        # Use ComfyUI's input directory for uploads
        input_dir = folder_paths.get_input_directory()
        webhook_dir = os.path.join(input_dir, "webhook_uploads")
    except ImportError:
        # Fallback for testing
        webhook_dir = os.path.join(os.getcwd(), "input", "webhook_uploads")

    os.makedirs(webhook_dir, exist_ok=True)
    return webhook_dir


def save_upload(data: bytes, filename: str, request_id: str) -> str:
    """
    Save an uploaded file to the webhook uploads directory.

    Args:
        data: File content as bytes
        filename: Original filename
        request_id: Request ID for organizing uploads

    Returns:
        Full path to the saved file
    """
    upload_dir = get_upload_directory()

    # Create request-specific subdirectory
    request_dir = os.path.join(upload_dir, request_id)
    os.makedirs(request_dir, exist_ok=True)

    # Sanitize filename and make unique
    safe_filename = sanitize_filename(filename)
    base, ext = os.path.splitext(safe_filename)

    # Add short UUID to avoid collisions
    unique_filename = f"{base}_{uuid.uuid4().hex[:8]}{ext}"
    file_path = os.path.join(request_dir, unique_filename)

    # Write file
    with open(file_path, "wb") as f:
        f.write(data)

    # Track for cleanup
    if request_id not in _upload_registry:
        _upload_registry[request_id] = []
    _upload_registry[request_id].append(file_path)

    logger.debug(f"Saved upload: {file_path} ({len(data)} bytes)")
    return file_path


def sanitize_filename(filename: str) -> str:
    """Sanitize a filename for safe storage."""
    # Remove path separators and null bytes
    filename = filename.replace("/", "_").replace("\\", "_").replace("\x00", "")

    # Remove leading dots (hidden files)
    while filename.startswith("."):
        filename = filename[1:]

    # Default name if empty
    if not filename:
        filename = "upload"

    # Limit length
    if len(filename) > 200:
        base, ext = os.path.splitext(filename)
        filename = base[:200 - len(ext)] + ext

    return filename


def cleanup_uploads(request_id: str) -> int:
    """
    Clean up uploaded files for a request.

    Args:
        request_id: Request ID to clean up

    Returns:
        Number of files removed
    """
    count = 0

    # Remove tracked files
    if request_id in _upload_registry:
        for file_path in _upload_registry[request_id]:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
                    count += 1
            except Exception as e:
                logger.warning(f"Failed to remove upload {file_path}: {e}")

        del _upload_registry[request_id]

    # Also try to remove the request directory
    upload_dir = get_upload_directory()
    request_dir = os.path.join(upload_dir, request_id)

    try:
        if os.path.exists(request_dir):
            shutil.rmtree(request_dir)
            logger.debug(f"Removed upload directory: {request_dir}")
    except Exception as e:
        logger.warning(f"Failed to remove upload directory {request_dir}: {e}")

    return count


def get_upload_path(request_id: str, field_name: str) -> Optional[str]:
    """
    Get the path to an uploaded file by field name.

    This is used by WebhookInput node to load files.

    Args:
        request_id: Request ID
        field_name: Field name from the upload

    Returns:
        File path if found, None otherwise
    """
    # This is primarily resolved via extra_data["webhook_uploads"]
    # But we can search the registry as fallback
    if request_id in _upload_registry:
        for path in _upload_registry[request_id]:
            if field_name in path:
                return path
    return None


def cleanup_old_uploads(max_age_hours: int = 24) -> int:
    """
    Clean up old upload directories.

    Args:
        max_age_hours: Remove directories older than this

    Returns:
        Number of directories removed
    """
    import time

    upload_dir = get_upload_directory()
    if not os.path.exists(upload_dir):
        return 0

    count = 0
    cutoff = time.time() - (max_age_hours * 3600)

    try:
        for entry in os.scandir(upload_dir):
            if entry.is_dir():
                try:
                    mtime = entry.stat().st_mtime
                    if mtime < cutoff:
                        shutil.rmtree(entry.path)
                        count += 1
                        logger.debug(f"Cleaned up old upload directory: {entry.path}")
                except Exception as e:
                    logger.warning(f"Failed to clean up {entry.path}: {e}")
    except Exception as e:
        logger.warning(f"Failed to scan upload directory: {e}")

    if count > 0:
        logger.info(f"Cleaned up {count} old upload directories")

    return count