Draupnir/apps/draupnir/test/integration/mjolnirSetupUtils.ts

// SPDX-FileCopyrightText: 2026 Catalan Lover <catalanlover@protonmail.com>
// Copyright 2022 - 2024 Gnuxie <Gnuxie@protonmail.com>
// Copyright 2021 - 2022 The Matrix.org Foundation C.I.C.
//
// SPDX-License-Identifier: Apache-2.0
//
// SPDX-FileAttributionText: <text>
// This modified file incorporates work from mjolnir
// https://github.com/matrix-org/mjolnir
// </text>

import {
  MatrixClient,
  PantalaimonClient,
  MemoryStorageProvider,
  LogService,
  LogLevel,
} from "@vector-im/matrix-bot-sdk";
import { overrideRatelimitForUser, registerUser } from "./clientHelper";
import { initializeSentry, patchMatrixClient } from "../../src/utils";
import { IConfig } from "../../src/config";
import { Draupnir } from "../../src/Draupnir";
import { DraupnirBotModeToggle } from "../../src/DraupnirBotMode";
import {
  MatrixSendClient,
  SafeMatrixEmitter,
  SafeMatrixEmitterWrapper,
} from "matrix-protection-suite-for-matrix-bot-sdk";
import {
  DefaultEventDecoder,
  MJOLNIR_PROTECTED_ROOMS_EVENT_TYPE,
  MJOLNIR_WATCHED_POLICY_ROOMS_EVENT_TYPE,
  OwnLifetime,
} from "matrix-protection-suite";
import { SafeModeDraupnir } from "../../src/safemode/DraupnirSafeMode";
import { TopLevelStores } from "../../src/backingstore/DraupnirStores";
import {
  isStringUserID,
  userLocalpart,
} from "@the-draupnir-project/matrix-basic-types";

patchMatrixClient();

// they are add [key: string]: any to their interface, amazing.
export type SafeMochaContext = Pick<
  Mocha.Context,
  | "test"
  | "currentTest"
  | "runnable"
  | "timeout"
  | "slow"
  | "skip"
  | "retries"
  | "done"
>;

export interface DraupnirTestContext extends SafeMochaContext {
  lifetime: OwnLifetime;
  draupnir?: Draupnir;
  toggle?: DraupnirBotModeToggle;
  config: IConfig;
  stores?: TopLevelStores;
}

/**
 * Ensures that a room exists with the alias, if it does not exist we create it.
 * @param client The MatrixClient to use to resolve or create the aliased room.
 * @param alias The alias of the room.
 * @returns The room ID of the aliased room.
 */
export async function ensureAliasedRoomExists(
  client: MatrixClient,
  alias: string
): Promise<string> {
  try {
    return await client.resolveRoom(alias);
  } catch (e) {
    if (typeof e !== "object" || e === null) {
      throw new TypeError("Something is throwing garbage");
    }
    if (!("body" in e) || typeof e.body !== "object" || e.body === null) {
      throw e;
    }
    if (!("errcode" in e.body) || e.body.errcode !== "M_NOT_FOUND") {
      throw e;
    }
    console.info(`${alias} hasn't been created yet, so we're making it now.`);
    const roomId = await client.createRoom({
      visibility: "public",
    });
    await client.createRoomAlias(alias, roomId);
    return roomId;
  }
}

async function configureMjolnir(config: IConfig) {
  // Initialize error monitoring as early as possible.
  initializeSentry(config);
  try {
    await registerUser(
      config.homeserverUrl,
      config.pantalaimon.username,
      config.pantalaimon.username,
      config.pantalaimon.password,
      true
    );
  } catch (e) {
    if (typeof e !== "object" || e === null) {
      throw new TypeError("Something is throwing garbage");
    }
    if (!("body" in e) || typeof e.body !== "object" || e.body === null) {
      throw e;
    }
    if (!("errcode" in e.body) || e.body.errcode !== "M_USER_IN_USE") {
      throw e;
    }
    console.log(`${config.pantalaimon.username} already registered, skipping`);
    return;
  }
}

export function draupnir(): Draupnir {
  if (globalMjolnir === null) {
    throw new TypeError("Setup code didn't run before you called `draupnir()`");
  }
  return globalMjolnir;
}
export function draupnirClient(): MatrixClient | null {
  return globalClient;
}
export function draupnirSafeEmitter(): SafeMatrixEmitter {
  if (globalSafeEmitter !== undefined) {
    return globalSafeEmitter;
  }
  throw new TypeError(`Setup code didn't run properly`);
}
let globalClient: MatrixClient | null;
let globalMjolnir: Draupnir | null;
let globalSafeEmitter: SafeMatrixEmitter | undefined;

/**
 * Return a test instance of Draupnir.
 */
export async function makeBotModeToggle(
  config: IConfig,
  {
    stores,
    eraseAccountData,
    allowSafeMode,
    deleteManagementRoomAliasOnStart: deleteAlias,
    ensureManagementRoomAlias,
  }: {
    stores: TopLevelStores;
    eraseAccountData?: boolean;
    allowSafeMode?: boolean;
    deleteManagementRoomAliasOnStart?: boolean;
    ensureManagementRoomAlias?: boolean;
  } = { stores: { dispose() {} } }
): Promise<DraupnirBotModeToggle> {
  await configureMjolnir(config);
  LogService.setLevel(LogLevel.fromString(config.logLevel, LogLevel.DEBUG));
  LogService.info("test/mjolnirSetupUtils", "Starting bot...");
  const pantalaimon = new PantalaimonClient(
    config.homeserverUrl,
    new MemoryStorageProvider()
  );
  const client = await pantalaimon.createClientWithCredentials(
    config.pantalaimon.username,
    config.pantalaimon.password
  );
  if (eraseAccountData) {
    await Promise.all([
      client.setAccountData(MJOLNIR_PROTECTED_ROOMS_EVENT_TYPE, { rooms: [] }),
      client.setAccountData(MJOLNIR_WATCHED_POLICY_ROOMS_EVENT_TYPE, {
        references: [],
      }),
      leaveAllRooms(client),
    ]);
  }
  await overrideRatelimitForUser(
    config.homeserverUrl,
    await client.getUserId()
  );
  if (config.initialManager !== undefined) {
    if (!isStringUserID(config.initialManager)) {
      throw new TypeError(
        `${config.initialManager} is not a valid initial manager mxid`
      );
    }
    try {
      await registerUser(
        config.homeserverUrl,
        userLocalpart(config.initialManager),
        userLocalpart(config.initialManager),
        userLocalpart(config.initialManager),
        true
      );
    } catch (e) {
      if (
        typeof e !== "object" ||
        e === null ||
        !("body" in e) ||
        typeof e.body !== "object" ||
        e.body === null ||
        !("errcode" in e.body) ||
        e.body.errcode !== "M_USER_IN_USE"
      ) {
        throw e;
      }
    }
  }
  // defaults to true
  if (ensureManagementRoomAlias ?? true) {
    if (config.managementRoom === undefined) {
      throw new TypeError("managementRoom is required when ensuring alias");
    }
    if (deleteAlias === undefined || deleteAlias) {
      await client
        .deleteRoomAlias(config.managementRoom)
        .catch((_: unknown) => undefined);
    }
    await ensureAliasedRoomExists(client, config.managementRoom);
  }
  const toggle = await DraupnirBotModeToggle.create(
    client,
    new SafeMatrixEmitterWrapper(client, DefaultEventDecoder),
    config,
    stores
  );
  // we don't want to send status on startup incase we want to test e2ee from the manual launch script.
  const mj = (
    await toggle.startFromScratch({ sendStatusOnStart: false })
  ).expect("Could not create Draupnir");
  if (mj instanceof SafeModeDraupnir && !allowSafeMode) {
    throw new TypeError(
      "Setup code is wrong, shouldn't be booting into safe mode"
    );
  }
  globalClient = client;
  if (mj instanceof Draupnir) {
    globalMjolnir = mj;
  }
  console.info(`management room ${mj.managementRoom.toPermalink()}`);
  globalSafeEmitter = new SafeMatrixEmitterWrapper(client, DefaultEventDecoder);
  return toggle;
}

/**
 * Remove the alias and leave the room, can't be implicitly provided from the config because Draupnir currently mutates it.
 * @param client The client to use to leave the room.
 * @param roomId The roomId of the room to leave.
 * @param alias The alias to remove from the room.
 */
export async function teardownManagementRoom(
  client: MatrixClient,
  roomId: string,
  alias: string
) {
  await client.deleteRoomAlias(alias);
  await client.leaveRoom(roomId);
}

export async function leaveAllRooms(client: MatrixSendClient): Promise<void> {
  const joinedRooms = await client.getJoinedRooms();
  await Promise.allSettled(
    joinedRooms.map((roomID) => client.leaveRoom(roomID))
  );
}