service-matrix/matrix/plugin.ts

import {
  Plugin,
  Configurable,
  EventListener,
  InjectService,
} from '@squeebot/core/lib/plugin';

import util from 'util';
import { promises as fs } from 'fs';

import {
  MatrixClient,
  AutojoinRoomsMixin,
  SimpleFsStorageProvider,
  LogService,
  LogLevel,
  RustSdkCryptoStorageProvider,
  RichReply,
} from 'matrix-bot-sdk';
import { StoreType } from '@matrix-org/matrix-sdk-crypto-nodejs';

import { logger } from '@squeebot/core/lib/core';
import {
  EMessageType,
  Formatter,
  IMessage,
  IMessageTarget,
  Protocol,
  ProtocolFeatureFlag,
} from '@squeebot/core/lib/types';

import { MatrixFormatter } from './format';
import { resolve, join } from 'path';

/**
 * Message adapter for Matrix protocol.
 */
class MatrixMessageAdapter implements IMessage {
  public time: Date = new Date();
  public resolved = false;
  public isReply = false;

  constructor(
    public type: EMessageType,
    public data: any,
    public source: Protocol,
    public sender: IMessageTarget,
    public target?: IMessageTarget,
    public direct = false
  ) {}

  public get fullSenderID(): string {
    return this.source.fullName + '/' + this.sender.id;
  }

  public get fullRoomID(): string {
    if (!this.target) {
      return this.source.fullName;
    }

    return this.source.fullName + '/' + this.target.id;
  }

  public get text(): string {
    return this.data.content.body;
  }

  public resolve(...args: any[]): void {
    this.resolved = true;
    this.source.resolve(this, ...args);
  }

  public reject(error: Error): void {
    this.resolved = true;
    this.source.resolve(this, error.message);
  }

  public reply(...args: any[]): void {
    this.isReply = true;
    this.resolve(...args);
  }

  public mention(user: IMessageTarget): string {
    return `<a href="https://matrix.to/#/${user.id}">${user.id}</a>`;
  }
}

/**
 * File system utilities for storing data outside of the configuration.
 */
class MatrixFileSystemUtils {
  constructor(private config: any) {}

  get baseDir() {
    return join(resolve('.matrix'), this.config.name);
  }

  get storageFile() {
    return join(this.baseDir, 'db.json');
  }

  get cryptoStorageDir() {
    return join(this.baseDir, 'sled');
  }

  get tokenFile() {
    return join(this.baseDir, 'token');
  }

  public async ensureMatrixDirectory() {
    try {
      await fs.mkdir(resolve(this.baseDir), { recursive: true });
    } catch {}
  }

  public async readAuth(): Promise<Record<string, string>> {
    try {
      await fs.stat(this.tokenFile);
    } catch {
      return {};
    }
    return JSON.parse(await fs.readFile(this.tokenFile, { encoding: 'utf-8' }));
  }

  public async writeAuth(data: Record<string, string>): Promise<void> {
    await fs.writeFile(this.tokenFile, JSON.stringify(data), {
      encoding: 'utf-8',
    });
  }
}

class MatrixProtocol extends Protocol {
  public static Features = [
    ProtocolFeatureFlag.COLORS,
    ProtocolFeatureFlag.FORMATTING,
    ProtocolFeatureFlag.PLAIN,
    ProtocolFeatureFlag.HTML,
    ProtocolFeatureFlag.EMOJI,
    ProtocolFeatureFlag.IMAGES,
    ProtocolFeatureFlag.VOICE,
    ProtocolFeatureFlag.VIDEO,
    ProtocolFeatureFlag.THREADS,
    ProtocolFeatureFlag.REACTIONS,
    ProtocolFeatureFlag.MENTION,
    ProtocolFeatureFlag.REPLY,
    ProtocolFeatureFlag.OPTIONAL_ENCRYPTION,
    ProtocolFeatureFlag.EVENT_MESSAGE,
    ProtocolFeatureFlag.EVENT_ROOM_JOIN,
    ProtocolFeatureFlag.EVENT_ROOM_LEAVE,
    ProtocolFeatureFlag.KICK,
    ProtocolFeatureFlag.BAN,
    ProtocolFeatureFlag.MUTE,
  ];

  public format: Formatter = new MatrixFormatter();
  public type = 'MatrixProtocol';
  public nameCache: Map<string, string> = new Map<string, string>();

  private client?: MatrixClient;
  private fs = new MatrixFileSystemUtils(this.config);

  public async handler(roomId: any, event: any): Promise<void> {
    // Don't handle events that don't have contents (they were probably redacted)
    if (!this.client || !this.me || !event.content) {
      return;
    }

    // Don't handle non-text events and quote replies
    if (
      !['m.text', 'm.emote'].includes(event.content.msgtype) ||
      event.content['m.relates_to']?.['m.in_reply_to']
    ) {
      return;
    }
    const msg = event.content.body;
    const direct = this.client.dms.isDm(roomId);

    // filter out events sent by the bot itself
    if (event.sender === this.me.id || !msg) {
      return;
    }
    const senderName = await this.getUserDisplayName(event.sender);
    const roomName = direct
      ? senderName
      : await this.getRoomDisplayName(roomId);

    const newMessage = new MatrixMessageAdapter(
      EMessageType.message,
      event,
      this,
      { id: event.sender, name: senderName },
      { id: roomId, name: roomName },
      direct
    );
    this.plugin.stream.emitTo('channel', 'message', newMessage);
  }

  public async getUserDisplayName(id: string): Promise<string> {
    if (this.nameCache.has(id)) {
      return this.nameCache.get(id) as string;
    }

    try {
      const profile = await this.client?.getUserProfile(id);
      if (!profile || !profile.displayname) {
        return id;
      }

      this.nameCache.set(id, profile.displayname);
      return profile.displayname;
    } catch {
      return id;
    }
  }

  public async getRoomDisplayName(id: string): Promise<string> {
    if (this.nameCache.has(id)) {
      return this.nameCache.get(id) as string;
    }

    try {
      const roomState = await this.client?.getRoomStateEvent(
        id,
        'm.room.name',
        ''
      );

      if (!roomState || !roomState.name) {
        return id;
      }

      this.nameCache.set(id, roomState.name);

      return roomState.name;
    } catch {
      return id;
    }
  }

  private async getSelf(): Promise<void> {
    if (!this.client) {
      return;
    }

    const id = await this.client.getUserId();
    const name = await this.getUserDisplayName(id);

    this.me = { id, name };
  }

  private async matrixLogin(
    homeserverUrl: string,
    username: string,
    password: string,
    deviceName: string
  ): Promise<{ accessToken: string; deviceId: string }> {
    const body = {
      type: 'm.login.password',
      identifier: {
        type: 'm.id.user',
        user: username,
      },
      password,
      initial_device_display_name: deviceName,
    };

    const response = await new MatrixClient(homeserverUrl, '').doRequest(
      'POST',
      '/_matrix/client/r0/login',
      null,
      body
    );
    const accessToken = response['access_token'];
    const deviceId = response['device_id'];
    if (!accessToken)
      throw new Error('Expected access token in response - got nothing');

    return { accessToken, deviceId };
  }

  public async start(...args: any[]): Promise<void> {
    await this.fs.ensureMatrixDirectory();

    LogService.setLevel(LogLevel.ERROR);

    const storage = new SimpleFsStorageProvider(this.fs.storageFile);
    const cryptoStorage = new RustSdkCryptoStorageProvider(
      this.fs.cryptoStorageDir,
      StoreType.Sled
    );

    let { accessToken, deviceId } = await this.fs.readAuth();
    if (!accessToken) {
      // No token, log in
      const loginInfo = await this.matrixLogin(
        this.config.homeserver,
        this.config.username,
        this.config.password,
        'Squeebot-Matrix'
      );
      accessToken = loginInfo.accessToken;
      deviceId = loginInfo.deviceId;
      await this.fs.writeAuth({ accessToken, deviceId });
      await cryptoStorage.setDeviceId(deviceId);
    }

    this.client = new MatrixClient(
      this.config.homeserver,
      accessToken,
      storage,
      !!this.config.encrypted ? cryptoStorage : undefined
    );
    AutojoinRoomsMixin.setupOnClient(this.client);

    this.client.on('room.message', (...dargs) =>
      this.handler(dargs[0], dargs[1]).catch((e) => logger.error(e))
    );

    this.client.on(
      'room.failed_decryption',
      async (roomId: string, event: any, e: Error) =>
        logger.error(
          `Failed to decrypt ${roomId} ${event['event_id']} because `,
          e
        )
    );

    try {
      await this.client.start();
      this.running = true;
      logger.info(
        '[%s] Protocol "%s" ready',
        this.plugin.manifest.name,
        this.config.name
      );
      this.getSelf();
    } catch (err) {
      this.stop(true);
      logger.error(
        '[%s] Protocol "%s" failed to start',
        this.plugin.manifest.name,
        this.config.name,
        (err as Error).message
      );
    }
  }

  public stop(force = false): void {
    if (!this.running && !force) {
      return;
    }

    if (this.client) {
      this.client.stop();
      this.client = undefined;
    }

    this.running = false;
    this.stopped = true;

    if (force) {
      this.failed = true;
    }

    this.emit('stopped');
  }

  public async resolve(
    msg: MatrixMessageAdapter,
    ...data: any[]
  ): Promise<void> {
    if (!msg.target) {
      return;
    }

    const parts = this.getMessageParts(
      `${this.fullName}/${msg.target.id}`,
      ...data
    );
    if (!parts) {
      return;
    }

    const { roomId, msgtype, text, html } = parts;
    const messageContents = {
      // Assign properties generated by reply
      ...(msg.isReply
        ? RichReply.createFor(roomId, msg.data, text, html)
        : {
            body: text,
            format: 'org.matrix.custom.html',
            formatted_body: html,
          }),
      // Override message type to support notice
      msgtype,
    };

    await this.client?.sendMessage(roomId, messageContents);
  }

  public async sendTo(target: string, ...data: any[]): Promise<boolean> {
    const parts = this.getMessageParts(target, ...data);
    if (!parts) {
      return false;
    }

    const { roomId, msgtype, text, html } = parts;
    await this.client?.sendMessage(roomId, {
      msgtype,
      body: text,
      format: 'org.matrix.custom.html',
      formatted_body: html,
    });

    return true;
  }

  private getMessageParts(target: string, ...data: any[]) {
    let response = util.format(data[0], ...data.slice(1));
    const rxSplit = target.split('/');
    if (!response || !target || rxSplit.length !== 3) {
      return undefined;
    }

    if (Array.isArray(data[0])) {
      try {
        response = this.format.compose(data[0]);
      } catch (e: any) {
        logger.error(
          '[%s] Failed to compose message:',
          this.fullName,
          e.message
        );
        return undefined;
      }
    }

    const msgtype =
      response.startsWith('<m.emote>') && response.endsWith('</m.emote>')
        ? 'm.emote'
        : 'm.text';

    // TODO: make more generic
    if (msgtype === 'm.emote') {
      response = response.substring(9, response.length - 10);
    }

    return {
      roomId: rxSplit[2],
      msgtype,
      response,
      text: this.format.strip(response),
      html: response.replace(/\n/g, '<br />'),
    };
  }
}

@InjectService(MatrixProtocol)
@Configurable({ instances: [] })
class MatrixServicePlugin extends Plugin {
  initialize(): void {
    const protoList = this.validateConfiguration();
    this.startAll(protoList);
  }

  private startAll(list: any[]): void {
    for (const ins of list) {
      const newProto = new MatrixProtocol(this, ins, MatrixProtocol.Features);
      logger.log('[%s] Starting Matrix service "%s".', this.name, ins.name);
      this.monitor(newProto);
      this.service?.use(newProto, true);
    }
  }

  private monitor(proto: Protocol): void {
    proto.on('running', () => this.emit('protocolNew', proto));
    proto.on('stopped', () => this.emit('protocolExit', proto));
  }

  private validateConfiguration(): any[] {
    if (!this.config.config.instances) {
      throw new Error('Configuration incomplete!');
    }

    const instances = this.config.config.instances;
    const runnables: any[] = [];
    for (const ins of instances) {
      if (ins.enabled === false) {
        continue;
      }

      if (!ins.name || !ins.username || !ins.password || !ins.homeserver) {
        throw new Error('Invalid instance configuration!');
      }

      runnables.push(ins);
    }

    return runnables;
  }

  @EventListener('pluginUnload')
  public unloadEventHandler(plugin: string | Plugin): void {
    if (plugin === this.name || plugin === this) {
      this.config
        .save()
        .then(() =>
          this.service?.stopAll().then(() => this.emit('pluginUnloaded', this))
        );
    }
  }
}

module.exports = MatrixServicePlugin;