diff --git a/Cargo.lock b/Cargo.lock index 5347af0..90729c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3426,6 +3426,7 @@ dependencies = [ "dashmap 5.5.3", "flume 0.10.14", "futures-util", + "poise", "rand", "serde", "serenity", diff --git a/youmubot-core/Cargo.toml b/youmubot-core/Cargo.toml index 8b8f26c..ac00c6b 100644 --- a/youmubot-core/Cargo.toml +++ b/youmubot-core/Cargo.toml @@ -19,6 +19,8 @@ futures-util = "0.3.21" tokio = { version = "1.19.2", features = ["time"] } flume = "0.10.13" dashmap = "5.3.4" +poise = { git = "https://github.com/serenity-rs/poise", branch = "current" } youmubot-db = { path = "../youmubot-db" } +youmubot-db-sql = { path = "../youmubot-db-sql" } youmubot-prelude = { path = "../youmubot-prelude" } diff --git a/youmubot-core/src/admin/ignore.rs b/youmubot-core/src/admin/ignore.rs new file mode 100644 index 0000000..379bf0e --- /dev/null +++ b/youmubot-core/src/admin/ignore.rs @@ -0,0 +1,152 @@ +use std::sync::Arc; + +use chrono::{DateTime, Utc}; +use dashmap::DashMap; +use serenity::all::User; +use youmubot_db_sql::models::ignore_list as model; +use youmubot_prelude::*; + +use crate::HasCoreEnv; + +// Should we ignore this user? +pub fn should_ignore(env: &impl HasCoreEnv, id: UserId) -> bool { + env.core_env().ignore.query(id).is_some() +} + +/// Ignore: make Youmu ignore all commands from an user. +#[poise::command( + slash_command, + subcommands("add", "remove", "list"), + owners_only, + install_context = "User", + interaction_context = "Guild|BotDm" +)] +pub async fn ignore(_ctx: CmdContext<'_, U>) -> Result<()> { + Ok(()) +} + +/// Add an user to ignore list. +#[poise::command(slash_command, owners_only)] +async fn add( + ctx: CmdContext<'_, U>, + #[description = "Discord username"] discord_name: User, +) -> Result<()> { + let env = ctx.data().core_env(); + ctx.defer().await?; + let msg = format!("User **{}** ignored!", discord_name.name); + env.ignore + .add(&env.prelude, UserId(discord_name.id), discord_name.name) + .await?; + ctx.say(msg).await?; + Ok(()) +} + +/// Remove an user from ignore list. +#[poise::command(slash_command, owners_only)] +async fn remove( + ctx: CmdContext<'_, U>, + #[description = "Discord username"] discord_name: User, +) -> Result<()> { + let env = ctx.data().core_env(); + ctx.defer().await?; + env.ignore + .remove(&env.prelude, UserId(discord_name.id)) + .await?; + let msg = format!("User **{}** removed from ignore list!", discord_name.name); + ctx.say(msg).await?; + Ok(()) +} + +/// List ignored users. +#[poise::command(slash_command, owners_only)] +async fn list(ctx: CmdContext<'_, U>) -> Result<()> { + let env = ctx.data().core_env(); + let is_dm = ctx.guild_id().is_none(); + ctx.defer().await?; + let users = env + .ignore + .list + .clone() + .iter() + .map(|v| { + format!( + "- {} ({}), since ", + v.username, + if is_dm { + v.id.0.mention().to_string() + } else { + format!("`{}`", v.id.0.get()) + }, + v.ignored_since.timestamp(), + ) + }) + .collect::>() + .join("\n"); + let users = if users == "" { + "No one is being ignored!" + } else { + &users[..] + }; + + let msg = format!("Ignored users:\n{}", users); + ctx.say(msg).await?; + Ok(()) +} + +#[derive(Debug, Clone)] +pub(crate) struct IgnoredUsers { + list: Arc>, +} + +impl IgnoredUsers { + pub async fn from_db(env: &Env) -> Result { + let list = model::IgnoredUser::get_all(&env.sql).await?; + let mp: DashMap<_, _> = list + .into_iter() + .map(|v| { + let id = (v.id as u64).into(); + ( + id, + IgnoredUser { + id, + username: v.username, + ignored_since: v.ignored_since, + }, + ) + }) + .collect(); + Ok(Self { list: Arc::new(mp) }) + } + + pub fn query<'a>( + &'a self, + id: UserId, + ) -> Option + 'a> { + self.list.get(&id) + } + + pub async fn add(&self, env: &Env, id: UserId, username: String) -> Result<()> { + let iu = model::IgnoredUser::add(&env.sql, id.0.get() as i64, username).await?; + self.list.insert( + id, + IgnoredUser { + id, + username: iu.username, + ignored_since: iu.ignored_since, + }, + ); + Ok(()) + } + + pub async fn remove(&self, env: &Env, id: UserId) -> Result { + model::IgnoredUser::remove(&env.sql, id.0.get() as i64).await?; + Ok(self.list.remove(&id).is_some()) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct IgnoredUser { + pub id: UserId, + pub username: String, + pub ignored_since: DateTime, +} diff --git a/youmubot-core/src/admin/mod.rs b/youmubot-core/src/admin/mod.rs index 2c766e2..6680c0a 100644 --- a/youmubot-core/src/admin/mod.rs +++ b/youmubot-core/src/admin/mod.rs @@ -13,6 +13,8 @@ use youmubot_prelude::*; mod soft_ban; pub use soft_ban::watch_soft_bans; +pub mod ignore; + #[group] #[description = "Administrative commands for the server."] #[commands(clean, ban, kick, soft_ban, soft_ban_init)] diff --git a/youmubot-core/src/lib.rs b/youmubot-core/src/lib.rs index ce32b17..00f2783 100644 --- a/youmubot-core/src/lib.rs +++ b/youmubot-core/src/lib.rs @@ -17,11 +17,32 @@ pub mod community; mod db; pub mod fun; +#[derive(Debug, Clone)] +pub struct CoreEnv { + pub(crate) prelude: Env, + pub(crate) ignore: admin::ignore::IgnoredUsers, +} + +impl CoreEnv { + async fn new(prelude: Env) -> Result { + let ignore = admin::ignore::IgnoredUsers::from_db(&prelude).await?; + Ok(Self { prelude, ignore }) + } +} + +/// Gets an [CoreEnv] from the current environment. +pub trait HasCoreEnv: Send + Sync { + fn core_env(&self) -> &CoreEnv; +} + +impl + Send + Sync> HasCoreEnv for T { + fn core_env(&self) -> &CoreEnv { + self.as_ref() + } +} + /// Sets up all databases in the client. -pub fn setup( - path: &std::path::Path, - data: &mut TypeMap, -) -> serenity::framework::standard::CommandResult { +pub async fn setup(path: &std::path::Path, data: &mut TypeMap, prelude: Env) -> Result { db::SoftBans::insert_into(&mut *data, &path.join("soft_bans.yaml"))?; db::load_role_list( &mut *data, @@ -32,7 +53,7 @@ pub fn setup( // Start reaction handlers data.insert::(community::ReactionWatchers::new(&*data)?); - Ok(()) + CoreEnv::new(prelude).await } pub fn ready_hook(ctx: &Context) -> CommandResult { diff --git a/youmubot-prelude/src/args.rs b/youmubot-prelude/src/args.rs index c6410ce..6803dfa 100644 --- a/youmubot-prelude/src/args.rs +++ b/youmubot-prelude/src/args.rs @@ -189,9 +189,15 @@ mod ids { use super::ParseError; /// An `UserId` parsed the old way. - #[derive(Debug, Clone, Copy, PartialEq, Eq)] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct UserId(pub id::UserId); + impl From for UserId { + fn from(value: u64) -> Self { + Self(id::UserId::new(value)) + } + } + impl FromStr for UserId { type Err = ParseError; diff --git a/youmubot/src/main.rs b/youmubot/src/main.rs index 9404011..ca2d25a 100644 --- a/youmubot/src/main.rs +++ b/youmubot/src/main.rs @@ -53,6 +53,8 @@ struct Env { prelude: youmubot_prelude::Env, #[cfg(feature = "osu")] osu: youmubot_osu::discord::OsuEnv, + #[cfg(feature = "core")] + core: youmubot_core::CoreEnv, } impl AsRef for Env { @@ -68,6 +70,13 @@ impl AsRef for Env { } } +#[cfg(feature = "core")] +impl AsRef for Env { + fn as_ref(&self) -> &youmubot_core::CoreEnv { + &self.core + } +} + impl TypeMapKey for Env { type Value = Env; } @@ -219,7 +228,9 @@ async fn main() { let prelude = setup::setup_prelude(&db_path, sql_path, &mut data).await; // Setup core #[cfg(feature = "core")] - youmubot_core::setup(&db_path, &mut data).expect("Setup db should succeed"); + let core = youmubot_core::setup(&db_path, &mut data, prelude.clone()) + .await + .expect("Setup db should succeed"); // osu! #[cfg(feature = "osu")] let osu = youmubot_osu::discord::setup(&mut data, prelude.clone(), &mut announcers) @@ -233,6 +244,8 @@ async fn main() { prelude, #[cfg(feature = "osu")] osu, + #[cfg(feature = "core")] + core, } }; @@ -266,20 +279,49 @@ async fn main() { case_insensitive_commands: true, ..Default::default() }, + command_check: Some(|ctx| { + Box::pin(command_check(ctx.data(), UserId(ctx.author().id)).map(Ok)) + }), on_error: |err| { Box::pin(async move { - if let poise::FrameworkError::Command { error, ctx, .. } = err { - let reply = format!( - "Command '{}' returned error: {:?}", - ctx.invoked_command_name(), - error - ); - eprintln!("{}\n{:?}", reply, error); - ctx.send(poise::CreateReply::default().content(reply).ephemeral(true)) + match err { + poise::FrameworkError::Command { error, ctx, .. } => { + let reply = format!( + "Command '{}' returned error: {:?}", + ctx.invoked_command_name(), + error + ); + eprintln!("{}\n{:?}", reply, error); + ctx.send(poise::CreateReply::default().content(reply).ephemeral(true)) + .await + .pls_ok(); + } + poise::FrameworkError::NotAnOwner { ctx, .. } => { + ctx.send( + poise::CreateReply::default() + .content("You have to be an owner to run this command!") + .ephemeral(true), + ) .await .pls_ok(); - } else { - eprintln!("Poise error: {:?}", err) + } + poise::FrameworkError::CommandCheckFailed { error: _, ctx, .. } + | poise::FrameworkError::CooldownHit { + remaining_cooldown: _, + ctx, + .. + } => { + ctx.send( + poise::CreateReply::default() + .content("You are being rate-limited, please try again later!") + .ephemeral(true), + ) + .await + .pls_ok(); + } + _ => { + eprintln!("Poise error: {:?}", err) + } } }) }, @@ -287,6 +329,8 @@ async fn main() { poise_register(), #[cfg(feature = "osu")] youmubot_osu::discord::osu_command(), + #[cfg(feature = "core")] + youmubot_core::admin::ignore::ignore(), ], ..Default::default() }) @@ -398,25 +442,37 @@ async fn poise_register(ctx: CmdContext<'_, Env>) -> Result<()> { Ok(()) } +async fn command_check(env: &Env, author: UserId) -> bool { + #[cfg(feature = "core")] + if youmubot_core::admin::ignore::should_ignore(env, author) { + tracing::info!("User is in ignore list, skipping..."); + return false; + } + true +} + // Hooks! #[hook] -async fn before_hook(_: &Context, msg: &Message, command_name: &str) -> bool { - println!( +async fn before_hook(ctx: &Context, msg: &Message, command_name: &str) -> bool { + let env = ctx.data.read().await; + let env = env.get::().unwrap(); + tracing::info!( "Got command '{}' by user '{}'", - command_name, msg.author.name + command_name, + msg.author.name ); - true + command_check(env, UserId(msg.author.id)).await } #[hook] async fn after_hook(ctx: &Context, msg: &Message, command_name: &str, error: CommandResult) { match error { - Ok(()) => println!("Processed command '{}'", command_name), + Ok(()) => tracing::info!("Processed command '{}'", command_name), Err(why) => { let reply = format!("Command '{}' returned error {:?}", command_name, why); msg.reply(&ctx, &reply).await.ok(); - println!("{}", reply) + tracing::info!("{}", reply) } } } @@ -441,6 +497,9 @@ async fn on_dispatch_error(ctx: &Context, msg: &Message, error: DispatchError, _ max, given ), DispatchError::OnlyForGuilds => "🔇 This command cannot be used in DMs.".to_owned(), + DispatchError::OnlyForOwners => { + "🔇 This command can only be used by bot owners.".to_owned() + } _ => return, }, )