From 066f448766eca202ad3d15221ff99a74d103cd7f Mon Sep 17 00:00:00 2001 From: Jakob Date: Mon, 15 Jul 2024 15:16:33 -0700 Subject: [PATCH] refactor: Update CredentialManager implementation --- .../credentialsManager/credentialsManager.ts | 209 ++++++++++-------- 1 file changed, 115 insertions(+), 94 deletions(-) diff --git a/packages/server/agents-v2/src/lib/services/credentialsManager/credentialsManager.ts b/packages/server/agents-v2/src/lib/services/credentialsManager/credentialsManager.ts index f483c341c4..abd9dd03f9 100644 --- a/packages/server/agents-v2/src/lib/services/credentialsManager/credentialsManager.ts +++ b/packages/server/agents-v2/src/lib/services/credentialsManager/credentialsManager.ts @@ -1,17 +1,16 @@ -import { prismaCore } from '@magickml/server-db' -import { decrypt, encrypt, PluginCredential } from '@magickml/credentials' -import { CREDENTIALS_ENCRYPTION_KEY } from '@magickml/server-config' -import { ICredentialManager } from '../../interfaces/ICredentialsManager' - -export type CredentialsType< - T extends object = Record -> = T -export class CredentialManager> - implements ICredentialManager -{ +import { CREDENTIALS_ENCRYPTION_KEY } from '../../../../../config/src' +import { decrypt, encrypt } from '../../../../../credentials/src' +import { prismaCore } from '../../../../../db/src' +import { + CredentialKeyValuePair, + ICredentialManager, + Credential, +} from '../../interfaces/ICredentialsManager' + +export class CredentialManager implements ICredentialManager { protected projectId: string protected agentId: string - protected currentCredentials: CredentialsType | undefined + protected cachedCredentials: CredentialKeyValuePair[] = [] constructor(agentId: string, projectId: string) { this.projectId = projectId @@ -19,16 +18,27 @@ export class CredentialManager> } async init(): Promise { - await this.update() + await this.refreshCredentialsCache() } - async update(): Promise { + async getCredentials(): Promise { + if (this.cachedCredentials.length === 0) { + await this.refreshCredentialsCache() + } + const credentials = this.cachedCredentials.map(credential => ({ + name: credential.name, + value: credential.value, + serviceType: 'core', + })) + return credentials + } + + async refreshCredentialsCache(): Promise { try { const creds = await prismaCore.agent_credentials.findMany({ where: { agentId: this.agentId, credentials: { - serviceType: 'core', projectId: this.projectId, }, }, @@ -42,29 +52,28 @@ export class CredentialManager> }, }) - const credentials = creds.reduce((acc, credential) => { - // @ts-ignore - acc[credential.credentials.name] = decrypt( + const credentials = creds.map(credential => ({ + name: credential.credentials.name, + value: decrypt( credential.credentials.value, CREDENTIALS_ENCRYPTION_KEY - ) - return acc - }, {}) + ), + })) as CredentialKeyValuePair[] - this.currentCredentials = credentials as CredentialsType + this.cachedCredentials = credentials } catch (error) { throw new Error( - `Error updating plugin credentials in agent: ${this.agentId}: ${error}` + `Error fetching plugin credentials for agent: ${this.agentId}: ${error}` ) } } - getCredentials(): CredentialsType | undefined { - return this.currentCredentials - } + getCredential(name: string): string | undefined { + const credential = this.cachedCredentials.find(cred => cred.name === name) - getCredential(name: keyof T): T[keyof T] | undefined { - return this.currentCredentials ? this.currentCredentials[name] : undefined + return credential + ? decrypt(credential.value, CREDENTIALS_ENCRYPTION_KEY) + : undefined } async getCustomCredential(name: string): Promise { @@ -94,87 +103,99 @@ export class CredentialManager> ) } - async addCredential( - credential: Partial, - pluginCredential: PluginCredential - ): Promise { - for (const [name, value] of Object.entries(credential)) { - // First, try to find an existing credential - const existingCredential = await prismaCore.credentials.findFirst({ - where: { - name: name, - serviceType: pluginCredential.serviceType, - projectId: this.projectId, - }, - }) - - if (existingCredential) { - // Update the existing credential - await prismaCore.credentials.update({ - where: { id: existingCredential.id }, - data: { - value: encrypt(value as string, CREDENTIALS_ENCRYPTION_KEY), - credentialType: pluginCredential.credentialType, - description: pluginCredential.description, - pluginName: pluginCredential.pluginName, - }, - }) - } else { - // Create a new credential - const createdCredential = await prismaCore.credentials.create({ - data: { - name: name, - value: encrypt(value as string, CREDENTIALS_ENCRYPTION_KEY), - serviceType: pluginCredential.serviceType, - projectId: this.projectId, - credentialType: pluginCredential.credentialType, - description: pluginCredential.description, - pluginName: pluginCredential.pluginName, - }, - }) + async addCredential(credential: Credential): Promise<{ id: string }> { + const existingCredential = await prismaCore.credentials.findFirst({ + where: { + name: credential.name, + serviceType: credential?.serviceType, + projectId: this.projectId, + }, + }) - // Link the new credential to the agent - await prismaCore.agent_credentials.create({ - data: { - agentId: this.agentId, - credentialId: createdCredential.id, - }, - }) - } + if (existingCredential) { + throw new Error(`Credential ${credential.name} already exists`) } - await this.update() - } - async deleteCredential(name: keyof T): Promise { - // First, delete the agent_credentials link - await prismaCore.agent_credentials.deleteMany({ - where: { + const createdCredential = await prismaCore.credentials.create({ + data: { + name: credential.name, + value: encrypt(credential.value, CREDENTIALS_ENCRYPTION_KEY), + serviceType: credential.serviceType, + projectId: this.projectId, + credentialType: credential.credentialType, + description: credential.description, + }, + }) + + await prismaCore.agent_credentials.create({ + data: { agentId: this.agentId, - credentials: { - name: name as string, - projectId: this.projectId, - }, + credentialId: createdCredential.id, }, }) - // Then, delete the credential itself - await prismaCore.credentials.deleteMany({ + await this.refreshCredentialsCache() + return { id: createdCredential.id } + } + + async updateCredential(credential: Credential): Promise { + const existingCredential = await prismaCore.credentials.findFirst({ where: { - name: name as string, + name: credential.name, + serviceType: credential?.serviceType, projectId: this.projectId, }, }) - await this.update() + if (!existingCredential) { + throw new Error(`Credential ${credential.name} not found`) + } + + await prismaCore.credentials.update({ + where: { id: existingCredential.id }, + data: { + value: encrypt(credential.value, CREDENTIALS_ENCRYPTION_KEY), + credentialType: credential.credentialType, + description: credential.description, + }, + }) + + await this.refreshCredentialsCache() + return true } - async validateCredential(name: keyof T): Promise { - const credential = this.getCredential(name) - return credential !== undefined && credential !== null + async deleteCredential(name: string): Promise { + try { + // First, delete the agent_credentials link + await prismaCore.agent_credentials.deleteMany({ + where: { + agentId: this.agentId, + credentials: { + name: name as string, + projectId: this.projectId, + }, + }, + }) + + // Then, delete the credential itself + await prismaCore.credentials.deleteMany({ + where: { + name: name as string, + projectId: this.projectId, + }, + }) + + await this.refreshCredentialsCache() + return true + } catch (error) { + throw new Error( + `Error deleting credential ${name} for agent: ${this.agentId}: ${error}` + ) + } } - getRequiredCredentials(): (keyof T)[] { - // Implement your logic to return required credentials - return [] as (keyof T)[] + async validateCredential(name: string): Promise { + const credential = this.getCredential(name) + return credential !== undefined && credential !== null } }