refacto(kb): Simplify code of aleph persistence

Libertai · Aug 7, 2024 · 744d121 · 744d121
1 parent cb2ef64
commit 744d121
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 58 deletions.
diff --git a/src/utils/aleph-persistent-storage.ts b/src/utils/aleph-persistent-storage.ts
@@ -12,9 +12,8 @@ import {
   knowledgeSchema,
 } from 'src/types/knowledge';
 import { decrypt, encrypt, generateIv, generateKey } from 'src/utils/encryption';
-import { decrypt as eciesDecrypt, encrypt as eciesEncrypt, PrivateKey } from 'eciesjs';
-// @ts-ignore
-import { BufferEncoding } from 'vite-plugin-checker/dist/cjs/checkers/vueTsc/typescript-vue-tsc';
+import { PrivateKey } from 'eciesjs';
+import { decryptKnowledgeBaseIdentifiers, encryptKnowledgeBaseIdentifiers } from 'src/utils/knowledge/encryption';
 
 // Aleph keys and channels settings
 const SECURITY_AGGREGATE_KEY = 'security';
@@ -23,7 +22,6 @@ const LIBERTAI_CHANNEL = 'libertai-chat-ui';
 const LIBERTAI_SETTINGS_KEY = `${LIBERTAI_CHANNEL}-settings`;
 const LIBERTAI_KNOWLEDGE_BASE_IDENTIFIERS_KEY = `${LIBERTAI_CHANNEL}-knowledge-base-identifiers-test-12`;
 const LIBERTAI_KNOWLEDGE_BASE_POST_TYPE = `${LIBERTAI_CHANNEL}-knowledge-base-test-12`;
-const BUFFER_ENCODING: BufferEncoding = 'hex';
 
 export class AlephPersistentStorage {
   constructor(
@@ -171,18 +169,7 @@ export class AlephPersistentStorage {
         throw new Error(`Zod parsing failed: ${parsedKnowledgeBaseIdentifiers.error}`);
       }
 
-      return parsedKnowledgeBaseIdentifiers.data.data.map((kbIdentifier) => {
-        const decryptedKey = eciesDecrypt(
-          this.encryptionPrivateKey.secret,
-          Buffer.from(kbIdentifier.encryption.key, BUFFER_ENCODING),
-        ).toString();
-        const decryptedIv = eciesDecrypt(
-          this.encryptionPrivateKey.secret,
-          Buffer.from(kbIdentifier.encryption.iv, BUFFER_ENCODING),
-        ).toString();
-
-        return { ...kbIdentifier, encryption: { key: decryptedKey, iv: decryptedIv } };
-      });
+      return decryptKnowledgeBaseIdentifiers(parsedKnowledgeBaseIdentifiers.data.data, this.encryptionPrivateKey);
     } catch (error) {
       console.error(`Fetching Knowledge base identifiers from Aleph failed: ${error}`);
       return undefined;
@@ -215,28 +202,12 @@ export class AlephPersistentStorage {
         post_hash: response.item_hash,
       };
 
-      const newKbIdentifiers: KnowledgeBaseIdentifier[] = [...currentKbIdentifiers, identifier].map((kbIdentifier) => ({
-        ...kbIdentifier,
-        encryption: {
-          key: eciesEncrypt(
-            this.encryptionPrivateKey.publicKey.toHex(),
-            Buffer.from(kbIdentifier.encryption.key),
-          ).toString(BUFFER_ENCODING),
-          iv: eciesEncrypt(
-            this.encryptionPrivateKey.publicKey.toHex(),
-            Buffer.from(kbIdentifier.encryption.iv),
-          ).toString(BUFFER_ENCODING),
-        },
-      }));
+      const newKbIdentifiers = encryptKnowledgeBaseIdentifiers(
+        [...currentKbIdentifiers, identifier],
+        this.encryptionPrivateKey,
+      );
 
-      await this.subAccountClient.createAggregate({
-        key: LIBERTAI_KNOWLEDGE_BASE_IDENTIFIERS_KEY,
-        content: {
-          data: newKbIdentifiers,
-        },
-        address: this.account.address,
-        channel: LIBERTAI_CHANNEL,
-      });
+      await this.updateKnowledgeBaseIdentifiers(newKbIdentifiers);
       return identifier;
     } catch (error) {
       console.error(`Creating knowledge base failed: ${error}`);
@@ -289,32 +260,27 @@ export class AlephPersistentStorage {
         hashes: [kbIdentifier.post_hash],
       });
 
-      const newKbIdentifiers: KnowledgeBaseIdentifier[] = [
-        ...currentKbIdentifiers.filter((kbi) => kbi.id !== kbIdentifier.id),
-      ].map((kbi) => ({
-        ...kbi,
-        encryption: {
-          key: eciesEncrypt(this.encryptionPrivateKey.publicKey.toHex(), Buffer.from(kbi.encryption.key)).toString(
-            BUFFER_ENCODING,
-          ),
-          iv: eciesEncrypt(this.encryptionPrivateKey.publicKey.toHex(), Buffer.from(kbi.encryption.iv)).toString(
-            BUFFER_ENCODING,
-          ),
-        },
-      }));
+      const newKbIdentifiers = encryptKnowledgeBaseIdentifiers(
+        currentKbIdentifiers.filter((kbi) => kbi.id !== kbIdentifier.id),
+        this.encryptionPrivateKey,
+      );
 
-      await this.subAccountClient.createAggregate({
-        key: LIBERTAI_KNOWLEDGE_BASE_IDENTIFIERS_KEY,
-        content: {
-          data: newKbIdentifiers,
-        },
-        address: this.account.address,
-        channel: LIBERTAI_CHANNEL,
-      });
+      await this.updateKnowledgeBaseIdentifiers(newKbIdentifiers);
       return true;
     } catch (error) {
       console.error(`Deleting knowledge base failed: ${error}`);
       return false;
     }
   }
+
+  private async updateKnowledgeBaseIdentifiers(kbIdentifiers: KnowledgeBaseIdentifier[]) {
+    return this.subAccountClient.createAggregate({
+      key: LIBERTAI_KNOWLEDGE_BASE_IDENTIFIERS_KEY,
+      content: {
+        data: kbIdentifiers,
+      },
+      address: this.account.address,
+      channel: LIBERTAI_CHANNEL,
+    });
+  }
 }
diff --git a/src/utils/encryption.ts b/src/utils/encryption.ts
@@ -1,9 +1,12 @@
 import { createCipheriv, createDecipheriv, randomBytes } from 'crypto';
+// @ts-ignore
+import { BufferEncoding } from 'vite-plugin-checker/dist/cjs/checkers/vueTsc/typescript-vue-tsc';
 
 // Global settings
 const ALGORITHM = 'aes-256-cbc';
 const INPUT_ENCODING = 'utf8';
 const OUTPUT_ENCODING = 'hex';
+export const BUFFER_ENCODING: BufferEncoding = 'hex';
 
 export const encrypt = (data: string, key: Buffer, iv: Buffer): string => {
   const cipher = createCipheriv(ALGORITHM, key, iv);

diff --git a/src/utils/knowledge/encryption.ts b/src/utils/knowledge/encryption.ts
@@ -0,0 +1,38 @@
+import { KnowledgeBaseIdentifier } from 'src/types/knowledge';
+import { decrypt as eciesDecrypt, encrypt as eciesEncrypt, PrivateKey } from 'eciesjs';
+import { BUFFER_ENCODING } from 'src/utils/encryption';
+
+export const encryptKnowledgeBaseIdentifiers = (
+  identifiers: KnowledgeBaseIdentifier[],
+  encryptionKey: PrivateKey,
+): KnowledgeBaseIdentifier[] =>
+  identifiers.map(
+    (kbIdentifier): KnowledgeBaseIdentifier => ({
+      ...kbIdentifier,
+      encryption: {
+        key: eciesEncrypt(encryptionKey.publicKey.toHex(), Buffer.from(kbIdentifier.encryption.key)).toString(
+          BUFFER_ENCODING,
+        ),
+        iv: eciesEncrypt(encryptionKey.publicKey.toHex(), Buffer.from(kbIdentifier.encryption.iv)).toString(
+          BUFFER_ENCODING,
+        ),
+      },
+    }),
+  );
+
+export const decryptKnowledgeBaseIdentifiers = (
+  identifiers: KnowledgeBaseIdentifier[],
+  encryptionKey: PrivateKey,
+): KnowledgeBaseIdentifier[] =>
+  identifiers.map((kbIdentifier): KnowledgeBaseIdentifier => {
+    const decryptedKey = eciesDecrypt(
+      encryptionKey.secret,
+      Buffer.from(kbIdentifier.encryption.key, BUFFER_ENCODING),
+    ).toString();
+    const decryptedIv = eciesDecrypt(
+      encryptionKey.secret,
+      Buffer.from(kbIdentifier.encryption.iv, BUFFER_ENCODING),
+    ).toString();
+
+    return { ...kbIdentifier, encryption: { key: decryptedKey, iv: decryptedIv } };
+  });