| | import json
|
| | import os
|
| | import re
|
| | import uuid
|
| | import glob
|
| | import shutil
|
| | from aiohttp import web
|
| | from comfy.cli_args import args
|
| | from folder_paths import user_directory
|
| | from .app_settings import AppSettings
|
| |
|
| | default_user = "default"
|
| | users_file = os.path.join(user_directory, "users.json")
|
| |
|
| |
|
| | class UserManager():
|
| | def __init__(self):
|
| | global user_directory
|
| |
|
| | self.settings = AppSettings(self)
|
| | if not os.path.exists(user_directory):
|
| | os.mkdir(user_directory)
|
| | if not args.multi_user:
|
| | print("****** User settings have been changed to be stored on the server instead of browser storage. ******")
|
| | print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******")
|
| |
|
| | if args.multi_user:
|
| | if os.path.isfile(users_file):
|
| | with open(users_file) as f:
|
| | self.users = json.load(f)
|
| | else:
|
| | self.users = {}
|
| | else:
|
| | self.users = {"default": "default"}
|
| |
|
| | def get_request_user_id(self, request):
|
| | user = "default"
|
| | if args.multi_user and "comfy-user" in request.headers:
|
| | user = request.headers["comfy-user"]
|
| |
|
| | if user not in self.users:
|
| | raise KeyError("Unknown user: " + user)
|
| |
|
| | return user
|
| |
|
| | def get_request_user_filepath(self, request, file, type="userdata", create_dir=True):
|
| | global user_directory
|
| |
|
| | if type == "userdata":
|
| | root_dir = user_directory
|
| | else:
|
| | raise KeyError("Unknown filepath type:" + type)
|
| |
|
| | user = self.get_request_user_id(request)
|
| | path = user_root = os.path.abspath(os.path.join(root_dir, user))
|
| |
|
| |
|
| | if os.path.commonpath((root_dir, user_root)) != root_dir:
|
| | return None
|
| |
|
| | if file is not None:
|
| |
|
| | path = os.path.abspath(os.path.join(user_root, file))
|
| | if os.path.commonpath((user_root, path)) != user_root:
|
| | return None
|
| |
|
| | parent = os.path.split(path)[0]
|
| |
|
| | if create_dir and not os.path.exists(parent):
|
| | os.makedirs(parent, exist_ok=True)
|
| |
|
| | return path
|
| |
|
| | def add_user(self, name):
|
| | name = name.strip()
|
| | if not name:
|
| | raise ValueError("username not provided")
|
| | user_id = re.sub("[^a-zA-Z0-9-_]+", '-', name)
|
| | user_id = user_id + "_" + str(uuid.uuid4())
|
| |
|
| | self.users[user_id] = name
|
| |
|
| | global users_file
|
| | with open(users_file, "w") as f:
|
| | json.dump(self.users, f)
|
| |
|
| | return user_id
|
| |
|
| | def add_routes(self, routes):
|
| | self.settings.add_routes(routes)
|
| |
|
| | @routes.get("/users")
|
| | async def get_users(request):
|
| | if args.multi_user:
|
| | return web.json_response({"storage": "server", "users": self.users})
|
| | else:
|
| | user_dir = self.get_request_user_filepath(request, None, create_dir=False)
|
| | return web.json_response({
|
| | "storage": "server",
|
| | "migrated": os.path.exists(user_dir)
|
| | })
|
| |
|
| | @routes.post("/users")
|
| | async def post_users(request):
|
| | body = await request.json()
|
| | username = body["username"]
|
| | if username in self.users.values():
|
| | return web.json_response({"error": "Duplicate username."}, status=400)
|
| |
|
| | user_id = self.add_user(username)
|
| | return web.json_response(user_id)
|
| |
|
| | @routes.get("/userdata")
|
| | async def listuserdata(request):
|
| | directory = request.rel_url.query.get('dir', '')
|
| | if not directory:
|
| | return web.Response(status=400)
|
| |
|
| | path = self.get_request_user_filepath(request, directory)
|
| | if not path:
|
| | return web.Response(status=403)
|
| |
|
| | if not os.path.exists(path):
|
| | return web.Response(status=404)
|
| |
|
| | recurse = request.rel_url.query.get('recurse', '').lower() == "true"
|
| | results = glob.glob(os.path.join(
|
| | glob.escape(path), '**/*'), recursive=recurse)
|
| | results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)]
|
| |
|
| | split_path = request.rel_url.query.get('split', '').lower() == "true"
|
| | if split_path:
|
| | results = [[x] + x.split(os.sep) for x in results]
|
| |
|
| | return web.json_response(results)
|
| |
|
| | def get_user_data_path(request, check_exists = False, param = "file"):
|
| | file = request.match_info.get(param, None)
|
| | if not file:
|
| | return web.Response(status=400)
|
| |
|
| | path = self.get_request_user_filepath(request, file)
|
| | if not path:
|
| | return web.Response(status=403)
|
| |
|
| | if check_exists and not os.path.exists(path):
|
| | return web.Response(status=404)
|
| |
|
| | return path
|
| |
|
| | @routes.get("/userdata/{file}")
|
| | async def getuserdata(request):
|
| | path = get_user_data_path(request, check_exists=True)
|
| | if not isinstance(path, str):
|
| | return path
|
| |
|
| | return web.FileResponse(path)
|
| |
|
| | @routes.post("/userdata/{file}")
|
| | async def post_userdata(request):
|
| | path = get_user_data_path(request)
|
| | if not isinstance(path, str):
|
| | return path
|
| |
|
| | overwrite = request.query["overwrite"] != "false"
|
| | if not overwrite and os.path.exists(path):
|
| | return web.Response(status=409)
|
| |
|
| | body = await request.read()
|
| |
|
| | with open(path, "wb") as f:
|
| | f.write(body)
|
| |
|
| | resp = os.path.relpath(path, self.get_request_user_filepath(request, None))
|
| | return web.json_response(resp)
|
| |
|
| | @routes.delete("/userdata/{file}")
|
| | async def delete_userdata(request):
|
| | path = get_user_data_path(request, check_exists=True)
|
| | if not isinstance(path, str):
|
| | return path
|
| |
|
| | os.remove(path)
|
| |
|
| | return web.Response(status=204)
|
| |
|
| | @routes.post("/userdata/{file}/move/{dest}")
|
| | async def move_userdata(request):
|
| | source = get_user_data_path(request, check_exists=True)
|
| | if not isinstance(source, str):
|
| | return source
|
| |
|
| | dest = get_user_data_path(request, check_exists=False, param="dest")
|
| | if not isinstance(source, str):
|
| | return dest
|
| |
|
| | overwrite = request.query["overwrite"] != "false"
|
| | if not overwrite and os.path.exists(dest):
|
| | return web.Response(status=409)
|
| |
|
| | print(f"moving '{source}' -> '{dest}'")
|
| | shutil.move(source, dest)
|
| |
|
| | resp = os.path.relpath(dest, self.get_request_user_filepath(request, None))
|
| | return web.json_response(resp)
|
| |
|
