link-stack/apps/bridge-whatsapp/src/service.ts

import { Server } from "@hapi/hapi";
import { Service } from "@hapipal/schmervice";
import makeWASocket, {
  DisconnectReason,
  proto,
  downloadContentFromMessage,
  fetchLatestBaileysVersion,
  isJidBroadcast,
  isJidStatusBroadcast,
  useMultiFileAuthState,
} from "@whiskeysockets/baileys";

type MediaType = "audio" | "document" | "image" | "video" | "sticker";
import fs from "fs";
import { createLogger } from "@link-stack/logger";
import {
  getMaxAttachmentSize,
  getMaxTotalAttachmentSize,
  MAX_ATTACHMENTS,
} from "@link-stack/bridge-common";

const logger = createLogger("bridge-whatsapp-service");

export type AuthCompleteCallback = (error?: string) => void;

export default class WhatsappService extends Service {
  connections: { [key: string]: any } = {};
  loginConnections: { [key: string]: any } = {};

  static browserDescription: [string, string, string] = ["Bridge", "Chrome", "2.0"];

  constructor(server: Server, options: never) {
    super(server, options);
  }

  getBaseDirectory(): string {
    return `/home/node/baileys`;
  }

  getBotDirectory(id: string): string {
    // Validate that ID contains only safe characters (alphanumeric, dash, underscore)
    if (!/^[a-zA-Z0-9_-]+$/.test(id)) {
      throw new Error(`Invalid bot ID format: ${id}`);
    }

    // Prevent path traversal by checking for suspicious patterns
    if (id.includes("..") || id.includes("/") || id.includes("\\")) {
      throw new Error(`Path traversal detected in bot ID: ${id}`);
    }

    const botPath = `${this.getBaseDirectory()}/${id}`;

    // Ensure the resolved path is still within the base directory
    if (!botPath.startsWith(this.getBaseDirectory())) {
      throw new Error(`Invalid bot path: ${botPath}`);
    }

    return botPath;
  }

  getAuthDirectory(id: string): string {
    return `${this.getBotDirectory(id)}/auth`;
  }

  async initialize(): Promise<void> {
    this.updateConnections();
  }

  async teardown(): Promise<void> {
    this.resetConnections();
  }

  private async sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
  }

  private async resetConnections() {
    for (const connection of Object.values(this.connections)) {
      try {
        connection.end(null);
      } catch (error) {
        logger.error({ error }, "Connection reset error");
      }
    }
    this.connections = {};
  }

  private async createConnection(
    botID: string,
    server: Server,
    options: any,
    authCompleteCallback?: any,
  ) {
    const authDirectory = this.getAuthDirectory(botID);
    const { state, saveCreds } = await useMultiFileAuthState(authDirectory);
    const msgRetryCounterMap: any = {};
    const socket = makeWASocket({
      ...options,
      auth: state,
      generateHighQualityLinkPreview: false,
      msgRetryCounterMap,
      shouldIgnoreJid: (jid) => isJidBroadcast(jid) || isJidStatusBroadcast(jid),
    });
    let pause = 5000;

    socket.ev.process(async (events) => {
      if (events["connection.update"]) {
        const update = events["connection.update"];
        const { connection: connectionState, lastDisconnect, qr, isNewLogin } = update;
        if (qr) {
          logger.info("got qr code");
          const botDirectory = this.getBotDirectory(botID);
          const qrPath = `${botDirectory}/qr.txt`;
          fs.writeFileSync(qrPath, qr, "utf8");
        } else if (isNewLogin) {
          logger.info("got new login");
          const botDirectory = this.getBotDirectory(botID);
          const verifiedFile = `${botDirectory}/verified`;
          fs.writeFileSync(verifiedFile, "");
        } else if (connectionState === "open") {
          logger.info("opened connection");
        } else if (connectionState === "close") {
          logger.info({ lastDisconnect }, "connection closed");
          const disconnectStatusCode = (lastDisconnect?.error as any)?.output?.statusCode;
          if (disconnectStatusCode === DisconnectReason.restartRequired) {
            logger.info("reconnecting after got new login");
            await this.createConnection(botID, server, options);
            authCompleteCallback?.();
          } else if (disconnectStatusCode !== DisconnectReason.loggedOut) {
            logger.info("reconnecting");
            await this.sleep(pause);
            pause *= 2;
            this.createConnection(botID, server, options);
          }
        }
      }

      if (events["creds.update"]) {
        logger.info("creds update");
        await saveCreds();
      }

      if (events["messages.upsert"]) {
        logger.info("messages upsert");
        const upsert = events["messages.upsert"];
        const { messages } = upsert;
        if (messages) {
          await this.queueUnreadMessages(botID, messages);
        }
      }
    });

    this.connections[botID] = { socket, msgRetryCounterMap };
  }

  private async updateConnections() {
    this.resetConnections();

    const baseDirectory = this.getBaseDirectory();
    const botIDs = fs.readdirSync(baseDirectory);

    for await (const botID of botIDs) {
      const directory = this.getBotDirectory(botID);
      const verifiedFile = `${directory}/verified`;
      if (fs.existsSync(verifiedFile)) {
        const { version, isLatest } = await fetchLatestBaileysVersion();
        logger.info({ version: version.join("."), isLatest }, "using WA version");

        await this.createConnection(botID, this.server, {
          browser: WhatsappService.browserDescription,
          printQRInTerminal: true,
          version,
        });
      }
    }
  }

  private async queueMessage(botID: string, webMessageInfo: proto.IWebMessageInfo) {
    const { key, message, messageTimestamp } = webMessageInfo;
    if (!key) {
      logger.warn("Message missing key, skipping");
      return;
    }
    const { id, fromMe, remoteJid } = key;
    logger.info("Message type debug");
    for (const key in message) {
      logger.info(
        { key, exists: !!message[key as keyof proto.IMessage] },
        "Message field",
      );
    }
    const isValidMessage = message && remoteJid !== "status@broadcast" && !fromMe;
    if (isValidMessage) {
      const { audioMessage, documentMessage, imageMessage, videoMessage } = message;
      const isMediaMessage =
        audioMessage || documentMessage || imageMessage || videoMessage;

      const messageContent = Object.values(message)[0];
      let messageType: MediaType;
      let attachment: string | null | undefined;
      let filename: string | null | undefined;
      let mimeType: string | null | undefined;
      if (isMediaMessage) {
        if (audioMessage) {
          messageType = "audio";
          const extension = audioMessage.mimetype?.split("/").pop() || "audio";
          filename = `${id}.${extension}`;
          mimeType = audioMessage.mimetype;
        } else if (documentMessage) {
          messageType = "document";
          filename = documentMessage.fileName || `${id}.bin`;
          mimeType = documentMessage.mimetype;
        } else if (imageMessage) {
          messageType = "image";
          const extension = imageMessage.mimetype?.split("/").pop() || "jpg";
          filename = `${id}.${extension}`;
          mimeType = imageMessage.mimetype;
        } else if (videoMessage) {
          messageType = "video";
          const extension = videoMessage.mimetype?.split("/").pop() || "mp4";
          filename = `${id}.${extension}`;
          mimeType = videoMessage.mimetype;
        }

        const stream = await downloadContentFromMessage(
          messageContent,
          // @ts-ignore
          messageType,
        );
        let buffer = Buffer.from([]);
        for await (const chunk of stream) {
          buffer = Buffer.concat([buffer, chunk]);
        }
        attachment = buffer.toString("base64");
      }

      if (messageContent || attachment) {
        const conversation = message?.conversation;
        const extendedTextMessage = message?.extendedTextMessage?.text;
        const imageMessage = message?.imageMessage?.caption;
        const videoMessage = message?.videoMessage?.caption;
        const messageText = [
          conversation,
          extendedTextMessage,
          imageMessage,
          videoMessage,
        ].find((text) => text && text !== "");

        const payload = {
          to: botID,
          from: remoteJid?.split("@")[0],
          messageId: id,
          sentAt: new Date((messageTimestamp as number) * 1000).toISOString(),
          message: messageText,
          attachment,
          filename,
          mimeType,
        };

        await fetch(
          `${process.env.BRIDGE_FRONTEND_URL}/api/whatsapp/bots/${botID}/receive`,
          {
            method: "POST",
            headers: {
              "Content-Type": "application/json",
            },
            body: JSON.stringify(payload),
          },
        );
      }
    }
  }

  private async queueUnreadMessages(botID: string, messages: proto.IWebMessageInfo[]) {
    for await (const message of messages) {
      await this.queueMessage(botID, message);
    }
  }

  getBot(botID: string): Record<string, any> {
    const botDirectory = this.getBotDirectory(botID);
    const qrPath = `${botDirectory}/qr.txt`;
    const verifiedFile = `${botDirectory}/verified`;
    const qr = fs.existsSync(qrPath) ? fs.readFileSync(qrPath, "utf8") : null;
    const verified = fs.existsSync(verifiedFile);

    return { qr, verified };
  }

  async unverify(botID: string): Promise<void> {
    // Step 1: Close and remove the active connection if it exists
    const connection = this.connections[botID];
    if (connection?.socket) {
      try {
        // Properly close the WebSocket connection
        await connection.socket.logout();
      } catch (error) {
        logger.warn({ botID, error }, "Error during logout, forcing disconnect");
        try {
          connection.socket.end(undefined);
        } catch (endError) {
          logger.warn({ botID, endError }, "Error ending socket connection");
        }
      }
    }

    // Step 2: Remove from in-memory connections
    delete this.connections[botID];

    // Step 3: Remove the bot directory (auth state, QR code, verified marker)
    const botDirectory = this.getBotDirectory(botID);
    if (fs.existsSync(botDirectory)) {
      fs.rmSync(botDirectory, { recursive: true, force: true });
    }
  }

  async register(botID: string, callback?: AuthCompleteCallback): Promise<void> {
    const { version } = await fetchLatestBaileysVersion();
    await this.createConnection(
      botID,
      this.server,
      { version, browser: WhatsappService.browserDescription },
      callback,
    );
    callback?.();
  }

  async send(
    botID: string,
    phoneNumber: string,
    message: string,
    attachments?: Array<{ data: string; filename: string; mime_type: string }>,
  ): Promise<void> {
    const connection = this.connections[botID]?.socket;
    const recipient = `${phoneNumber.replace(/\D+/g, "")}@s.whatsapp.net`;

    // Send text message if provided
    if (message) {
      await connection.sendMessage(recipient, { text: message });
    }

    // Send attachments if provided with size validation
    if (attachments && attachments.length > 0) {
      const MAX_ATTACHMENT_SIZE = getMaxAttachmentSize();
      const MAX_TOTAL_SIZE = getMaxTotalAttachmentSize();

      if (attachments.length > MAX_ATTACHMENTS) {
        throw new Error(
          `Too many attachments: ${attachments.length} (max ${MAX_ATTACHMENTS})`,
        );
      }

      let totalSize = 0;

      for (const attachment of attachments) {
        // Calculate size before converting to buffer
        const estimatedSize = (attachment.data.length * 3) / 4;

        if (estimatedSize > MAX_ATTACHMENT_SIZE) {
          logger.warn(
            {
              filename: attachment.filename,
              size: estimatedSize,
              maxSize: MAX_ATTACHMENT_SIZE,
            },
            "Attachment exceeds size limit, skipping",
          );
          continue;
        }

        totalSize += estimatedSize;
        if (totalSize > MAX_TOTAL_SIZE) {
          logger.warn(
            {
              totalSize,
              maxTotalSize: MAX_TOTAL_SIZE,
            },
            "Total attachment size exceeds limit, skipping remaining",
          );
          break;
        }

        const buffer = Buffer.from(attachment.data, "base64");

        if (attachment.mime_type.startsWith("image/")) {
          await connection.sendMessage(recipient, {
            image: buffer,
            caption: attachment.filename,
          });
        } else if (attachment.mime_type.startsWith("video/")) {
          await connection.sendMessage(recipient, {
            video: buffer,
            caption: attachment.filename,
          });
        } else if (attachment.mime_type.startsWith("audio/")) {
          await connection.sendMessage(recipient, {
            audio: buffer,
            mimetype: attachment.mime_type,
          });
        } else {
          await connection.sendMessage(recipient, {
            document: buffer,
            fileName: attachment.filename,
            mimetype: attachment.mime_type,
          });
        }
      }
    }
  }

  async receive(
    botID: string,
    _lastReceivedDate: Date,
  ): Promise<proto.IWebMessageInfo[]> {
    const connection = this.connections[botID]?.socket;
    const messages = await connection.loadAllUnreadMessages();

    return messages;
  }
}