Dr.Lt.Data
2024-05-14 23:54:06 +09:00
parent ef6b6ff690
commit 1c313ee822
9 changed files with 198 additions and 106 deletions

View File

@@ -42,6 +42,36 @@ from comfy.cli_args import args
import latent_preview
is_local_mode = args.listen.startswith('127.')
def is_allowed_security_level(level):
if level == 'high':
if is_local_mode:
return core.get_config()['security_level'].lower() in ['weak', 'normal']
else:
return core.get_config()['security_level'].lower() == 'weak'
elif level == 'middle':
return core.get_config()['security_level'].lower() in ['weak', 'normal']
else:
return True
async def get_risky_level(files):
json_data1 = await core.get_data_by_mode('local', 'custom-node-list.json')
json_data2 = await core.get_data_by_mode('cache', 'custom-node-list.json', channel_url='https://github.com/ltdrdata/ComfyUI-Manager/raw/main/custom-node-list.json')
all_urls = set()
for x in json_data1['custom_nodes'] + json_data2['custom_nodes']:
all_urls.update(x['files'])
for x in files:
if x not in all_urls:
return "high"
return "middle"
class ManagerFuncsInComfyUI(core.ManagerFuncs):
def get_current_preview_method(self):
if args.preview_method == latent_preview.LatentPreviewMethod.Auto:
@@ -358,6 +388,10 @@ async def fetch_updates(request):
@PromptServer.instance.routes.get("/customnode/update_all")
async def update_all(request):
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
try:
core.save_snapshot_with_postfix('autosave')
@@ -551,9 +585,9 @@ async def get_snapshot_list(request):
@PromptServer.instance.routes.get("/snapshot/remove")
async def remove_snapshot(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the remove feature. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
try:
target = request.rel_url.query["target"]
@@ -569,9 +603,9 @@ async def remove_snapshot(request):
@PromptServer.instance.routes.get("/snapshot/restore")
async def remove_snapshot(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the restore feature. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
try:
target = request.rel_url.query["target"]
@@ -737,12 +771,17 @@ def copy_set_active(files, is_disable, js_path_name='.'):
@PromptServer.instance.routes.post("/customnode/install")
async def install_custom_node(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the installation of custom nodes. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
json_data = await request.json()
risky_level = await get_risky_level(json_data['files'])
if not is_allowed_security_level(risky_level):
print(f"ERROR: This installation is not allowed in this security_level. Please contact the administrator.")
return web.Response(status=404)
install_type = json_data['install_type']
print(f"Install custom node '{json_data['title']}'")
@@ -779,9 +818,9 @@ async def install_custom_node(request):
@PromptServer.instance.routes.post("/customnode/fix")
async def fix_custom_node(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the fix feature. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
json_data = await request.json()
@@ -813,9 +852,9 @@ async def fix_custom_node(request):
@PromptServer.instance.routes.post("/customnode/install/git_url")
async def install_custom_node_git_url(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the installation of custom nodes. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('high'):
print(f"ERROR: To use this feature, you must set '--listen' to a local IP and set the security level to 'middle' or 'weak'. Please contact the administrator.")
return web.Response(status=403)
url = await request.text()
res = core.gitclone_install([url])
@@ -829,9 +868,9 @@ async def install_custom_node_git_url(request):
@PromptServer.instance.routes.post("/customnode/install/pip")
async def install_custom_node_git_url(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the installation of pip package. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('high'):
print(f"ERROR: To use this feature, you must set '--listen' to a local IP and set the security level to 'middle' or 'weak'. Please contact the administrator.")
return web.Response(status=403)
packages = await request.text()
core.pip_install(packages.split(' '))
@@ -841,9 +880,9 @@ async def install_custom_node_git_url(request):
@PromptServer.instance.routes.post("/customnode/uninstall")
async def uninstall_custom_node(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the uninstallation of custom nodes. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
json_data = await request.json()
@@ -869,6 +908,10 @@ async def uninstall_custom_node(request):
@PromptServer.instance.routes.post("/customnode/update")
async def update_custom_node(request):
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
json_data = await request.json()
install_type = json_data['install_type']
@@ -983,9 +1026,9 @@ manager_terminal_hook = ManagerTerminalHook()
@PromptServer.instance.routes.get("/manager/terminal")
async def terminal_mode(request):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the terminal feature. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('high'):
print(f"ERROR: To use this action, a security_level of `weak` is required. Please contact the administrator.")
return web.Response(status=403)
if "mode" in request.rel_url.query:
if request.rel_url.query['mode'] == 'true':
@@ -1110,9 +1153,9 @@ async def get_notice(request):
@PromptServer.instance.routes.get("/manager/reboot")
def restart(self):
if core.is_unsecure_features_disabled():
print(f"ERROR: The unsecure feature is disabled, restricting the reboot feature. Please contact the administrator.")
return web.Response(status=400)
if not is_allowed_security_level('middle'):
print(f"ERROR: To use this action, a security_level of `middle or below` is required. Please contact the administrator.")
return web.Response(status=403)
try:
sys.stdout.close_log()