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

@@ -23,7 +23,7 @@ sys.path.append(glob_path)
import cm_global
from manager_util import *
version = [2, 32, 7]
version = [2, 33]
version_str = f"V{version[0]}.{version[1]}" + (f'.{version[2]}' if len(version) > 2 else '')
comfyui_manager_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
@@ -205,7 +205,7 @@ def write_config():
'windows_selector_event_loop_policy': get_config()['windows_selector_event_loop_policy'],
'model_download_by_agent': get_config()['model_download_by_agent'],
'downgrade_blacklist': get_config()['downgrade_blacklist'],
'disable_unsecure_features': get_config()['disable_unsecure_features'],
'security_level': get_config()['security_level'],
}
with open(config_path, 'w') as configfile:
config.write(configfile)
@@ -217,6 +217,15 @@ def read_config():
config.read(config_path)
default_conf = config['default']
# policy migration: disable_unsecure_features -> security_level
if 'disable_unsecure_features' in default_conf:
if default_conf['disable_unsecure_features'].lower() == 'true':
security_level = 'strong'
else:
security_level = 'normal'
else:
security_level = default_conf['security_level'] if 'security_level' in default_conf else 'normal'
return {
'preview_method': default_conf['preview_method'] if 'preview_method' in default_conf else manager_funcs.get_current_preview_method(),
'badge_mode': default_conf['badge_mode'] if 'badge_mode' in default_conf else 'none',
@@ -231,7 +240,7 @@ def read_config():
'windows_selector_event_loop_policy': default_conf['windows_selector_event_loop_policy'].lower() == 'true' if 'windows_selector_event_loop_policy' in default_conf else False,
'model_download_by_agent': default_conf['model_download_by_agent'].lower() == 'true' if 'model_download_by_agent' in default_conf else False,
'downgrade_blacklist': default_conf['downgrade_blacklist'] if 'downgrade_blacklist' in default_conf else '',
'disable_unsecure_features': default_conf['disable_unsecure_features'].lower() == 'true' if 'disable_unsecure_features' in default_conf else False,
'security_level': security_level
}
except Exception:
@@ -249,7 +258,7 @@ def read_config():
'windows_selector_event_loop_policy': False,
'model_download_by_agent': False,
'downgrade_blacklist': '',
'disable_unsecure_features': False,
'security_level': 'normal',
}
@@ -1190,6 +1199,3 @@ def unzip(model_path):
os.remove(model_path)
return True
def is_unsecure_features_disabled():
return get_config()['disable_unsecure_features']

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()