magicwang1111-ComfyUI-Account-Manager/account_manager.py

import os
import logging
import jwt
import uuid
import folder_paths
from aiohttp import web

from server import PromptServer, is_loopback

from .utils import *
from .utils.api_audit import install_api_audit
from .utils.frontend import create_jobs_middleware, install_origin_middleware
from .utils.frontend_reload import install_frontend_reload

instance = PromptServer.instance
app = instance.app
routes = instance.routes
install_origin_middleware(app, is_loopback)
install_frontend_reload(app, routes)

logger = Logger(LOG_FILE, LOG_LEVELS)
sanitizer = Sanitizer()
ip_filter = IPFilter(WHITELIST, BLACKLIST)
timeout = Timeout(ip_filter, BLACKLIST_AFTER_ATTEMPTS)
users_db = UsersDB(USERS_FILE)
scheduler = None
cloud_archive = None
event_relay = None
internal_signer = None
if DISTRIBUTED_QUEUE_ENABLED:
    scheduler = SchedulerStore(
        SCHEDULER_FILE,
        max_concurrent_jobs_per_user=MAX_CONCURRENT_JOBS_PER_USER,
        admin_concurrency_limit=ADMIN_CONCURRENCY_LIMIT,
        resource_concurrency_limits=RESOURCE_CONCURRENCY_LIMITS,
        gpu_node_types=GPU_NODE_TYPES,
    )
    internal_signer = InternalSigner(SCHEDULER_SECRET_FILE)
    install_api_audit(scheduler.record_api_call)

    if CLOUD_ARCHIVE_ENABLED:
        try:
            cloud_archive_config = CloudArchiveConfig(
                enabled=True,
                server_id=CLOUD_SERVER_ID,
                region=CLOUD_OSS_REGION,
                endpoint=CLOUD_OSS_ENDPOINT,
                bucket=CLOUD_OSS_BUCKET,
                prefix=CLOUD_OSS_PREFIX,
                public_base_url=CLOUD_PUBLIC_BASE_URL,
                max_attempts=CLOUD_MAX_ATTEMPTS,
                upload_concurrency=CLOUD_UPLOAD_CONCURRENCY,
                remote_max_bytes=CLOUD_REMOTE_MAX_BYTES,
                manifest_max_bytes=CLOUD_MANIFEST_MAX_BYTES,
                staging_dir=CLOUD_STAGING_DIR,
                checkpoint_dir=CLOUD_CHECKPOINT_DIR,
            )
            cloud_archive = CloudArchiveManager(
                cloud_archive_config,
                scheduler_database=SCHEDULER_FILE,
                history_database=HISTORY_FILE,
                output_dir=folder_paths.get_output_directory(),
                temp_dir=folder_paths.get_temp_directory(),
                project_dir=os.path.dirname(__file__),
                logger=logging.getLogger("ComfyUI-Account-Manager"),
            )
        except Exception:
            logging.getLogger("ComfyUI-Account-Manager").exception(
                "Failed to initialize cloud archive"
            )

access_control = AccessControl(
    users_db,
    instance,
    HISTORY_FILE,
    scheduler=scheduler,
    instance_port=INSTANCE_PORT,
    heartbeat_seconds=WORKER_HEARTBEAT_SECONDS,
    stale_seconds=WORKER_STALE_SECONDS,
    worker_resource_class=WORKER_RESOURCE_CLASS,
    cloud_archive=cloud_archive,
)
jwt_auth = JWTAuth(
    users_db, access_control, logger, SECRET_KEY, TOKEN_EXPIRE_MINUTES, TOKEN_ALGORITHM
)
access_control.sync_comfy_users()

if scheduler:
    event_relay = EventRelay(instance, scheduler, internal_signer, INSTANCE_PORT)
    event_relay.install()


@routes.post(EventRelay.PATH)
async def post_internal_event(request: web.Request) -> web.Response:
    if not event_relay:
        return web.json_response({"error": "Distributed queue is disabled"}, status=404)
    return await event_relay.handle_event(request)


@routes.get("/register")
async def get_register(request: web.Request) -> web.Response:
    with open(os.path.join(HTML_DIR, "register.html"), "r") as f:
        html_content = f.read()

    if not users_db.load_users():
        html_content = html_content.replace("{{ X-Admin-User }}", "true")
    html_content = html_content.replace("{{ X-Admin-User }}", "false")

    return web.Response(body=html_content, content_type="text/html")


@routes.post("/register")
async def post_register(request: web.Request) -> web.Response:
    sanitized_data = request.get("_sanitized_data", {})
    ip = get_ip(request)
    new_user_username = sanitized_data.get("new_user_username")
    new_user_password = sanitized_data.get("new_user_password")
    username = sanitized_data.get("username")
    password = sanitized_data.get("password")

    username_valid, username_invalid_message = validate_username(new_user_username)
    if not username_valid:
        return web.json_response({"error": username_invalid_message}, status=400)

    password_valid, password_invalid_message = validate_password(new_user_password)
    if not password_valid:
        return web.json_response({"error": password_invalid_message}, status=400)

    admin_user = users_db.get_admin_user()
    admin_user_id = None

    if admin_user[0] and (not new_user_username or not new_user_password):
        return web.json_response(
            {"error": "Missing new user registration details"}, status=400
        )

    if admin_user[0]:
        if not username or not password:
            return web.json_response(
                {"error": "Missing admin user authentication details"}, status=400
            )

        admin_user_id = admin_user[0]

        if admin_user_id is not None:
            if not (
                users_db.get_user(username)[0] == admin_user_id
                and users_db.check_username_password(username, password)
            ):
                logger.registration_attempt(
                    ip, username, password, new_user_username, new_user_password
                )
                timeout.add_failed_attempt(ip)
                return web.json_response(
                    {"message": "Invalid username or password"}, status=403
                )

    if None not in users_db.get_user(new_user_username):
        return web.json_response({"error": "Username already exists"}, status=400)

    new_user_id = str(uuid.uuid4())
    try:
        users_db.add_user(
            new_user_id,
            new_user_username,
            new_user_password,
            not bool(admin_user_id),
        )
    except ValueError:
        return web.json_response({"error": "Username already exists"}, status=400)
    access_control.sync_comfy_users()

    logger.registration_success(
        ip, new_user_username, username if admin_user_id is not None else None
    )
    timeout.remove_failed_attempts(ip)
    return web.json_response({"message": "User registered successfully"})


@routes.get("/login")
async def get_login(request: web.Request) -> web.Response:
    if not users_db.load_users():
        return web.HTTPFound("/register")

    token = jwt_auth.get_token_from_request(request)
    if token:
        return web.HTTPFound("/logout")
    return web.FileResponse(os.path.join(HTML_DIR, "login.html"))


@routes.post("/login")
async def post_login(request: web.Request) -> web.Response:
    sanitized_data = request.get("_sanitized_data", {})
    ip = get_ip(request)
    username = sanitized_data.get("username")
    password = sanitized_data.get("password")

    if not username or not password:
        return web.json_response(
            {"error": "Missing login credentials (username and password)"}, status=400
        )

    if users_db.check_username_password(username, password):
        timeout.remove_failed_attempts(ip)

        user_id, _ = users_db.get_user(username)
        token = jwt_auth.create_access_token({"id": user_id, "username": username})
        response = web.json_response(
            {
                "message": "Login successful",
                "jwt_token": token,
                "user_settings_id": user_id,
            }
        )
        secure_flag = request.headers.get("X-Forwarded-Proto", "http") == "https"
        response.set_cookie(
            "jwt_token",
            token,
            max_age=TOKEN_EXPIRE_MINUTES * 60,
            httponly=True,
            secure=secure_flag,
            samesite="Strict",
        )
        logger.login_success(ip, username)
        return response

    logger.login_attempt(ip, username, password)
    timeout.add_failed_attempt(ip)
    return web.json_response({"error": "Invalid username or password"}, status=401)


@routes.get("/generate_token")
async def get_generate_token(request: web.Request) -> web.Response:
    if not users_db.load_users():
        return web.HTTPFound("/register")

    token = jwt_auth.get_token_from_request(request)
    if token:
        return web.HTTPFound("/logout")
    return web.FileResponse(os.path.join(HTML_DIR, "generate_token.html"))


@routes.post("/generate_token")
async def post_generate_token(request: web.Request) -> web.Response:
    sanitized_data = request.get("_sanitized_data", {})
    ip = get_ip(request)
    username = sanitized_data.get("username")
    password = sanitized_data.get("password")

    try:
        expire_hours = int(
            sanitized_data.get("expire_hours", TOKEN_EXPIRE_MINUTES / 60)
        )

    except ValueError:
        return web.json_response(
            {"error": "Expiration hours must be a number"},
            status=400,
        )

    if expire_hours > MAX_TOKEN_EXPIRE_MINUTES / 60:
        return web.json_response(
            {
                "error": f"Expiration hours must be smaller than {MAX_TOKEN_EXPIRE_MINUTES / 60}"
            },
            status=400,
        )

    if not username or not password:
        return web.json_response(
            {"error": "Missing login credentials (username and password)"}, status=400
        )

    if users_db.check_username_password(username, password):
        timeout.remove_failed_attempts(ip)

        user_id, _ = users_db.get_user(username)
        token = jwt_auth.create_access_token(
            {"id": user_id, "username": username}, expire_minutes=(expire_hours * 60)
        )
        response = web.json_response(
            {
                "message": "JWT Token successfully generated",
                "jwt_token": token,
            }
        )
        secure_flag = request.headers.get("X-Forwarded-Proto", "http") == "https"
        response.set_cookie(
            "jwt_token",
            token,
            max_age=expire_hours * 60 * 60,
            httponly=True,
            secure=secure_flag,
            samesite="Strict",
        )
        
        logger.generate_success(ip, username, expire_hours)
        
        return response

    logger.generate_attempt(ip, username, password, expire_hours)
    timeout.add_failed_attempt(ip)
    return web.json_response({"error": "Invalid username or password"}, status=401)


@routes.get("/logout")
async def get_logout(request: web.Request) -> web.Response:
    ip = get_ip(request)
    free_memory = request.query.get("free_memory", "false").lower() == "true"
    unload_models = request.query.get("unload_models", "false").lower() == "true"

    token = jwt_auth.get_token_from_request(request)
    if token and FREE_MEMORY_ON_LOGOUT:
        try:
            username = jwt_auth.decode_access_token(token).get("username")
            if free_memory or unload_models:
                if hasattr(instance, "post_free"):
                    json_data = {
                        "unload_models": unload_models,
                        "free_memory": free_memory,
                    }
                    mock_request = web.Request(
                        app=app,
                        method="POST",
                        path="/free",
                        headers={},
                        match_info={},
                        payload=None,
                    )
                    mock_request._post = json_data
                    await instance.post_free(mock_request)
                    logger.memory_free(ip, username, free_memory, unload_models)

            logger.logout(ip, username)
        except jwt.ExpiredSignatureError:
            pass
        except jwt.InvalidTokenError:
            pass
        except Exception as e:
            logger.error(f"Unexpected error during logout: {e}")

    response = web.HTTPFound("/login")
    response.del_cookie("jwt_token", path="/")

    return response


app.add_routes(
    [
        web.static("/account-manager/css", CSS_DIR),
        web.static("/account-manager/js", JS_DIR),
        web.static("/account-manager/assets", ASSETS_DIR),
    ]
)

if FORCE_HTTPS:
    app.middlewares.append(create_https_middleware(MATCH_HEADERS))

app.middlewares.append(ip_filter.create_ip_filter_middleware())
app.middlewares.append(sanitizer.create_sanitizer_middleware())
app.middlewares.append(
    timeout.create_time_out_middleware(
        limited=("/login", "/register", "/generate_token")
    )
)
app.middlewares.append(
    jwt_auth.create_jwt_middleware(
        public=("/login", "/logout", "/register", "/generate_token"),
        public_prefixes=("/account-manager"),
    )
)

if SEPERATE_USERS:
    access_control.patch_comfy_user_manager()
    access_control.patch_assets()
    if access_control.asset_access is not None:
        app.middlewares.append(access_control.asset_access.middleware())

if scheduler:
    distributed_routes = DistributedRoutes(
        scheduler,
        internal_signer,
        users_db,
        INSTANCE_PORT,
        WORKER_STALE_SECONDS,
    )
    app.middlewares.append(distributed_routes.middleware())

if SEPERATE_USERS:
    app.middlewares.append(access_control.create_view_access_control_middleware())
    app.middlewares.append(access_control.create_folder_access_control_middleware())

    access_control.patch_folder_paths()
    access_control.patch_prompt_queue()
    app.middlewares.append(create_jobs_middleware(instance.prompt_queue, access_control.user_job_history))

if MANAGER_ADMIN_ONLY:
    app.middlewares.append(
        access_control.create_manager_access_control_middleware(
            manager_directory="/extensions/comfyui-manager",
            manager_routes=(
                "api/customnode",
                "api/snapshot",
                "/api/manager",
                "api/comfyui_manager",
                "api/externalmodel",
            ),
        )
    )