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

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 "@digiresilience/metamigo-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,
  };
};