From e964f267983073196247639c2fe17275eaff426d Mon Sep 17 00:00:00 2001 From: Hayden <48267247+hayden-fr@users.noreply.github.com> Date: Tue, 29 Apr 2025 17:30:36 +0800 Subject: [PATCH] Secure storage API keys (#181) * Migrate api key to private.key * Optimize API Key setting --- .gitignore | 3 ++ py/download.py | 65 ++++++++++++++++++++++- src/components/SettingApiKey.vue | 69 +++++++++++++++++++++++++ src/hooks/config.ts | 88 ++++++++++++++++++++++++++++++-- src/hooks/download.ts | 9 ++++ 5 files changed, 229 insertions(+), 5 deletions(-) create mode 100644 src/components/SettingApiKey.vue diff --git a/.gitignore b/.gitignore index 2585b6a..317eb59 100644 --- a/.gitignore +++ b/.gitignore @@ -197,3 +197,6 @@ web/ # config config/ + +# private info +private.key \ No newline at end of file diff --git a/py/download.py b/py/download.py index 178a9f9..77e7d1b 100644 --- a/py/download.py +++ b/py/download.py @@ -2,6 +2,7 @@ import os import uuid import time import requests +import base64 import folder_paths @@ -94,8 +95,68 @@ class TaskContent: } +class ApiKey: + + __store: dict[str, str] = {} + + def __init__(self): + self.__cache_file = os.path.join(config.extension_uri, "private.key") + + def init(self, request): + # Try to migrate api key from user setting + if not os.path.exists(self.__cache_file): + self.__store = { + "civitai": utils.get_setting_value(request, "api_key.civitai"), + "huggingface": utils.get_setting_value(request, "api_key.huggingface"), + } + self.__update__() + # Remove api key from user setting + utils.set_setting_value(request, "api_key.civitai", None) + utils.set_setting_value(request, "api_key.huggingface", None) + self.__store = utils.load_dict_pickle_file(self.__cache_file) + # Desensitization returns + result: dict[str, str] = {} + for key in self.__store: + v = self.__store[key] + if v is not None: + result[key] = v[:4] + "****" + v[-4:] + return result + + def get_value(self, key: str): + return self.__store.get(key, None) + + def set_value(self, key: str, value: str): + self.__store[key] = value + self.__update__() + + def __update__(self): + utils.save_dict_pickle_file(self.__cache_file, self.__store) + + class ModelDownload: + def __init__(self): + self.api_key = ApiKey() + def add_routes(self, routes): + @routes.post("/model-manager/download/init") + async def init_download(request): + """ + Init download setting. + """ + result = self.api_key.init(request) + return web.json_response({"success": True, "data": result}) + + @routes.post("/model-manager/download/setting") + async def set_download_setting(request): + """ + Set download setting. + """ + json_data = await request.json() + key = json_data.get("key", None) + value = json_data.get("value", None) + value = base64.b64decode(value).decode("utf-8") if value is not None else None + self.api_key.set_value(key, value) + return web.json_response({"success": True}) @routes.get("/model-manager/download/task") async def scan_download_tasks(request): @@ -331,12 +392,12 @@ class ModelDownload: download_platform = task_status.platform if download_platform == "civitai": - api_key = utils.get_setting_value(request, "api_key.civitai") + api_key = self.api_key.get_value("civitai") if api_key: headers["Authorization"] = f"Bearer {api_key}" elif download_platform == "huggingface": - api_key = utils.get_setting_value(request, "api_key.huggingface") + api_key = self.api_key.get_value("huggingface") if api_key: headers["Authorization"] = f"Bearer {api_key}" diff --git a/src/components/SettingApiKey.vue b/src/components/SettingApiKey.vue new file mode 100644 index 0000000..e7d0bb8 --- /dev/null +++ b/src/components/SettingApiKey.vue @@ -0,0 +1,69 @@ + + + diff --git a/src/hooks/config.ts b/src/hooks/config.ts index fadb9dc..cd9da55 100644 --- a/src/hooks/config.ts +++ b/src/hooks/config.ts @@ -1,6 +1,9 @@ +import SettingApiKey from 'components/SettingApiKey.vue' import SettingCardSize from 'components/SettingCardSize.vue' +import { request } from 'hooks/request' import { defineStore } from 'hooks/store' -import { app } from 'scripts/comfyAPI' +import { useToast } from 'hooks/toast' +import { $el, app } from 'scripts/comfyAPI' import { computed, onMounted, onUnmounted, readonly, ref, watch } from 'vue' import { useI18n } from 'vue-i18n' @@ -65,6 +68,7 @@ export const useConfig = defineStore('config', (store) => { }, }, flat: flatLayout, + apiKeyInfo: ref>({}), } watch(cardSizeFlag, (val) => { @@ -97,6 +101,84 @@ export const configSetting = { function useAddConfigSettings(store: import('hooks/store').StoreProvider) { const { t } = useI18n() + const { confirm } = useToast() + + const iconButton = (opt: { + icon: string + onClick: () => void | Promise + }) => { + return $el( + 'span.h-4.cursor-pointer', + { onclick: opt.onClick }, + $el(`i.${opt.icon.replace(/\s/g, '.')}`), + ) + } + + const setApiKey = async (key: string, setter: (val: string) => void) => { + store.dialog.open({ + key: `setting.api_key.${key}`, + title: t(`setting.api_key.${key}`), + content: SettingApiKey, + modal: true, + defaultSize: { + width: 500, + height: 200, + }, + contentProps: { + keyField: key, + setter: setter, + }, + }) + } + + const removeApiKey = async (key: string) => { + await new Promise((resolve, reject) => { + confirm.require({ + message: t('deleteAsk'), + header: 'Danger', + icon: 'pi pi-info-circle', + accept: () => resolve(true), + reject: reject, + }) + }) + await request('/download/setting', { + method: 'POST', + body: JSON.stringify({ key, value: null }), + }) + } + + const renderApiKey = (key: string) => { + return () => { + const apiKey = store.config.apiKeyInfo.value[key] || 'None' + const apiKeyDisplayEl = $el('div.text-sm.text-gray-500.flex-1', { + textContent: apiKey, + }) + + const setter = (val: string) => { + store.config.apiKeyInfo.value[key] = val + apiKeyDisplayEl.textContent = val || 'None' + } + return $el('div.flex.gap-4', [ + apiKeyDisplayEl, + iconButton({ + icon: 'pi pi-pencil text-blue-400', + onClick: () => { + setApiKey(key, setter) + }, + }), + iconButton({ + icon: 'pi pi-trash text-red-400', + onClick: async () => { + const value = store.config.apiKeyInfo.value[key] + if (value) { + await removeApiKey(key) + setter('') + } + }, + }), + ]) + } + } onMounted(() => { // API keys @@ -104,16 +186,16 @@ function useAddConfigSettings(store: import('hooks/store').StoreProvider) { id: 'ModelManager.APIKey.HuggingFace', category: [t('modelManager'), t('setting.apiKey'), 'HuggingFace'], name: 'HuggingFace API Key', - type: 'text', defaultValue: undefined, + type: renderApiKey('huggingface'), }) app.ui?.settings.addSetting({ id: 'ModelManager.APIKey.Civitai', category: [t('modelManager'), t('setting.apiKey'), 'Civitai'], name: 'Civitai API Key', - type: 'text', defaultValue: undefined, + type: renderApiKey('civitai'), }) const defaultCardSize = store.config.defaultCardSizeMap diff --git a/src/hooks/download.ts b/src/hooks/download.ts index e58fbb2..6afe39c 100644 --- a/src/hooks/download.ts +++ b/src/hooks/download.ts @@ -84,7 +84,16 @@ export const useDownload = defineStore('download', (store) => { }) }) + // Initial download settings + // Migrate API keys from user settings to private key + const init = async () => { + const res = await request('/download/init', { method: 'POST' }) + store.config.apiKeyInfo.value = res + } + onBeforeMount(() => { + init() + api.addEventListener('reconnected', () => { refresh() })