81 lines
2.2 KiB
Python
81 lines
2.2 KiB
Python
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}")
|