Add custom node

basic_auth.py

@@ -0,0 +1,80 @@
+import os
+import base64
+from aiohttp import web
+
+class BasicAuthMiddleware:
+    def __init__(self):
+        self.username = os.getenv('COMFYUI_USERNAME', '')
+        self.password = os.getenv('COMFYUI_PASSWORD', '')
+        self.enabled = bool(self.username and self.password)
+
+    @web.middleware
+    async def handle(self, request, handler):
+        if not self.enabled:
+            return await handler(request)
+
+        # Skip auth for WebSocket connections
+        if request.path == '/ws':
+            return await handler(request)
+
+        # Get Authorization header
+        auth_header = request.headers.get('Authorization')
+
+        if not auth_header:
+            return web.Response(
+                status=401,
+                headers={'WWW-Authenticate': 'Basic realm="ComfyUI Server"'}
+            )
+
+        try:
+            auth_type, auth_string = auth_header.split(' ', 1)
+            if auth_type.lower() != 'basic':
+                raise ValueError('Invalid auth type')
+
+            decoded = base64.b64decode(auth_string).decode('utf-8')
+            provided_username, provided_password = decoded.split(':', 1)
+
+            if provided_username == self.username and provided_password == self.password:
+                return await handler(request)
+
+        except Exception:
+            pass
+
+        return web.Response(
+            status=401,
+            headers={'WWW-Authenticate': 'Basic realm="ComfyUI Server"'}
+        )
+
+class BasicAuthSetup:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "enabled": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("BASIC_AUTH",)
+    FUNCTION = "setup_auth"
+    CATEGORY = "utils"
+    OUTPUT_NODE = True
+
+    def setup_auth(self, enabled):
+        return ({"enabled": enabled},)
+
+NODE_CLASS_MAPPINGS = {
+    "BasicAuthSetup": BasicAuthSetup
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BasicAuthSetup": "Basic Auth Setup"
+}
+
+# Register the middleware
+try:
+    import server
+    app = server.PromptServer.instance.app
+    middleware = BasicAuthMiddleware()
+    app.middlewares.insert(0, middleware.handle)
+except Exception as e:
+    print(f"Failed to register basic auth middleware: {e}")