Allow certain or all roles to be removed upon mute and readded on unmute (#140)

backend/src/data/GuildMutes.ts

@@ -34,7 +34,7 @@ export class GuildMutes extends BaseGuildRepository {
     return mute != null;
   }
 
-  async addMute(userId, expiryTime): Promise<Mute> {
+  async addMute(userId, expiryTime, rolesToRestore?: string[]): Promise<Mute> {
     const expiresAt = expiryTime
       ? moment
           .utc()
@@ -46,12 +46,13 @@ export class GuildMutes extends BaseGuildRepository {
       guild_id: this.guildId,
       user_id: userId,
       expires_at: expiresAt,
+      roles_to_restore: rolesToRestore ?? [],
     });
 
     return (await this.mutes.findOne({ where: result.identifiers[0] }))!;
   }
 
-  async updateExpiryTime(userId, newExpiryTime) {
+  async updateExpiryTime(userId, newExpiryTime, rolesToRestore?: string[]) {
     const expiresAt = newExpiryTime
       ? moment
           .utc()
@@ -59,15 +60,28 @@ export class GuildMutes extends BaseGuildRepository {
           .format("YYYY-MM-DD HH:mm:ss")
       : null;
 
-    return this.mutes.update(
+    if (rolesToRestore && rolesToRestore.length) {
-      {
+      return this.mutes.update(
-        guild_id: this.guildId,
+        {
-        user_id: userId,
+          guild_id: this.guildId,
-      },
+          user_id: userId,
-      {
+        },
-        expires_at: expiresAt,
+        {
-      },
+          expires_at: expiresAt,
-    );
+          roles_to_restore: rolesToRestore,
+        },
+      );
+    } else {
+      return this.mutes.update(
+        {
+          guild_id: this.guildId,
+          user_id: userId,
+        },
+        {
+          expires_at: expiresAt,
+        },
+      );
+    }
   }
 
   async getActiveMutes(): Promise<Mute[]> {

backend/src/data/entities/Mute.ts

@@ -15,4 +15,6 @@ export class Mute {
   @Column({ type: String, nullable: true }) expires_at: string | null;
 
   @Column() case_id: number;
+
+  @Column("simple-array") roles_to_restore: string[];
 }

backend/src/migrations/1608608903570-CreateRestoredRolesColumn.ts

@@ -0,0 +1,18 @@
+import { MigrationInterface, QueryRunner, TableColumn } from "typeorm";
+
+export class CreateRestoredRolesColumn1608608903570 implements MigrationInterface {
+  public async up(queryRunner: QueryRunner): Promise<any> {
+    await queryRunner.addColumn(
+      "mutes",
+      new TableColumn({
+        name: "roles_to_restore",
+        type: "text",
+        isNullable: true,
+      }),
+    );
+  }
+
+  public async down(queryRunner: QueryRunner): Promise<any> {
+    await queryRunner.dropColumn("mutes", "roles_to_restore");
+  }
+}

backend/src/plugins/Automod/actions/mute.ts

@@ -22,6 +22,8 @@ export const MuteAction = automodAction({
     duration: tNullable(tDelayString),
     notify: tNullable(t.string),
     notifyChannel: tNullable(t.string),
+    remove_roles_on_mute: tNullable(t.union([t.boolean, t.array(t.string)])),
+    restore_roles_on_mute: tNullable(t.union([t.boolean, t.array(t.string)])),
   }),
 
   defaultConfig: {
@@ -32,6 +34,8 @@ export const MuteAction = automodAction({
     const duration = actionConfig.duration ? convertDelayStringToMS(actionConfig.duration)! : undefined;
     const reason = actionConfig.reason || "Muted automatically";
     const contactMethods = actionConfig.notify ? resolveActionContactMethods(pluginData, actionConfig) : undefined;
+    const rolesToRemove = actionConfig.remove_roles_on_mute;
+    const rolesToRestore = actionConfig.restore_roles_on_mute;
 
     const caseArgs = {
       modId: pluginData.client.user.id,
@@ -43,7 +47,7 @@ export const MuteAction = automodAction({
     const mutes = pluginData.getPlugin(MutesPlugin);
     for (const userId of userIdsToMute) {
       try {
-        await mutes.muteUser(userId, duration, reason, { contactMethods, caseArgs });
+        await mutes.muteUser(userId, duration, reason, { contactMethods, caseArgs }, rolesToRemove, rolesToRestore);
       } catch (e) {
         if (e instanceof RecoverablePluginError && e.code === ERRORS.NO_MUTE_ROLE_IN_CONFIG) {
           pluginData.getPlugin(LogsPlugin).log(LogType.BOT_ALERT, {

backend/src/plugins/Mutes/MutesPlugin.ts

@@ -13,7 +13,6 @@ import { ClearMutesWithoutRoleCmd } from "./commands/ClearMutesWithoutRoleCmd";
 import { ClearMutesCmd } from "./commands/ClearMutesCmd";
 import { muteUser } from "./functions/muteUser";
 import { unmuteUser } from "./functions/unmuteUser";
-import { CaseArgs } from "../Cases/types";
 import { Member } from "eris";
 import { ClearActiveMuteOnMemberBanEvt } from "./events/ClearActiveMuteOnMemberBanEvt";
 import { ReapplyActiveMuteOnJoinEvt } from "./events/ReapplyActiveMuteOnJoinEvt";
@@ -32,6 +31,8 @@ const defaultOptions = {
     mute_message: "You have been muted on the {guildName} server. Reason given: {reason}",
     timed_mute_message: "You have been muted on the {guildName} server for {time}. Reason given: {reason}",
     update_mute_message: "Your mute on the {guildName} server has been updated to {time}.",
+    remove_roles_on_mute: false,
+    restore_roles_on_mute: false,
 
     can_view_list: false,
     can_cleanup: false,

backend/src/plugins/Mutes/functions/clearExpiredMutes.ts

@@ -2,6 +2,7 @@ import { GuildPluginData } from "knub";
 import { MutesPluginType } from "../types";
 import { LogType } from "../../../data/LogType";
 import { resolveMember, stripObjectToScalars, UnknownUser } from "../../../utils";
+import { MemberOptions } from "eris";
 
 export async function clearExpiredMutes(pluginData: GuildPluginData<MutesPluginType>) {
   const expiredMutes = await pluginData.state.mutes.getExpiredMutes();
@@ -14,6 +15,14 @@ export async function clearExpiredMutes(pluginData: GuildPluginData<MutesPluginT
         if (muteRole) {
           await member.removeRole(muteRole);
         }
+        if (mute.roles_to_restore) {
+          const memberOptions: MemberOptions = {};
+          const guildRoles = pluginData.guild.roles;
+          memberOptions.roles = Array.from(
+            new Set([...mute.roles_to_restore, ...member.roles.filter(x => x !== muteRole && guildRoles.has(x))]),
+          );
+          member.edit(memberOptions);
+        }
       } catch (e) {
         pluginData.state.serverLogs.log(LogType.BOT_ALERT, {
           body: `Failed to remove mute role from {userMention(member)}`,

backend/src/plugins/Mutes/functions/muteUser.ts

@@ -12,7 +12,7 @@ import {
   UserNotificationMethod,
 } from "../../../utils";
 import { renderTemplate } from "../../../templateFormatter";
-import { TextChannel, User } from "eris";
+import { MemberOptions, TextChannel, User } from "eris";
 import { CasesPlugin } from "../../Cases/CasesPlugin";
 import { CaseTypes } from "../../../data/CaseTypes";
 import { LogType } from "../../../data/LogType";
@@ -26,6 +26,8 @@ export async function muteUser(
   muteTime?: number,
   reason?: string,
   muteOptions: MuteOptions = {},
+  removeRolesOnMuteOverride: boolean | string[] | null = null,
+  restoreRolesOnMuteOverride: boolean | string[] | null = null,
 ) {
   const lock = await pluginData.locks.acquire(`mute-${userId}`);
 
@@ -52,8 +54,37 @@ export async function muteUser(
   const member = await resolveMember(pluginData.client, pluginData.guild, user.id, true); // Grab the fresh member so we don't have stale role info
   const config = pluginData.config.getMatchingConfig({ member, userId });
 
+  let rolesToRestore: string[] = [];
   if (member) {
     const logs = pluginData.getPlugin(LogsPlugin);
+    // remove and store any roles to be removed/restored
+    const currentUserRoles = member.roles;
+    const memberOptions: MemberOptions = {};
+    const removeRoles = removeRolesOnMuteOverride ?? config.remove_roles_on_mute;
+    const restoreRoles = restoreRolesOnMuteOverride ?? config.restore_roles_on_mute;
+
+    // remove roles
+    if (!Array.isArray(removeRoles)) {
+      if (removeRoles) {
+        // exclude managed roles from being removed
+        const managedRoles = pluginData.guild.roles.filter(x => x.managed).map(y => y.id);
+        memberOptions.roles = managedRoles.filter(x => member.roles.includes(x));
+        await member.edit(memberOptions);
+      }
+    } else {
+      memberOptions.roles = currentUserRoles.filter(x => !(<string[]>removeRoles).includes(x));
+      await member.edit(memberOptions);
+    }
+
+    // set roles to be restored
+    if (!Array.isArray(restoreRoles)) {
+      if (restoreRoles) {
+        rolesToRestore = currentUserRoles;
+      }
+    } else {
+      rolesToRestore = currentUserRoles.filter(x => (<string[]>restoreRoles).includes(x));
+    }
+
     // Apply mute role if it's missing
     if (!member.roles.includes(muteRole)) {
       try {
@@ -103,9 +134,12 @@ export async function muteUser(
   let notifyResult: UserNotificationResult = { method: null, success: true };
 
   if (existingMute) {
-    await pluginData.state.mutes.updateExpiryTime(user.id, muteTime);
+    if (existingMute.roles_to_restore?.length || rolesToRestore?.length) {
+      rolesToRestore = Array.from(new Set([...existingMute.roles_to_restore, ...rolesToRestore]));
+    }
+    await pluginData.state.mutes.updateExpiryTime(user.id, muteTime, rolesToRestore);
   } else {
-    await pluginData.state.mutes.addMute(user.id, muteTime);
+    await pluginData.state.mutes.addMute(user.id, muteTime, rolesToRestore);
   }
 
   const template = existingMute

backend/src/plugins/Mutes/functions/unmuteUser.ts

@@ -7,7 +7,7 @@ import humanizeDuration from "humanize-duration";
 import { CasesPlugin } from "../../Cases/CasesPlugin";
 import { CaseTypes } from "../../../data/CaseTypes";
 import { LogType } from "../../../data/LogType";
-import { WithRequiredProps } from "../../../utils/typeUtils";
+import { MemberOptions } from "eris";
 
 export async function unmuteUser(
   pluginData: GuildPluginData<MutesPluginType>,
@@ -36,6 +36,14 @@ export async function unmuteUser(
       if (muteRole && member.roles.includes(muteRole)) {
         await member.removeRole(muteRole);
       }
+      if (existingMute?.roles_to_restore) {
+        const memberOptions: MemberOptions = {};
+        const guildRoles = pluginData.guild.roles;
+        memberOptions.roles = Array.from(
+          new Set([...existingMute.roles_to_restore, ...member.roles.filter(x => x !== muteRole && guildRoles.has(x))]),
+        );
+        member.edit(memberOptions);
+      }
     } else {
       console.warn(
         `Member ${userId} not found in guild ${pluginData.guild.name} (${pluginData.guild.id}) when attempting to unmute`,

backend/src/plugins/Mutes/types.ts

@@ -23,6 +23,8 @@ export const ConfigSchema = t.type({
   mute_message: tNullable(t.string),
   timed_mute_message: tNullable(t.string),
   update_mute_message: tNullable(t.string),
+  remove_roles_on_mute: t.union([t.boolean, t.array(t.string)]),
+  restore_roles_on_mute: t.union([t.boolean, t.array(t.string)]),
 
   can_view_list: t.boolean,
   can_cleanup: t.boolean,

backend/src/plugins/Spam/types.ts

@@ -11,6 +11,8 @@ const BaseSingleSpamConfig = t.type({
   count: t.number,
   mute: tNullable(t.boolean),
   mute_time: tNullable(t.number),
+  remove_roles_on_mute: tNullable(t.union([t.boolean, t.array(t.string)])),
+  restore_roles_on_mute: tNullable(t.union([t.boolean, t.array(t.string)])),
   clean: tNullable(t.boolean),
 });
 export type TBaseSingleSpamConfig = t.TypeOf<typeof BaseSingleSpamConfig>;

backend/src/plugins/Spam/util/logAndDetectMessageSpam.ts

@@ -82,12 +82,19 @@ export async function logAndDetectMessageSpam(
             (spamConfig.mute_time && convertDelayStringToMS(spamConfig.mute_time.toString())) ?? 120 * 1000;
 
           try {
-            muteResult = await mutesPlugin.muteUser(member.id, muteTime, "Automatic spam detection", {
+            muteResult = await mutesPlugin.muteUser(
-              caseArgs: {
+              member.id,
-                modId: pluginData.client.user.id,
+              muteTime,
-                postInCaseLogOverride: false,
+              "Automatic spam detection",
+              {
+                caseArgs: {
+                  modId: pluginData.client.user.id,
+                  postInCaseLogOverride: false,
+                },
               },
-            });
+              spamConfig.remove_roles_on_mute,
+              spamConfig.restore_roles_on_mute,
+            );
           } catch (e) {
             if (e instanceof RecoverablePluginError && e.code === ERRORS.NO_MUTE_ROLE_IN_CONFIG) {
               logs.log(LogType.BOT_ALERT, {

backend/src/plugins/Spam/util/logAndDetectOtherSpam.ts

@@ -41,12 +41,19 @@ export async function logAndDetectOtherSpam(
           (spamConfig.mute_time && convertDelayStringToMS(spamConfig.mute_time.toString())) ?? 120 * 1000;
 
         try {
-          await mutesPlugin.muteUser(member.id, muteTime, "Automatic spam detection", {
+          await mutesPlugin.muteUser(
-            caseArgs: {
+            member.id,
-              modId: pluginData.client.user.id,
+            muteTime,
-              extraNotes: [`Details: ${details}`],
+            "Automatic spam detection",
+            {
+              caseArgs: {
+                modId: pluginData.client.user.id,
+                extraNotes: [`Details: ${details}`],
+              },
             },
-          });
+            spamConfig.remove_roles_on_mute,
+            spamConfig.restore_roles_on_mute,
+          );
         } catch (e) {
           if (e instanceof RecoverablePluginError && e.code === ERRORS.NO_MUTE_ROLE_IN_CONFIG) {
             logs.log(LogType.BOT_ALERT, {