refactor(addmod): allow adding all users with a specific role

crates/rustmail/src/commands/add_staff/common.rs

@@ -1,10 +1,28 @@
 use crate::prelude::config::*;
 use crate::prelude::errors::*;
 use serenity::all::{
-    ChannelId, Context, Message, PermissionOverwrite, PermissionOverwriteType, UserId,
+    ChannelId, Context, GuildId, Message, PermissionOverwrite, PermissionOverwriteType, RoleId,
+    UserId,
 };
 use serenity::model::Permissions;
 
+pub const MAX_ROLE_MEMBERS_PER_ADD: usize = 50;
+
+pub enum AddTarget {
+    User(UserId),
+    Role(RoleId),
+}
+
+pub enum AddTargetParse {
+    Explicit(AddTarget),
+    AmbiguousId(u64),
+}
+
+pub struct AddRoleOutcome {
+    pub added: Vec<UserId>,
+    pub failed: Vec<UserId>,
+}
+
 pub async fn add_user_to_channel(
     ctx: &Context,
     channel_id: ChannelId,
@@ -26,6 +44,83 @@ pub async fn add_user_to_channel(
     Ok(())
 }
 
+pub fn parse_add_target(raw: &str) -> Option<AddTargetParse> {
+    let s = raw.trim();
+    if let Some(inner) = s.strip_prefix("<@&").and_then(|s| s.strip_suffix('>')) {
+        return inner
+            .parse::<u64>()
+            .ok()
+            .map(|id| AddTargetParse::Explicit(AddTarget::Role(RoleId::new(id))));
+    }
+    if let Some(inner) = s.strip_prefix("<@").and_then(|s| s.strip_suffix('>')) {
+        let inner = inner.strip_prefix('!').unwrap_or(inner);
+        return inner
+            .parse::<u64>()
+            .ok()
+            .map(|id| AddTargetParse::Explicit(AddTarget::User(UserId::new(id))));
+    }
+    s.parse::<u64>().ok().map(AddTargetParse::AmbiguousId)
+}
+
+pub async fn members_with_role(
+    ctx: &Context,
+    guild_id: GuildId,
+    role_id: RoleId,
+) -> ModmailResult<Vec<UserId>> {
+    const PAGE_LIMIT: u64 = 1000;
+
+    let mut result = Vec::new();
+    let mut after: Option<UserId> = None;
+
+    loop {
+        let page = guild_id
+            .members(&ctx.http, Some(PAGE_LIMIT), after)
+            .await
+            .map_err(|_| {
+                ModmailError::Discord(DiscordError::ApiError(
+                    "Failed to fetch guild members".to_string(),
+                ))
+            })?;
+
+        if page.is_empty() {
+            break;
+        }
+
+        let last_id = page.last().map(|m| m.user.id);
+        let page_len = page.len();
+
+        for member in page {
+            if member.roles.contains(&role_id) {
+                result.push(member.user.id);
+            }
+        }
+
+        if page_len < PAGE_LIMIT as usize {
+            break;
+        }
+
+        after = last_id;
+    }
+
+    Ok(result)
+}
+
+pub async fn add_role_members_to_channel(
+    ctx: &Context,
+    channel_id: ChannelId,
+    members: Vec<UserId>,
+) -> AddRoleOutcome {
+    let mut added = Vec::new();
+    let mut failed = Vec::new();
+    for user_id in members {
+        match add_user_to_channel(ctx, channel_id, user_id).await {
+            Ok(_) => added.push(user_id),
+            Err(_) => failed.push(user_id),
+        }
+    }
+    AddRoleOutcome { added, failed }
+}
+
 pub async fn extract_staff_id(msg: &Message, config: &Config) -> String {
     let content = msg.content.trim();
     let prefix = &config.command.prefix;

crates/rustmail/src/commands/add_staff/slash_command/add_staff.rs

@@ -8,7 +8,7 @@ use crate::prelude::utils::*;
 use serenity::FutureExt;
 use serenity::all::{
     CommandDataOptionValue, CommandInteraction, CommandOptionType, Context, CreateCommand,
-    CreateCommandOption, ResolvedOption,
+    CreateCommandOption, GuildId, ResolvedOption, RoleId, UserId,
 };
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -41,9 +41,9 @@ impl RegistrableCommand for AddStaffCommand {
             )
             .await;
 
-            let user_id_desc = get_translated_message(
+            let target_desc = get_translated_message(
                 &config,
-                "slash_command.add_staff_user_id_argument",
+                "slash_command.add_staff_target_argument",
                 None,
                 None,
                 None,
@@ -53,7 +53,7 @@ impl RegistrableCommand for AddStaffCommand {
 
             vec![
                 CreateCommand::new(name).description(cmd_desc).add_option(
-                    CreateCommandOption::new(CommandOptionType::User, "user_id", user_id_desc)
+                    CreateCommandOption::new(CommandOptionType::Mentionable, "target", target_desc)
                         .required(true),
                 ),
             ]
@@ -80,51 +80,144 @@ impl RegistrableCommand for AddStaffCommand {
 
             defer_response(&ctx, &command).await?;
 
-            let user_id = match command
-                .data
-                .options
-                .iter()
-                .find(|opt| opt.name == "user_id")
-            {
-                Some(opt) => match &opt.value {
-                    CommandDataOptionValue::User(user_id) => *user_id,
-                    _ => {
-                        return Err(ModmailError::Command(CommandError::InvalidArguments(
-                            "user_id".to_string(),
-                        )));
-                    }
-                },
-                None => {
-                    if let Some(user_id) = command.data.target_id {
-                        user_id.to_user_id()
-                    } else {
-                        return Err(ModmailError::Command(CommandError::InvalidArguments(
-                            "user_id".to_string(),
-                        )));
-                    }
-                }
-            };
-
-            if thread_exists_by_channel(command.channel_id, pool).await {
-                match add_user_to_channel(&ctx, command.channel_id, user_id).await {
-                    Ok(_) => {
-                        let mut params = HashMap::new();
-                        params.insert("user".to_string(), format!("<@{}>", user_id));
-
-                        let _ = MessageBuilder::system_message(&ctx, &config)
-                            .translated_content("add_staff.add_success", Some(&params), None, None)
-                            .await
-                            .to_channel(command.channel_id)
-                            .send_interaction_followup(&command, true)
-                            .await;
-
-                        Ok(())
-                    }
-                    Err(..) => Err(ModmailError::Command(CommandError::InvalidFormat)),
-                }
-            } else {
-                Err(ModmailError::Thread(ThreadError::NotAThreadChannel))
+            if !thread_exists_by_channel(command.channel_id, pool).await {
+                return Err(ModmailError::Thread(ThreadError::NotAThreadChannel));
+            }
+
+            let target = resolve_target(&command)?;
+
+            match target {
+                AddTarget::User(user_id) => add_single_user(&ctx, &config, &command, user_id).await,
+                AddTarget::Role(role_id) => add_role(&ctx, &config, &command, role_id).await,
             }
         })
     }
 }
+
+fn resolve_target(command: &CommandInteraction) -> ModmailResult<AddTarget> {
+    let opt = command
+        .data
+        .options
+        .iter()
+        .find(|opt| opt.name == "target")
+        .ok_or_else(|| {
+            ModmailError::Command(CommandError::InvalidArguments("target".to_string()))
+        })?;
+
+    let id = match &opt.value {
+        CommandDataOptionValue::Mentionable(id) => id.get(),
+        CommandDataOptionValue::User(user_id) => return Ok(AddTarget::User(*user_id)),
+        CommandDataOptionValue::Role(role_id) => return Ok(AddTarget::Role(*role_id)),
+        _ => {
+            return Err(ModmailError::Command(CommandError::InvalidArguments(
+                "target".to_string(),
+            )));
+        }
+    };
+
+    if command.data.resolved.users.contains_key(&UserId::new(id)) {
+        Ok(AddTarget::User(UserId::new(id)))
+    } else if command.data.resolved.roles.contains_key(&RoleId::new(id)) {
+        Ok(AddTarget::Role(RoleId::new(id)))
+    } else {
+        Err(ModmailError::Command(CommandError::InvalidArguments(
+            "target".to_string(),
+        )))
+    }
+}
+
+async fn add_single_user(
+    ctx: &Context,
+    config: &Config,
+    command: &CommandInteraction,
+    user_id: UserId,
+) -> ModmailResult<()> {
+    add_user_to_channel(ctx, command.channel_id, user_id).await?;
+
+    let mut params = HashMap::new();
+    params.insert("user".to_string(), format!("<@{}>", user_id));
+
+    let _ = MessageBuilder::system_message(ctx, config)
+        .translated_content("add_staff.add_success", Some(&params), None, None)
+        .await
+        .to_channel(command.channel_id)
+        .send_interaction_followup(command, true)
+        .await;
+
+    Ok(())
+}
+
+async fn add_role(
+    ctx: &Context,
+    config: &Config,
+    command: &CommandInteraction,
+    role_id: RoleId,
+) -> ModmailResult<()> {
+    let guild_id = command
+        .guild_id
+        .unwrap_or_else(|| GuildId::new(config.bot.get_staff_guild_id()));
+
+    if role_id.get() == guild_id.get() {
+        let _ = MessageBuilder::system_message(ctx, config)
+            .translated_content("add_staff.role_everyone_forbidden", None, None, None)
+            .await
+            .to_channel(command.channel_id)
+            .send_interaction_followup(command, true)
+            .await;
+        return Ok(());
+    }
+
+    let role_mention = format!("<@&{}>", role_id);
+    let members = members_with_role(ctx, guild_id, role_id).await?;
+
+    if members.is_empty() {
+        let mut params = HashMap::new();
+        params.insert("role".to_string(), role_mention);
+        let _ = MessageBuilder::system_message(ctx, config)
+            .translated_content("add_staff.role_no_members", Some(&params), None, None)
+            .await
+            .to_channel(command.channel_id)
+            .send_interaction_followup(command, true)
+            .await;
+        return Ok(());
+    }
+
+    if members.len() > MAX_ROLE_MEMBERS_PER_ADD {
+        let mut params = HashMap::new();
+        params.insert("role".to_string(), role_mention);
+        params.insert("count".to_string(), members.len().to_string());
+        params.insert("max".to_string(), MAX_ROLE_MEMBERS_PER_ADD.to_string());
+        let _ = MessageBuilder::system_message(ctx, config)
+            .translated_content("add_staff.role_too_many", Some(&params), None, None)
+            .await
+            .to_channel(command.channel_id)
+            .send_interaction_followup(command, true)
+            .await;
+        return Ok(());
+    }
+
+    let total = members.len();
+    let outcome = add_role_members_to_channel(ctx, command.channel_id, members).await;
+
+    let key = if outcome.failed.is_empty() {
+        "add_staff.role_add_success"
+    } else {
+        "add_staff.role_add_partial"
+    };
+
+    let mut params = HashMap::new();
+    params.insert("role".to_string(), role_mention);
+    params.insert("count".to_string(), outcome.added.len().to_string());
+    params.insert("added".to_string(), outcome.added.len().to_string());
+    params.insert("total".to_string(), total.to_string());
+    params.insert("failed".to_string(), outcome.failed.len().to_string());
+
+    let _ = MessageBuilder::system_message(ctx, config)
+        .translated_content(key, Some(&params), None, None)
+        .await
+        .to_channel(command.channel_id)
+        .send_interaction_followup(command, true)
+        .await;
+
+    Ok(())
+}