Add support for extras, conditionally displaying commands in the help command, and start working on moderation commands

src/groups/moderation/arguments.rs

@@ -0,0 +1,148 @@
+use std::ops::{Deref, DerefMut};
+
+use scylla::client::session::Session;
+
+use crate::{
+    database::models::{Message, User},
+    extras::Extras,
+    framework::commands::{
+        arguments::{ArgumentParsingError, FromArguments, NeedyString},
+        context::CommandContext,
+    },
+    state::State,
+};
+
+/// Fetch a user by their ID from the database.
+///
+/// Arguments:
+/// * `id` - The user's ID.
+/// * `db` - The database session.
+///
+/// Returns:
+/// [`Result<User, ArgumentParser>`] - The user from the DB, or an [`ArgumentParsingError`] ready
+/// to be returned by the parser.
+async fn get_user_by_id(id: u64, db: &Session) -> Result<User, ArgumentParsingError> {
+    let Some(user) = User::get_by_id(id, db)
+        .await
+        .map_err(|_| ArgumentParsingError::DatabaseFailed)?
+    else {
+        return Err(ArgumentParsingError::InvalidArgument);
+    };
+    Ok(user)
+}
+
+/// Fetches the user from the message's reply and falls back to the command executor.
+///
+/// This function takes into account messages on call and traces them back to the original author.
+///
+/// Arguments:
+/// * `ctx` - The command context.
+///
+/// Returns:
+/// [`Result<User, ArgumentParsingError>`] - The user from the message's reply, or an
+/// [`ArgumentParsingError`] ready to be returned by the parser.
+async fn user_from_reply_or_command_executor(
+    ctx: &CommandContext<'_, State, Extras>,
+) -> Result<User, ArgumentParsingError> {
+    if let Some(referenced) = &ctx.event.referenced_message {
+        if !referenced.author.system.unwrap_or(false) {
+            if referenced.author.bot {
+                if let Some(message) =
+                    Message::get_by_propagation_message_id(referenced.id.get(), &ctx.rt.state.db)
+                        .await
+                        .map_err(|_| ArgumentParsingError::DatabaseFailed)?
+                {
+                    return get_user_by_id(message.author_id, &ctx.rt.state.db).await;
+                }
+            }
+
+            // If the replied message's author is not a bot or system user, use them.
+            return get_user_by_id(referenced.author.id.get(), &ctx.rt.state.db).await;
+        }
+    }
+
+    // fallback to command executor
+    get_user_by_id(ctx.event.author.id.get(), &ctx.rt.state.db).await
+}
+
+/// Parses a user mention or ID into a UserArgument.
+///
+/// It parses from multiple sources, attempting in the following way:
+///
+/// 1. If the argument is a Discord mention, parsing it.
+/// 2. If the argument is a Discord ID, finding it in the DB.
+/// 3. If the argument is a Discord mention, finding it in the DB.
+/// 4. If there's no argument and the message replies to another message:
+///    a. If the replied message's author is a bot, attempts to query the user who sent it from the
+///       other side of the call. If it fails, the command executor is used.
+///    b. If the replied message's author is not a bot, they're used.
+/// 5. If there's no argument and the message is not replying to any message, the command executor
+///    is used.
+pub struct UserArgument(pub User);
+
+impl FromArguments<State, Extras> for UserArgument {
+    async fn from_args<'a>(
+        args: &'a str,
+        ctx: &'a CommandContext<'_, State, Extras>,
+    ) -> Result<(Self, &'a str), ArgumentParsingError> {
+        let user: User;
+
+        if !args.is_empty() {
+            // There's an argument, so that argument is used to find the user.
+
+            let (identifier, args) = NeedyString::from_args(args, ctx).await?;
+
+            if identifier.starts_with("<@") && identifier.ends_with(">") {
+                // The argument is a Discord mention.
+
+                let user_id: u64 = identifier[2..identifier.len() - 1]
+                    .to_string()
+                    .parse()
+                    .map_err(|_| ArgumentParsingError::InvalidArgument)?;
+
+                user = get_user_by_id(user_id, &ctx.rt.state.db).await?;
+            } else if identifier.chars().all(|c| c.is_ascii_digit()) {
+                // The argument is a user ID.
+
+                let user_id: u64 = identifier
+                    .parse()
+                    .map_err(|_| ArgumentParsingError::InvalidArgument)?;
+
+                user = get_user_by_id(user_id, &ctx.rt.state.db).await?;
+            } else {
+                // The argument is (probably) a username.
+
+                let Some(db_user) = User::get_by_name(identifier.0, &ctx.rt.state.db)
+                    .await
+                    .map_err(|_| ArgumentParsingError::DatabaseFailed)?
+                else {
+                    // Check in the message's reply, and if not use the command executor.
+                    user = user_from_reply_or_command_executor(ctx).await?;
+
+                    return Ok((UserArgument(user), args));
+                };
+
+                user = db_user;
+            }
+        } else {
+            // Check in the message's reply, and if not use the command executor.
+            user = user_from_reply_or_command_executor(ctx).await?;
+        }
+
+        Ok((UserArgument(user), args))
+    }
+}
+
+impl Deref for UserArgument {
+    type Target = User;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for UserArgument {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}