Move packages/apps back

apps/metamigo-api/app/lib/whatsapp-key-store.ts

@@ -0,0 +1,198 @@
+import { Boom } from "@hapi/boom";
+import { Server } from "@hapi/hapi";
+import { randomBytes } from "crypto";
+import type { Logger } from "pino";
+import {
+  proto,
+  BufferJSON,
+  generateRegistrationId,
+  Curve,
+  signedKeyPair,
+  AuthenticationCreds,
+  AuthenticationState,
+  AccountSettings,
+  SignalDataSet,
+  SignalDataTypeMap,
+  SignalKeyStore,
+  SignalKeyStoreWithTransaction,
+} from "@adiwajshing/baileys";
+import { SavedWhatsappBot as Bot } from "db";
+
+const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = {
+  "pre-key": "preKeys",
+  session: "sessions",
+  "sender-key": "senderKeys",
+  "app-state-sync-key": "appStateSyncKeys",
+  "app-state-sync-version": "appStateVersions",
+  "sender-key-memory": "senderKeyMemory",
+};
+
+export const addTransactionCapability = (
+  state: SignalKeyStore,
+  logger: Logger
+): SignalKeyStoreWithTransaction => {
+  let inTransaction = false;
+  let transactionCache: SignalDataSet = {};
+  let mutations: SignalDataSet = {};
+
+  const prefetch = async (type: keyof SignalDataTypeMap, ids: string[]) => {
+    if (!inTransaction) {
+      throw new Boom("Cannot prefetch without transaction");
+    }
+
+    const dict = transactionCache[type];
+    const idsRequiringFetch = dict
+      ? ids.filter((item) => !(item in dict))
+      : ids;
+    // only fetch if there are any items to fetch
+    if (idsRequiringFetch.length) {
+      const result = await state.get(type, idsRequiringFetch);
+
+      transactionCache[type] = transactionCache[type] || {};
+      // @ts-expect-error
+      Object.assign(transactionCache[type], result);
+    }
+  };
+
+  return {
+    get: async (type, ids) => {
+      if (inTransaction) {
+        await prefetch(type, ids);
+        return ids.reduce((dict, id) => {
+          const value = transactionCache[type]?.[id];
+          if (value) {
+            // @ts-expect-error
+            dict[id] = value;
+          }
+
+          return dict;
+        }, {});
+      } else {
+        return state.get(type, ids);
+      }
+    },
+    set: (data) => {
+      if (inTransaction) {
+        logger.trace({ types: Object.keys(data) }, "caching in transaction");
+        for (const key in data) {
+          // @ts-expect-error
+          transactionCache[key] = transactionCache[key] || {};
+          // @ts-expect-error
+          Object.assign(transactionCache[key], data[key]);
+          // @ts-expect-error
+          mutations[key] = mutations[key] || {};
+          // @ts-expect-error
+          Object.assign(mutations[key], data[key]);
+        }
+      } else {
+        return state.set(data);
+      }
+    },
+    isInTransaction: () => inTransaction,
+    // @ts-expect-error
+    prefetch: (type, ids) => {
+      logger.trace({ type, ids }, "prefetching");
+      return prefetch(type, ids);
+    },
+    transaction: async (work) => {
+      if (inTransaction) {
+        await work();
+      } else {
+        logger.debug("entering transaction");
+        inTransaction = true;
+        try {
+          await work();
+          if (Object.keys(mutations).length) {
+            logger.debug("committing transaction");
+            await state.set(mutations);
+          } else {
+            logger.debug("no mutations in transaction");
+          }
+        } finally {
+          inTransaction = false;
+          transactionCache = {};
+          mutations = {};
+        }
+      }
+    },
+  };
+};
+
+export const initAuthCreds = (): AuthenticationCreds => {
+  const identityKey = Curve.generateKeyPair();
+  return {
+    noiseKey: Curve.generateKeyPair(),
+    signedIdentityKey: identityKey,
+    signedPreKey: signedKeyPair(identityKey, 1),
+    registrationId: generateRegistrationId(),
+    advSecretKey: randomBytes(32).toString("base64"),
+    nextPreKeyId: 1,
+    firstUnuploadedPreKeyId: 1,
+
+    processedHistoryMessages: [],
+    accountSettings: {
+      unarchiveChats: false,
+    },
+  } as any;
+};
+
+export const useDatabaseAuthState = (
+  bot: Bot,
+  server: Server
+): { state: AuthenticationState; saveState: () => void } => {
+  let { logger }: any = server;
+  let creds: AuthenticationCreds;
+  let keys: any = {};
+
+  const saveState = async () => {
+    logger && logger.trace("saving auth state");
+    const authInfo = JSON.stringify({ creds, keys }, BufferJSON.replacer, 2);
+    await server.db().whatsappBots.updateAuthInfo(bot, authInfo);
+  };
+
+  if (bot.authInfo) {
+    console.log("Auth info exists");
+    const result = JSON.parse(bot.authInfo, BufferJSON.reviver);
+    creds = result.creds;
+    keys = result.keys;
+  } else {
+    console.log("Auth info does not exist");
+    creds = initAuthCreds();
+    keys = {};
+  }
+
+  return {
+    state: {
+      creds,
+      keys: {
+        get: (type, ids) => {
+          const key = KEY_MAP[type];
+          return ids.reduce((dict, id) => {
+            let value = keys[key]?.[id];
+            if (value) {
+              if (type === "app-state-sync-key") {
+                // @ts-expect-error
+                value = proto.AppStateSyncKeyData.fromObject(value);
+              }
+              // @ts-expect-error
+              dict[id] = value;
+            }
+
+            return dict;
+          }, {});
+        },
+        set: (data) => {
+          for (const _key in data) {
+            const key = KEY_MAP[_key as keyof SignalDataTypeMap];
+            keys[key] = keys[key] || {};
+            // @ts-expect-error
+            Object.assign(keys[key], data[_key]);
+          }
+
+          saveState();
+        },
+      },
+    },
+    saveState,
+  };
+};