added functionality for reusable hook functions that will execute on commands

src/hooks.rs

@@ -0,0 +1,217 @@
+use log::warn;
+use regex_command_attr::check;
+use serenity::{client::Context, model::channel::Channel};
+
+use crate::{
+    framework::{CommandInvoke, CommandOptions, CreateGenericResponse, HookResult},
+    moderation_cmds, RecordingMacros, SQLPool,
+};
+
+#[check]
+pub async fn macro_check(
+    ctx: &Context,
+    invoke: &CommandInvoke,
+    args: &CommandOptions,
+) -> HookResult {
+    if let Some(guild_id) = invoke.guild_id() {
+        if args.command != moderation_cmds::MACRO_CMD_COMMAND.names[0] {
+            let active_recordings =
+                ctx.data.read().await.get::<RecordingMacros>().cloned().unwrap();
+            let mut lock = active_recordings.write().await;
+
+            if let Some(command_macro) = lock.get_mut(&(guild_id, invoke.author_id())) {
+                command_macro.commands.push(args.clone());
+
+                let _ = invoke
+                    .respond(
+                        &ctx,
+                        CreateGenericResponse::new().content("Command recorded to macro"),
+                    )
+                    .await;
+
+                HookResult::Halt
+            } else {
+                HookResult::Continue
+            }
+        } else {
+            HookResult::Continue
+        }
+    } else {
+        HookResult::Continue
+    }
+}
+
+#[check]
+pub async fn check_self_permissions(
+    ctx: &Context,
+    invoke: &CommandInvoke,
+    _args: &CommandOptions,
+) -> HookResult {
+    if let Some(guild) = invoke.guild(&ctx) {
+        let user_id = ctx.cache.current_user_id();
+
+        let manage_webhooks =
+            guild.member_permissions(&ctx, user_id).await.map_or(false, |p| p.manage_webhooks());
+        let (send_messages, embed_links) = invoke
+            .channel_id()
+            .to_channel_cached(&ctx)
+            .map(|c| {
+                if let Channel::Guild(channel) = c {
+                    channel.permissions_for_user(ctx, user_id).ok()
+                } else {
+                    None
+                }
+            })
+            .flatten()
+            .map_or((false, false), |p| (p.send_messages(), p.embed_links()));
+
+        if manage_webhooks && send_messages && embed_links {
+            HookResult::Continue
+        } else {
+            if send_messages {
+                let _ = invoke
+                    .respond(
+                        &ctx,
+                        CreateGenericResponse::new().content(format!(
+                            "Please ensure the bot has the correct permissions:
+
+✅     **Send Message**
+{}     **Embed Links**
+{}     **Manage Webhooks**",
+                            if manage_webhooks { "✅" } else { "❌" },
+                            if embed_links { "✅" } else { "❌" },
+                        )),
+                    )
+                    .await;
+            } else {
+                warn!("Missing permissions in guild {}", guild.id);
+            }
+
+            HookResult::Halt
+        }
+    } else {
+        HookResult::Continue
+    }
+}
+
+#[check]
+pub async fn check_managed_permissions(
+    ctx: &Context,
+    invoke: &CommandInvoke,
+    args: &CommandOptions,
+) -> HookResult {
+    if let Some(guild) = invoke.guild(&ctx) {
+        let permissions = guild.member_permissions(&ctx, invoke.author_id()).await.unwrap();
+
+        if permissions.manage_messages() {
+            return HookResult::Continue;
+        }
+
+        let member = invoke.member(&ctx).await.unwrap();
+
+        let pool = ctx
+            .data
+            .read()
+            .await
+            .get::<SQLPool>()
+            .cloned()
+            .expect("Could not get SQLPool from data");
+
+        match sqlx::query!(
+            "
+SELECT
+    role
+FROM
+    roles
+INNER JOIN
+    command_restrictions ON roles.id = command_restrictions.role_id
+WHERE
+    command_restrictions.command = ? AND
+    roles.guild_id = (
+        SELECT
+            id
+        FROM
+            guilds
+        WHERE
+            guild = ?)
+                    ",
+            args.command,
+            guild.id.as_u64()
+        )
+        .fetch_all(&pool)
+        .await
+        {
+            Ok(rows) => {
+                let role_ids = member.roles.iter().map(|r| *r.as_u64()).collect::<Vec<u64>>();
+
+                for row in rows {
+                    if role_ids.contains(&row.role) {
+                        return HookResult::Continue;
+                    }
+                }
+
+                let _ = invoke
+                    .respond(
+                        &ctx,
+                        CreateGenericResponse::new().content(
+                            "You must have \"Manage Messages\" or have a role capable of sending reminders to that channel. \
+Please talk to your server admin, and ask them to use the `/restrict` command to specify allowed roles.",
+                        ),
+                    )
+                    .await;
+
+                HookResult::Halt
+            }
+
+            Err(sqlx::Error::RowNotFound) => {
+                let _ = invoke
+                    .respond(
+                        &ctx,
+                        CreateGenericResponse::new().content(
+                            "You must have \"Manage Messages\" or have a role capable of sending reminders to that channel. \
+Please talk to your server admin, and ask them to use the `/restrict` command to specify allowed roles.",
+                        ),
+                    )
+                    .await;
+
+                HookResult::Halt
+            }
+
+            Err(e) => {
+                warn!("Unexpected error occurred querying command_restrictions: {:?}", e);
+
+                HookResult::Halt
+            }
+        }
+    } else {
+        HookResult::Continue
+    }
+}
+
+#[check]
+pub async fn check_guild_permissions(
+    ctx: &Context,
+    invoke: &CommandInvoke,
+    _args: &CommandOptions,
+) -> HookResult {
+    if let Some(guild) = invoke.guild(&ctx) {
+        let permissions = guild.member_permissions(&ctx, invoke.author_id()).await.unwrap();
+
+        if !permissions.manage_guild() {
+            let _ = invoke
+                .respond(
+                    &ctx,
+                    CreateGenericResponse::new().content(
+                        "You must have the \"Manage Server\" permission to use this command",
+                    ),
+                )
+                .await;
+
+            HookResult::Halt
+        } else {
+            HookResult::Continue
+        }
+    } else {
+        HookResult::Continue
+    }
+}