From 600dabb2a5b3cd4d1c787fa85756affd13278135 Mon Sep 17 00:00:00 2001 From: millerlight Date: Sun, 10 Aug 2025 16:39:10 +0200 Subject: [PATCH] whitelist check for files --- __init__.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 62bfe65..edb2cde 100644 --- a/__init__.py +++ b/__init__.py @@ -83,20 +83,37 @@ try: # POST /kvtools/peek {file_name|path, key} -> {"ok":true,"value": "..."} # reads a value from a JSON in BASE_DIR, no paths outside allowed. + # POST /kvtools/peek {file_name|path, key} -> {"ok":true,"value": "..."} + # Reads ONLY whitelisted (file,key) pairs that exist in the registry. @PromptServer.instance.routes.post("/kvtools/peek") async def kvtools_peek(request): try: data = await request.json() - except: + except Exception: data = {} file_name = data.get("file_name") or data.get("path") or "" key = data.get("key") or "" if not file_name or not key: return web.json_response({"ok": False, "error": "missing file_name or key"}, status=400) + # --- load registry (prefer on-disk JSON; fallback to fresh scan) + try: + with open(REG_PATH, "r", encoding="utf-8") as fh: + reg = json.load(fh) + except Exception: + reg = _scan() + + files = (reg or {}).get("files", {}) + # normalize to registry key (json basename with extension) json_name = os.path.basename(str(file_name).strip()) + # require: listed in registry, key exists + entry = files.get(json_name) + if not entry or "keys" not in entry or str(key) not in entry["keys"]: + return web.json_response({"ok": False, "error": "not whitelisted (file/key)"}, status=403) + + # path under BASE_DIR, file must exist json_path = os.path.join(BASE_DIR, json_name) - if not (json_path.startswith(BASE_DIR) and os.path.isfile(json_path)): + if not (os.path.isfile(json_path) and os.path.abspath(json_path).startswith(os.path.abspath(BASE_DIR))): return web.json_response({"ok": False, "error": "invalid file"}, status=400) try: @@ -109,6 +126,7 @@ try: except Exception as e: return web.json_response({"ok": False, "error": str(e)}, status=500) + # GET /kvtools/image?file=&key=&ext=png # Serves ONLY images from /custom_kv_stores/images//. @PromptServer.instance.routes.get("/kvtools/image")