diff --git a/src/LLMProviders/brevilabsClient.ts b/src/LLMProviders/brevilabsClient.ts index 8fd85166..baeadb27 100644 --- a/src/LLMProviders/brevilabsClient.ts +++ b/src/LLMProviders/brevilabsClient.ts @@ -94,7 +94,7 @@ export class BrevilabsClient { method, headers: { "Content-Type": "application/json", - Authorization: `Bearer ${getDecryptedKey(getSettings().plusLicenseKey)}`, + Authorization: `Bearer ${await getDecryptedKey(getSettings().plusLicenseKey)}`, }, ...(method === "POST" && { body: JSON.stringify(body) }), }); diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts index 2e1e4056..9510c8cc 100644 --- a/src/LLMProviders/chainManager.ts +++ b/src/LLMProviders/chainManager.ts @@ -53,8 +53,12 @@ export default class ChainManager { this.memoryManager = MemoryManager.getInstance(); this.chatModelManager = ChatModelManager.getInstance(); this.promptManager = PromptManager.getInstance(); - this.createChainWithNewModel(); - subscribeToModelKeyChange(() => this.createChainWithNewModel()); + + // Initialize async operations + this.initialize(); + + // Set up subscriptions + subscribeToModelKeyChange(async () => await this.createChainWithNewModel()); subscribeToChainTypeChange(() => this.setChain(getChainType(), { refreshIndex: @@ -63,7 +67,11 @@ export default class ChainManager { getChainType() === ChainType.COPILOT_PLUS_CHAIN), }) ); - subscribeToSettingsChange(() => this.createChainWithNewModel()); + subscribeToSettingsChange(async () => await this.createChainWithNewModel()); + } + + private async initialize() { + await this.createChainWithNewModel(); } static getChain(): RunnableSequence { @@ -102,7 +110,7 @@ export default class ChainManager { * Update the active model and create a new chain with the specified model * name. */ - createChainWithNewModel(): void { + async createChainWithNewModel(): Promise { let newModelKey = getModelKey(); try { let customModel = findCustomModel(newModelKey, getSettings().activeModels); @@ -112,7 +120,7 @@ export default class ChainManager { customModel = BUILTIN_CHAT_MODELS[0]; newModelKey = customModel.name + "|" + customModel.provider; } - this.chatModelManager.setChatModel(customModel); + await this.chatModelManager.setChatModel(customModel); // Must update the chatModel for chain because ChainFactory always // retrieves the old chain without the chatModel change if it exists! // Create a new chain with the new chatModel diff --git a/src/LLMProviders/chatModelManager.ts b/src/LLMProviders/chatModelManager.ts index 60dad5ca..728452aa 100644 --- a/src/LLMProviders/chatModelManager.ts +++ b/src/LLMProviders/chatModelManager.ts @@ -70,7 +70,7 @@ export default class ChatModelManager { return ChatModelManager.instance; } - private getModelConfig(customModel: CustomModel): ModelConfig { + private async getModelConfig(customModel: CustomModel): Promise { const settings = getSettings(); // Check if the model starts with "o1" @@ -90,7 +90,7 @@ export default class ChatModelManager { } = { [ChatModelProviders.OPENAI]: { modelName: modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey), + openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey), configuration: { baseURL: customModel.baseUrl, fetch: customModel.enableCors ? safeFetch : undefined, @@ -100,7 +100,7 @@ export default class ChatModelManager { ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature), }, [ChatModelProviders.ANTHROPIC]: { - anthropicApiKey: getDecryptedKey(customModel.apiKey || settings.anthropicApiKey), + anthropicApiKey: await getDecryptedKey(customModel.apiKey || settings.anthropicApiKey), modelName: modelName, anthropicApiUrl: customModel.baseUrl, clientOptions: { @@ -110,7 +110,7 @@ export default class ChatModelManager { }, }, [ChatModelProviders.AZURE_OPENAI]: { - azureOpenAIApiKey: getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey), + azureOpenAIApiKey: await getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey), azureOpenAIApiInstanceName: settings.azureOpenAIApiInstanceName, azureOpenAIApiDeploymentName: settings.azureOpenAIApiDeploymentName, azureOpenAIApiVersion: settings.azureOpenAIApiVersion, @@ -121,11 +121,11 @@ export default class ChatModelManager { ...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature), }, [ChatModelProviders.COHEREAI]: { - apiKey: getDecryptedKey(customModel.apiKey || settings.cohereApiKey), + apiKey: await getDecryptedKey(customModel.apiKey || settings.cohereApiKey), model: modelName, }, [ChatModelProviders.GOOGLE]: { - apiKey: getDecryptedKey(customModel.apiKey || settings.googleApiKey), + apiKey: await getDecryptedKey(customModel.apiKey || settings.googleApiKey), modelName: modelName, safetySettings: [ { @@ -149,14 +149,14 @@ export default class ChatModelManager { }, [ChatModelProviders.OPENROUTERAI]: { modelName: modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openRouterAiApiKey), + openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openRouterAiApiKey), configuration: { baseURL: customModel.baseUrl || "https://openrouter.ai/api/v1", fetch: customModel.enableCors ? safeFetch : undefined, }, }, [ChatModelProviders.GROQ]: { - apiKey: getDecryptedKey(customModel.apiKey || settings.groqApiKey), + apiKey: await getDecryptedKey(customModel.apiKey || settings.groqApiKey), modelName: modelName, }, [ChatModelProviders.OLLAMA]: { @@ -177,7 +177,7 @@ export default class ChatModelManager { }, [ChatModelProviders.OPENAI_FORMAT]: { modelName: modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey), + openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey), configuration: { baseURL: customModel.baseUrl, fetch: customModel.enableCors ? safeFetch : undefined, @@ -251,7 +251,7 @@ export default class ChatModelManager { return ChatModelManager.chatModel; } - setChatModel(model: CustomModel): void { + async setChatModel(model: CustomModel): Promise { const modelKey = `${model.name}|${model.provider}`; if (!ChatModelManager.modelMap.hasOwnProperty(modelKey)) { throw new Error(`No model found for: ${modelKey}`); @@ -266,7 +266,7 @@ export default class ChatModelManager { throw new Error(errorMessage); } - const modelConfig = this.getModelConfig(model); + const modelConfig = await this.getModelConfig(model); setModelKey(`${model.name}|${model.provider}`); try { @@ -312,7 +312,7 @@ export default class ChatModelManager { async ping(model: CustomModel): Promise { const tryPing = async (enableCors: boolean) => { const modelToTest = { ...model, enableCors }; - const modelConfig = this.getModelConfig(modelToTest); + const modelConfig = await this.getModelConfig(modelToTest); // eslint-disable-next-line @typescript-eslint/no-unused-vars const { streaming, temperature, ...pingConfig } = modelConfig; pingConfig.maxTokens = 10; diff --git a/src/LLMProviders/embeddingManager.ts b/src/LLMProviders/embeddingManager.ts index c6245baf..9ca26ded 100644 --- a/src/LLMProviders/embeddingManager.ts +++ b/src/LLMProviders/embeddingManager.ts @@ -131,7 +131,7 @@ export default class EmbeddingManager { })[0]; } - getEmbeddingsAPI(): Embeddings { + async getEmbeddingsAPI(): Promise { const { embeddingModelKey } = getSettings(); if (!EmbeddingManager.modelMap.hasOwnProperty(embeddingModelKey)) { @@ -146,7 +146,7 @@ export default class EmbeddingManager { } const customModel = this.getCustomModel(embeddingModelKey); - const config = this.getEmbeddingConfig(customModel); + const config = await this.getEmbeddingConfig(customModel); try { EmbeddingManager.embeddingModel = new selectedModel.EmbeddingConstructor(config); @@ -158,7 +158,7 @@ export default class EmbeddingManager { } } - private getEmbeddingConfig(customModel: CustomModel): any { + private async getEmbeddingConfig(customModel: CustomModel): Promise { const settings = getSettings(); const modelName = customModel.name; @@ -187,7 +187,7 @@ export default class EmbeddingManager { } = { [EmbeddingModelProviders.COPILOT_PLUS]: { modelName, - apiKey: getDecryptedKey(settings.plusLicenseKey), + apiKey: await getDecryptedKey(settings.plusLicenseKey), timeout: 10000, configuration: { baseURL: BREVILABS_API_BASE_URL, @@ -196,7 +196,7 @@ export default class EmbeddingManager { }, [EmbeddingModelProviders.COPILOT_PLUS_JINA]: { model: modelName, - apiKey: getDecryptedKey(settings.plusLicenseKey), + apiKey: await getDecryptedKey(settings.plusLicenseKey), timeout: 10000, batchSize: 128, dimensions: 512, @@ -207,7 +207,7 @@ export default class EmbeddingManager { }, [EmbeddingModelProviders.OPENAI]: { modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey), + openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey), timeout: 10000, configuration: { baseURL: customModel.baseUrl, @@ -216,14 +216,14 @@ export default class EmbeddingManager { }, [EmbeddingModelProviders.COHEREAI]: { model: modelName, - apiKey: getDecryptedKey(customModel.apiKey || settings.cohereApiKey), + apiKey: await getDecryptedKey(customModel.apiKey || settings.cohereApiKey), }, [EmbeddingModelProviders.GOOGLE]: { modelName: modelName, - apiKey: getDecryptedKey(settings.googleApiKey), + apiKey: await getDecryptedKey(settings.googleApiKey), }, [EmbeddingModelProviders.AZURE_OPENAI]: { - azureOpenAIApiKey: getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey), + azureOpenAIApiKey: await getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey), azureOpenAIApiInstanceName: settings.azureOpenAIApiInstanceName, azureOpenAIApiDeploymentName: settings.azureOpenAIApiEmbeddingDeploymentName, azureOpenAIApiVersion: settings.azureOpenAIApiVersion, @@ -239,7 +239,7 @@ export default class EmbeddingManager { }, [EmbeddingModelProviders.LM_STUDIO]: { modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || "default-key"), + openAIApiKey: await getDecryptedKey(customModel.apiKey || "default-key"), configuration: { baseURL: customModel.baseUrl || "http://localhost:1234/v1", fetch: customModel.enableCors ? safeFetch : undefined, @@ -247,7 +247,7 @@ export default class EmbeddingManager { }, [EmbeddingModelProviders.OPENAI_FORMAT]: { modelName, - openAIApiKey: getDecryptedKey(customModel.apiKey || ""), + openAIApiKey: await getDecryptedKey(customModel.apiKey || ""), configuration: { baseURL: customModel.baseUrl, fetch: customModel.enableCors ? safeFetch : undefined, @@ -265,7 +265,7 @@ export default class EmbeddingManager { async ping(model: CustomModel): Promise { const tryPing = async (enableCors: boolean) => { const modelToTest = { ...model, enableCors }; - const config = this.getEmbeddingConfig(modelToTest); + const config = await this.getEmbeddingConfig(modelToTest); const testModel = new (this.getProviderConstructor(modelToTest))(config); await testModel.embedQuery("test"); }; diff --git a/src/encryptionService.ts b/src/encryptionService.ts index f29e387e..c28ddeda 100644 --- a/src/encryptionService.ts +++ b/src/encryptionService.ts @@ -14,10 +14,27 @@ function getSafeStorage() { return safeStorageInternal; } +// Add new prefixes to distinguish encryption methods +const DESKTOP_PREFIX = "enc_desk_"; +const WEBCRYPTO_PREFIX = "enc_web_"; +// Keep old prefix for backward compatibility const ENCRYPTION_PREFIX = "enc_"; const DECRYPTION_PREFIX = "dec_"; -export function encryptAllKeys(settings: Readonly): Readonly { +// Add these constants for the Web Crypto implementation +const ENCRYPTION_KEY = new TextEncoder().encode("obsidian-copilot-v1"); +const ALGORITHM = { name: "AES-GCM", iv: new Uint8Array(12) }; + +async function getEncryptionKey(): Promise { + return await crypto.subtle.importKey("raw", ENCRYPTION_KEY, ALGORITHM.name, false, [ + "encrypt", + "decrypt", + ]); +} + +export async function encryptAllKeys( + settings: Readonly +): Promise> { if (!settings.enableEncryption) { return settings; } @@ -28,20 +45,22 @@ export function encryptAllKeys(settings: Readonly): Readonly ({ - ...model, - apiKey: getEncryptedKey(model.apiKey || ""), - })); + newSettings.activeModels = await Promise.all( + settings.activeModels.map(async (model) => ({ + ...model, + apiKey: await getEncryptedKey(model.apiKey || ""), + })) + ); } return newSettings; } -export function getEncryptedKey(apiKey: string): string { +export async function getEncryptedKey(apiKey: string): Promise { if (!apiKey || apiKey.startsWith(ENCRYPTION_PREFIX)) { return apiKey; } @@ -50,20 +69,25 @@ export function getEncryptedKey(apiKey: string): string { apiKey = apiKey.replace(DECRYPTION_PREFIX, ""); } - if (getSafeStorage() && getSafeStorage().isEncryptionAvailable()) { - // Convert the encrypted buffer to a Base64 string and prepend the prefix - const encryptedBuffer = getSafeStorage().encryptString(apiKey) as Buffer; - // Convert the encrypted buffer to a Base64 string and prepend the prefix - return ENCRYPTION_PREFIX + encryptedBuffer.toString("base64"); - } else { - // Simple fallback for mobile (just for demonstration) - const encoder = new TextEncoder(); - const data = encoder.encode(apiKey); - return ENCRYPTION_PREFIX + arrayBufferToBase64(data); + try { + // Try desktop encryption first + if (getSafeStorage()?.isEncryptionAvailable()) { + const encryptedBuffer = getSafeStorage().encryptString(apiKey) as Buffer; + return DESKTOP_PREFIX + encryptedBuffer.toString("base64"); + } + + // Fallback to Web Crypto API + const key = await getEncryptionKey(); + const encodedData = new TextEncoder().encode(apiKey); + const encryptedData = await crypto.subtle.encrypt(ALGORITHM, key, encodedData); + return WEBCRYPTO_PREFIX + arrayBufferToBase64(encryptedData); + } catch (error) { + console.error("Encryption failed:", error); + return apiKey; } } -export function getDecryptedKey(apiKey: string): string { +export async function getDecryptedKey(apiKey: string): Promise { if (!apiKey || isPlainText(apiKey)) { return apiKey; } @@ -71,17 +95,41 @@ export function getDecryptedKey(apiKey: string): string { return apiKey.replace(DECRYPTION_PREFIX, ""); } + // Handle different encryption methods + if (apiKey.startsWith(DESKTOP_PREFIX)) { + const base64Data = apiKey.replace(DESKTOP_PREFIX, ""); + const buffer = Buffer.from(base64Data, "base64"); + return getSafeStorage().decryptString(buffer) as string; + } + + if (apiKey.startsWith(WEBCRYPTO_PREFIX)) { + const base64Data = apiKey.replace(WEBCRYPTO_PREFIX, ""); + const key = await getEncryptionKey(); + const encryptedData = base64ToArrayBuffer(base64Data); + const decryptedData = await crypto.subtle.decrypt(ALGORITHM, key, encryptedData); + return new TextDecoder().decode(decryptedData); + } + + // Legacy support for old enc_ prefix const base64Data = apiKey.replace(ENCRYPTION_PREFIX, ""); try { - if (getSafeStorage() && getSafeStorage().isEncryptionAvailable()) { - const buffer = Buffer.from(base64Data, "base64"); - return getSafeStorage().decryptString(buffer) as string; - } else { - // Simple fallback for mobile (just for demonstration) - const data = base64ToArrayBuffer(base64Data); - const decoder = new TextDecoder(); - return decoder.decode(data); + // Try desktop decryption first + if (getSafeStorage()?.isEncryptionAvailable()) { + try { + const buffer = Buffer.from(base64Data, "base64"); + return getSafeStorage().decryptString(buffer) as string; + } catch { + // Silent catch is intentional - if desktop decryption fails, + // it means this key was likely encrypted with Web Crypto. + // We'll fall through to the Web Crypto decryption below. + } } + + // Fallback to Web Crypto API + const key = await getEncryptionKey(); + const encryptedData = base64ToArrayBuffer(base64Data); + const decryptedData = await crypto.subtle.decrypt(ALGORITHM, key, encryptedData); + return new TextDecoder().decode(decryptedData); } catch (err) { console.error("Decryption failed:", err); return "Copilot failed to decrypt API keys!"; diff --git a/src/main.ts b/src/main.ts index e58a7092..3612c6ef 100644 --- a/src/main.ts +++ b/src/main.ts @@ -54,12 +54,12 @@ export default class CopilotPlugin extends Plugin { async onload(): Promise { await this.loadSettings(); - this.settingsUnsubscriber = subscribeToSettingsChange(() => { + this.settingsUnsubscriber = subscribeToSettingsChange(async () => { const settings = getSettings(); if (settings.enableEncryption) { - this.saveData(encryptAllKeys(settings)); + await this.saveData(await encryptAllKeys(settings)); } else { - this.saveData(settings); + await this.saveData(settings); } registerBuiltInCommands(this); }); diff --git a/src/search/dbOperations.ts b/src/search/dbOperations.ts index d7cdae11..a5f7a840 100644 --- a/src/search/dbOperations.ts +++ b/src/search/dbOperations.ts @@ -45,7 +45,7 @@ export class DBOperations { this.oramaDb = undefined; } else if (Platform.isMobile && !settings.disableIndexOnMobile && !this.oramaDb) { // Re-initialize DB if mobile setting is enabled - await this.initializeDB(EmbeddingsManager.getInstance().getEmbeddingsAPI()); + await this.initializeDB(await EmbeddingsManager.getInstance().getEmbeddingsAPI()); } // Handle index sync setting change @@ -55,7 +55,7 @@ export class DBOperations { console.log("Path change detected, reinitializing database..."); this.dbPath = newPath; await this.initializeChunkedStorage(); - await this.initializeDB(EmbeddingsManager.getInstance().getEmbeddingsAPI()); + await this.initializeDB(await EmbeddingsManager.getInstance().getEmbeddingsAPI()); console.log("Database reinitialized with new path:", newPath); } }); diff --git a/src/search/hybridRetriever.ts b/src/search/hybridRetriever.ts index 9992470a..eb9c2206 100644 --- a/src/search/hybridRetriever.ts +++ b/src/search/hybridRetriever.ts @@ -346,7 +346,8 @@ export class HybridRetriever extends BaseRetriever { } private async convertQueryToVector(query: string): Promise { - const vector = await EmbeddingManager.getInstance().getEmbeddingsAPI().embedQuery(query); + const embeddingsAPI = await EmbeddingManager.getInstance().getEmbeddingsAPI(); + const vector = await embeddingsAPI.embedQuery(query); if (vector.length === 0) { throw new Error("Query embedding returned an empty vector"); } diff --git a/src/search/indexOperations.ts b/src/search/indexOperations.ts index 1ca7a116..7b09b4f5 100644 --- a/src/search/indexOperations.ts +++ b/src/search/indexOperations.ts @@ -52,7 +52,7 @@ export class IndexOperations { let rateLimitNoticeShown = false; try { - const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI(); + const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingInstance) { console.error("Embedding instance not found."); return 0; @@ -169,7 +169,7 @@ export class IndexOperations { fileInfo: any; }> > { - const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI(); + const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingInstance) { console.error("Embedding instance not found."); return []; @@ -444,7 +444,7 @@ export class IndexOperations { public async reindexFile(file: TFile): Promise { try { - const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI(); + const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI(); if (!embeddingInstance) { return; } diff --git a/src/search/vectorStoreManager.ts b/src/search/vectorStoreManager.ts index d7f11c78..3af5b1ea 100644 --- a/src/search/vectorStoreManager.ts +++ b/src/search/vectorStoreManager.ts @@ -48,7 +48,7 @@ export default class VectorStoreManager { const oldPath = this.dbOps.getCurrentDbPath(); if (oldPath !== newPath) { - await this.dbOps.initializeDB(this.embeddingsManager.getEmbeddingsAPI()); + await this.dbOps.initializeDB(await this.embeddingsManager.getEmbeddingsAPI()); } } @@ -71,7 +71,7 @@ export default class VectorStoreManager { let retries = 3; while (retries > 0) { try { - await this.dbOps.initializeDB(this.embeddingsManager.getEmbeddingsAPI()); + await this.dbOps.initializeDB(await this.embeddingsManager.getEmbeddingsAPI()); break; } catch (error) { if ( @@ -84,6 +84,10 @@ export default class VectorStoreManager { continue; } } + new Notice( + "Failed to initialize vector store. Please make sure you have a valid API key " + + "for your embedding model and restart the plugin." + ); console.error("Failed to initialize vector store:", error); break; } @@ -109,7 +113,7 @@ export default class VectorStoreManager { public async clearIndex(): Promise { await this.waitForInitialization(); - await this.dbOps.clearIndex(this.embeddingsManager.getEmbeddingsAPI()); + await this.dbOps.clearIndex(await this.embeddingsManager.getEmbeddingsAPI()); } public async garbageCollectVectorStore(): Promise { diff --git a/tests/encryptionService.test.ts b/tests/encryptionService.test.ts index 693c744e..ecd25ffc 100644 --- a/tests/encryptionService.test.ts +++ b/tests/encryptionService.test.ts @@ -1,20 +1,50 @@ -import { getDecryptedKey, getEncryptedKey, encryptAllKeys } from "@/encryptionService"; -import { Platform } from "obsidian"; -import { type CopilotSettings } from "@/settings/model"; +import { TextDecoder, TextEncoder } from "util"; -// Mocking Electron's safeStorage -jest.mock("electron", () => { - return { - remote: { - safeStorage: { - encryptString: jest.fn().mockImplementation((text) => `encrypted_${text}`), - decryptString: jest - .fn() - .mockImplementation((buffer) => buffer.toString().replace("encrypted_", "")), - isEncryptionAvailable: jest.fn().mockReturnValue(true), - }, +// Mock electron module with proper types +const mockElectron = { + remote: { + safeStorage: { + encryptString: jest.fn().mockImplementation((text) => Buffer.from(`${text}_encrypted`)), + decryptString: jest + .fn() + .mockImplementation((buffer) => buffer.toString().replace("_encrypted", "")), + isEncryptionAvailable: jest.fn().mockReturnValue(true), }, - }; + }, +}; + +jest.mock("electron", () => mockElectron); + +global.TextEncoder = TextEncoder as any; +global.TextDecoder = TextDecoder as any; + +// Now we can import our modules +import { encryptAllKeys, getDecryptedKey, getEncryptedKey } from "@/encryptionService"; +import { type CopilotSettings } from "@/settings/model"; +import { Platform } from "obsidian"; + +// Mock window.btoa and window.atob for base64 encoding/decoding +global.btoa = jest.fn().mockImplementation((str) => Buffer.from(str).toString("base64")); +global.atob = jest.fn().mockImplementation((str) => Buffer.from(str, "base64").toString()); + +const mockSubtle = { + importKey: jest.fn().mockResolvedValue("mockCryptoKey"), + encrypt: jest.fn().mockImplementation((algorithm, key, data) => { + const originalText = new TextDecoder().decode(data); + const encryptedText = `${originalText}_encrypted`; + return Promise.resolve(new TextEncoder().encode(encryptedText).buffer); + }), + decrypt: jest.fn().mockImplementation((algorithm, key, data) => { + const encryptedText = new TextDecoder().decode(new Uint8Array(data)); + const originalText = encryptedText.replace("_encrypted", ""); + return Promise.resolve(new TextEncoder().encode(originalText).buffer); + }), +}; + +// Mock crypto.subtle instead of the entire crypto object +Object.defineProperty(global.crypto, "subtle", { + value: mockSubtle, + configurable: true, }); describe("Platform-specific Tests", () => { @@ -39,52 +69,61 @@ describe("EncryptionService", () => { }); describe("getEncryptedKey", () => { - it("should encrypt an API key", () => { + it("should encrypt an API key", async () => { const apiKey = "testApiKey"; - const encryptedKey = getEncryptedKey(apiKey); - expect(encryptedKey).toBe(`enc_encrypted_${apiKey}`); + const encryptedKey = await getEncryptedKey(apiKey); + // The key is base64 encoded, so we should expect that format + expect(encryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/); + // Verify we can decrypt it back + const decryptedKey = await getDecryptedKey(encryptedKey); + expect(decryptedKey).toBe(apiKey); }); - it("should return the original key if already encrypted", () => { + it("should return the original key if already encrypted", async () => { const apiKey = "enc_testApiKey"; - const encryptedKey = getEncryptedKey(apiKey); + const encryptedKey = await getEncryptedKey(apiKey); expect(encryptedKey).toBe(apiKey); }); }); describe("getDecryptedKey", () => { - it("should decrypt an encrypted API key", () => { + it("should decrypt an encrypted API key", async () => { const apiKey = "testApiKey"; - const mockEncryptedKey = `encrypted_${apiKey}`; - const base64Encoded = Buffer.from(mockEncryptedKey).toString("base64"); - const encryptedKey = `enc_${base64Encoded}`; - - const decryptedKey = getDecryptedKey(encryptedKey); + const encryptedKey = await getEncryptedKey(apiKey); + const decryptedKey = await getDecryptedKey(encryptedKey); expect(decryptedKey).toBe(apiKey); }); - it("should return the original key if it is in plain text", () => { + it("should return the original key if it is in plain text", async () => { const apiKey = "testApiKey"; - const decryptedKey = getDecryptedKey(apiKey); + const decryptedKey = await getDecryptedKey(apiKey); expect(decryptedKey).toBe(apiKey); }); }); describe("encryptAllKeys", () => { - it('should encrypt all keys containing "apikey"', () => { - const newSettings = encryptAllKeys({ + it("should encrypt all keys containing 'apikey'", async () => { + const settings = { enableEncryption: true, openAIApiKey: "testApiKey", cohereApiKey: "anotherTestApiKey", userSystemPrompt: "shouldBeIgnored", - } as unknown as CopilotSettings); - expect(newSettings.openAIApiKey).toBe("enc_encrypted_testApiKey"); - expect(newSettings.cohereApiKey).toBe("enc_encrypted_anotherTestApiKey"); + } as unknown as CopilotSettings; + + const newSettings = await encryptAllKeys(settings); + expect(newSettings.openAIApiKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/); + expect(newSettings.cohereApiKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/); expect(newSettings.userSystemPrompt).toBe("shouldBeIgnored"); + + // Verify we can decrypt the keys back + const decryptedOpenAI = await getDecryptedKey(newSettings.openAIApiKey); + const decryptedCohere = await getDecryptedKey(newSettings.cohereApiKey); + expect(decryptedOpenAI).toBe("testApiKey"); + expect(decryptedCohere).toBe("anotherTestApiKey"); }); - it("should not encrypt keys when encryption is not enabled", () => { - const newSettings = encryptAllKeys({ + it("should not encrypt keys when encryption is not enabled", async () => { + const newSettings = await encryptAllKeys({ enableEncryption: false, openAIApiKey: "testApiKey", cohereApiKey: "anotherTestApiKey", @@ -96,3 +135,58 @@ describe("EncryptionService", () => { }); }); }); + +describe("Cross-platform compatibility", () => { + let originalConsoleError: typeof console.error; + + beforeEach(() => { + jest.clearAllMocks(); + // Save original console.error + originalConsoleError = console.error; + // Mock console.error to suppress expected encryption fallback messages + console.error = jest.fn(); + }); + + afterEach(() => { + // Restore original console.error + console.error = originalConsoleError; + }); + + it("should encrypt and decrypt consistently on mobile", async () => { + // Mock as mobile by making safeStorage unavailable + mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(false); + + const originalKey = "testApiKey"; + const encryptedKey = await getEncryptedKey(originalKey); + expect(encryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/); + + // Reset the mock counts before decryption + mockSubtle.encrypt.mockClear(); + mockSubtle.decrypt.mockClear(); + + const decryptedKey = await getDecryptedKey(encryptedKey); + expect(decryptedKey).toBe(originalKey); + + // On mobile, we should use Web Crypto API for decryption + expect(mockSubtle.decrypt).toHaveBeenCalled(); + }); + + it("should be able to decrypt mobile-encrypted keys on desktop", async () => { + // First encrypt on mobile + mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(false); + + const originalKey = "testApiKey"; + const mobileEncryptedKey = await getEncryptedKey(originalKey); + expect(mobileEncryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/); + expect(mockSubtle.encrypt).toHaveBeenCalled(); + + // Reset the mock counts before desktop decryption + mockSubtle.encrypt.mockClear(); + mockSubtle.decrypt.mockClear(); + + // Then decrypt on desktop + mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(true); + const decryptedKey = await getDecryptedKey(mobileEncryptedKey); + expect(decryptedKey).toBe(originalKey); + }); +});