feat(handlers): add maintenance mode checks for commands and DMs

Author: Akinator31

rustmail/src/handlers/guild_interaction_handler.rs

@@ -3,8 +3,10 @@ use crate::prelude::config::*;
 use crate::prelude::features::*;
 use crate::prelude::modules::*;
 use crate::prelude::types::*;
-use serenity::all::{Context, EventHandler, Interaction};
+use crate::utils::{MessageBuilder, defer_response};
+use serenity::all::{Context, EventHandler, Interaction, Permissions};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use tokio::sync::watch::Receiver;
 
 #[derive(Clone)]
@@ -13,6 +15,7 @@ pub struct InteractionHandler {
     pub registry: Arc<CommandRegistry>,
     pub shutdown: Arc<Receiver<bool>>,
     pub pagination: PaginationStore,
+    pub maintenance_mode: Arc<AtomicBool>,
 }
 
 impl InteractionHandler {
@@ -21,16 +24,59 @@ impl InteractionHandler {
         registry: Arc<CommandRegistry>,
         shutdown: Receiver<bool>,
         pagination: PaginationStore,
+        maintenance_mode: Arc<AtomicBool>,
     ) -> Self {
         Self {
             config: Arc::new(config.clone()),
             registry,
             shutdown: Arc::new(shutdown),
             pagination,
+            maintenance_mode,
         }
     }
 }
 
+async fn is_interaction_user_maintenance_exempt(
+    ctx: &Context,
+    guild_id: serenity::all::GuildId,
+    user_id: serenity::all::UserId,
+    config: &Config,
+) -> bool {
+    if config.bot.panel_super_admin_users.contains(&user_id.get()) {
+        return true;
+    }
+
+    let member = match guild_id.member(&ctx.http, user_id).await {
+        Ok(m) => m,
+        Err(_) => return false,
+    };
+
+    if !config.bot.panel_super_admin_roles.is_empty() {
+        for role_id in &member.roles {
+            if config.bot.panel_super_admin_roles.contains(&role_id.get()) {
+                return true;
+            }
+        }
+    }
+
+    let guild = match guild_id.to_partial_guild(&ctx.http).await {
+        Ok(g) => g,
+        Err(_) => return false,
+    };
+
+    if guild.owner_id == user_id {
+        return true;
+    }
+
+    member.roles.iter().any(|role_id| {
+        guild
+            .roles
+            .get(role_id)
+            .map(|role| role.permissions.contains(Permissions::ADMINISTRATOR))
+            .unwrap_or(false)
+    })
+}
+
 #[async_trait::async_trait]
 impl EventHandler for InteractionHandler {
     async fn interaction_create(&self, ctx: Context, interaction: Interaction) {
@@ -65,6 +111,33 @@ impl EventHandler for InteractionHandler {
                 }
             }
             Interaction::Command(command) => {
+                if self.maintenance_mode.load(Ordering::Relaxed) {
+                    if let Some(guild_id) = command.guild_id {
+                        if !is_interaction_user_maintenance_exempt(
+                            &ctx,
+                            guild_id,
+                            command.user.id,
+                            &self.config,
+                        )
+                        .await
+                        {
+                            defer_response(&ctx, &command).await.ok();
+
+                            let _ = MessageBuilder::system_message(&ctx, &self.config)
+                                .translated_content(
+                                    "status.maintenance_mode_active",
+                                    None,
+                                    Some(command.user.id),
+                                    command.guild_id.map(|g| g.get()),
+                                )
+                                .await
+                                .send_interaction_followup(&command, true)
+                                .await;
+                            return;
+                        }
+                    }
+                }
+
                 let ctx = ctx.clone();
                 let command = command.clone();
                 let options = command.data.options().clone();

rustmail/src/handlers/guild_messages_handler.rs

@@ -7,12 +7,13 @@ use crate::prelude::modules::*;
 use crate::prelude::types::*;
 use crate::prelude::utils::*;
 use crate::wrap_command;
-use serenity::all::{GuildId, MessageId, UserId};
+use serenity::all::{GuildId, MessageId, Permissions, UserId};
 use serenity::{
     all::{ChannelId, Context, EventHandler, Message, MessageUpdateEvent},
     async_trait,
 };
 use std::collections::HashSet;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{LazyLock, Mutex};
 use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
 use tokio::sync::Mutex as AsyncMutex;
@@ -39,6 +40,7 @@ pub struct GuildMessagesHandler {
     pub registry: Arc<CommandRegistry>,
     pub shutdown: Arc<Receiver<bool>>,
     pub pagination: PaginationStore,
+    pub maintenance_mode: Arc<AtomicBool>,
 }
 
 impl GuildMessagesHandler {
@@ -47,13 +49,15 @@ impl GuildMessagesHandler {
         registry: Arc<CommandRegistry>,
         shutdown: Receiver<bool>,
         pagination: PaginationStore,
+        maintenance_mode: Arc<AtomicBool>,
     ) -> Self {
         let h = Self {
             config: Arc::new(config.clone()),
             commands: Arc::new(AsyncMutex::new(HashMap::new())),
             registry,
             shutdown: Arc::new(shutdown),
             pagination,
+            maintenance_mode,
         };
 
         let mut lock = h.commands.lock().await;
@@ -198,10 +202,68 @@ async fn manage_incoming_message(
     Ok(())
 }
 
+async fn is_user_maintenance_exempt(ctx: &Context, msg: &Message, config: &Config) -> bool {
+    let user_id = msg.author.id.get();
+
+    if config.bot.panel_super_admin_users.contains(&user_id) {
+        return true;
+    }
+
+    let guild_id = match msg.guild_id {
+        Some(id) => id,
+        None => return false,
+    };
+
+    let member = match guild_id.member(&ctx.http, msg.author.id).await {
+        Ok(m) => m,
+        Err(_) => return false,
+    };
+
+    if !config.bot.panel_super_admin_roles.is_empty() {
+        for role_id in &member.roles {
+            if config.bot.panel_super_admin_roles.contains(&role_id.get()) {
+                return true;
+            }
+        }
+    }
+
+    let guild = match guild_id.to_partial_guild(&ctx.http).await {
+        Ok(g) => g,
+        Err(_) => return false,
+    };
+
+    if guild.owner_id == msg.author.id {
+        return true;
+    }
+
+    member.roles.iter().any(|role_id| {
+        guild
+            .roles
+            .get(role_id)
+            .map(|role| role.permissions.contains(Permissions::ADMINISTRATOR))
+            .unwrap_or(false)
+    })
+}
+
 #[async_trait]
 impl EventHandler for GuildMessagesHandler {
     async fn message(&self, ctx: Context, msg: Message) {
         if msg.guild_id.is_none() {
+            if self.maintenance_mode.load(Ordering::Relaxed) {
+                let _ = MessageBuilder::system_message(&ctx, &self.config)
+                    .translated_content(
+                        "status.maintenance_mode_active_user",
+                        None,
+                        Some(msg.author.id),
+                        None,
+                    )
+                    .await
+                    .to_user(msg.author.id)
+                    .send(true)
+                    .await;
+                return;
+            }
+
             if let Err(error) = manage_incoming_message(&ctx, &msg, &self.config).await {
                 if let Some(error_handler) = &self.config.error_handler {
                     let _ = error_handler
@@ -222,6 +284,18 @@ impl EventHandler for GuildMessagesHandler {
                 command_name = &message_content[self.config.command.prefix.len()..i];
             }
 
+            if self.maintenance_mode.load(Ordering::Relaxed) {
+                if !is_user_maintenance_exempt(&ctx, &msg, &self.config).await {
+                    let _ = MessageBuilder::system_message(&ctx, &self.config)
+                        .translated_content("status.maintenance_mode_active", None, None, None)
+                        .await
+                        .to_channel(msg.channel_id)
+                        .send(true)
+                        .await;
+                    return;
+                }
+            }
+
             let commands_lock = self.commands.lock().await;
 
             if let Some(command_func) = commands_lock.get(command_name)