diff --git a/.vs/settings.json b/.vs/settings.json index d210276..b7fe30b 100644 --- a/.vs/settings.json +++ b/.vs/settings.json @@ -1,5 +1,7 @@ { - "rust.features": [ + "rust-analyzer.cargo.features": [ "music" - ] + ], + "rust-analyzer.cargo.loadOutDirsFromCheck": true, + "rust-analyzer.procMacro.enable": true } \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index ee31f01..700d5f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,16 +10,17 @@ edition = "2018" music = ["serenity/voice"] [dependencies] -serenity = { version = "0.9.0-rc.2" } -toml = "0.5.6" +serenity = "0.9" +toml = "0.5" serde = { version = "1.0", features = ["derive"] } -reqwest = "0.10.7" -rand = "0.7.3" -lazy_static = "1.4.0" -async-trait = "0.1.36" +reqwest = "0.10" +rand = "0.7" +lazy_static = "1.4" +async-trait = "0.1" tokio = { version = "0.2", features = ["full"] } futures = "0.3" -chrono = "0.4.15" +chrono = "0.4" serde_json = "1.0" log = "0.4" -log4rs = "0.13.0" \ No newline at end of file +log4rs = "0.13" +ctrlc = "3.1" \ No newline at end of file diff --git a/log4rs.yaml b/log4rs.yaml index 71dda45..794b89c 100644 --- a/log4rs.yaml +++ b/log4rs.yaml @@ -5,7 +5,10 @@ appenders: kind: file path: "log/debug.log" root: - level: debug + level: warn appenders: - stdout - - file-debug \ No newline at end of file + - file-debug +loggers: + rusty_bot: + level: trace \ No newline at end of file diff --git a/src/commands/general.rs b/src/commands/general.rs index 33ea7df..a8a2f68 100644 --- a/src/commands/general.rs +++ b/src/commands/general.rs @@ -1,4 +1,4 @@ -use crate::{api, debugln, ShardManagerContainer}; +use crate::{api, data::ShardManagerContainer, debugln}; use futures::StreamExt; use log::error; use serenity::{ @@ -249,6 +249,14 @@ async fn ping(ctx: &Context, msg: &Message) -> CommandResult { #[description = "Image"] #[bucket = "image"] pub async fn image(ctx: &Context, msg: &Message, args: Args) -> CommandResult { + msg.author + .has_role( + ctx, + msg.guild_id.unwrap(), + msg.guild(ctx).await.unwrap().role_by_name("mute").unwrap(), + ) + .await + .unwrap(); if let Err(e) = _image(ctx, msg, args).await { error!("Error in image : {:?}", e); } @@ -263,8 +271,7 @@ async fn _image(ctx: &Context, msg: &Message, mut args: Args) -> crate::commands m.embed(|e| { image.embed(e); e - }); - m + }) }) .await?; } diff --git a/src/commands/music.rs b/src/commands/music.rs index 04f58c5..e20e498 100644 --- a/src/commands/music.rs +++ b/src/commands/music.rs @@ -1,3 +1,4 @@ +use crate::data::{GuildOptions, GuildOptionsKey}; use log::{error, info}; use serenity::{ client::{bridge::voice::ClientVoiceManager, Context}, @@ -5,7 +6,7 @@ use serenity::{ macros::{command, group}, Args, CommandResult, }, - model::{channel::Message, misc::Mentionable}, + model::{channel::Message, guild::PartialMember, id::GuildId, misc::Mentionable}, prelude::*, voice, Result as SerenityResult, }; @@ -33,6 +34,19 @@ async fn join(ctx: &Context, msg: &Message) -> CommandResult { } }; + if if let Some(member) = &msg.member { + is_mute(ctx, member, guild.id).await.unwrap_or(false) + } else { + false + } { + check_msg( + msg.channel_id + .say(&ctx.http, "Error, you cant play music") + .await, + ); + return Ok(()); + } + let guild_id = guild.id; let channel_id = guild @@ -91,6 +105,19 @@ async fn leave(ctx: &Context, msg: &Message) -> CommandResult { } }; + if if let Some(member) = &msg.member { + is_mute(ctx, member, guild_id).await.unwrap_or(false) + } else { + false + } { + check_msg( + msg.channel_id + .say(&ctx.http, "Error, you cant play music") + .await, + ); + return Ok(()); + } + let manager_lock = ctx .data .read() @@ -151,6 +178,19 @@ async fn play(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { } }; + if if let Some(member) = &msg.member { + is_mute(ctx, member, guild_id).await.unwrap_or(false) + } else { + false + } { + check_msg( + msg.channel_id + .say(&ctx.http, "Error, you cant play music") + .await, + ); + return Ok(()); + } + let manager_lock = ctx .data .read() @@ -202,6 +242,19 @@ async fn stop(ctx: &Context, msg: &Message) -> CommandResult { } }; + if if let Some(member) = &msg.member { + is_mute(ctx, member, guild_id).await.unwrap_or(false) + } else { + false + } { + check_msg( + msg.channel_id + .say(&ctx.http, "Error, you cant play music") + .await, + ); + return Ok(()); + } + let manager_lock = ctx .data .read() @@ -232,3 +285,22 @@ fn check_msg(result: SerenityResult) { error!("Error sending message: {:?}", why); } } + +async fn is_mute( + ctx: &Context, + member: &PartialMember, + guild_id: GuildId, +) -> tokio::io::Result { + let mut data = ctx.data.write().await; + let data = data + .get_mut::() + .expect("Failed to get guild cache"); + + let guild_options = data + .entry(guild_id) + .or_insert_with(|| GuildOptions::default().set_guild_id(guild_id)); + + Ok(member + .roles + .contains(&guild_options.get_mute_role(ctx).await?)) +} diff --git a/src/commands/roulette.rs b/src/commands/roulette.rs index 06c0efa..e7351fd 100644 --- a/src/commands/roulette.rs +++ b/src/commands/roulette.rs @@ -1,4 +1,8 @@ -use crate::{api, commands, debugln}; +use crate::{ + api, commands, + data::{BulletsContainer, GuildOptions, GuildOptionsKey}, + debugln, +}; use log::error; use rand::Rng; use serenity::{ @@ -9,19 +13,6 @@ use serenity::{ model::prelude::*, prelude::*, }; -use std::collections::{HashMap, HashSet}; - -pub(crate) struct BulletsContainer; - -impl TypeMapKey for BulletsContainer { - type Value = HashMap; -} - -pub(crate) struct NonKickGuildsContainer; - -impl TypeMapKey for NonKickGuildsContainer { - type Value = HashSet; -} #[group] #[default_command(shot)] @@ -70,7 +61,7 @@ async fn _shot(ctx: &Context, msg: &Message) -> commands::Result<()> { ) .await?; } - debugln!("Bullets Map : {:?}", bullets_map); + log::trace!("Bullets Map : {:?}", bullets_map); Ok(()) } @@ -141,10 +132,12 @@ async fn kick(ctx: &Context, msg: &Message) -> CommandResult { async fn _kick(ctx: &Context, msg: &Message) -> commands::Result<()> { if let Some(guild_id) = &msg.guild_id { let mut data = ctx.data.write().await; - let non_kick_guilds = data - .get_mut::() + let guilds_options = data + .get_mut::() .expect("Expected NonKickGuildsContainer in TypeMap."); - if non_kick_guilds.contains(guild_id.as_u64()) { + + let guild_options = guilds_options.entry(*guild_id).or_default(); + if !guild_options.roulette_options.kick_enabled { msg.channel_id .say( ctx, @@ -193,19 +186,23 @@ async fn disable_kick(ctx: &Context, msg: &Message, mut args: Args) -> CommandRe _ => args.single::()?, }; let mut data = ctx.data.write().await; - let non_kick_guilds = data - .get_mut::() + let guilds_options = data + .get_mut::() .expect("Expected NonKickGuildsContainer in TypeMap."); if let Some(guild_id) = msg.guild_id { - let id = *guild_id.as_u64(); + let entry = guilds_options + .entry(guild_id) + .or_insert_with(|| GuildOptions::default().set_guild_id(guild_id)); + entry.roulette_options.kick_enabled = !disable; if disable { - non_kick_guilds.insert(id); msg.channel_id.say(ctx, "No fun allowed").await?; } else { - non_kick_guilds.remove(&id); msg.channel_id.say(ctx, "Done").await?; } + + entry.save_async(guild_id.0).await?; } + log::debug!("{:?}", guilds_options); Ok(()) } diff --git a/src/data/guilds_options.rs b/src/data/guilds_options.rs new file mode 100644 index 0000000..7c86e88 --- /dev/null +++ b/src/data/guilds_options.rs @@ -0,0 +1,159 @@ +use log::error; +use serde::{Deserialize, Serialize}; +use serenity::{ + model::prelude::{GuildId, RoleId}, + prelude::TypeMapKey, +}; +use std::{ + collections::HashMap, + fs, + io::{Result as IoResult, Write}, + path::{Path, PathBuf}, +}; +use tokio::io::AsyncWriteExt; + +#[cfg(feature = "music")] +use serenity::{http::CacheHttp, model::prelude::PartialGuild}; + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct GuildOptions { + #[serde(skip_serializing)] + pub(crate) guild_id: Option, + #[serde(skip_serializing)] + pub(crate) mute_id: Option, + pub(crate) roulette_options: RouletteOptions, +} + +impl GuildOptions { + pub async fn save_async(&mut self, guild_id: u64) -> tokio::io::Result<()> { + let path = PathBuf::from(format!("./data/guilds_options/{}.json", guild_id)); + + if !path.parent().unwrap().exists() { + tokio::fs::create_dir_all(path.parent().unwrap()).await?; + } + + let mut file = tokio::fs::File::create(path).await?; + let serialized = serde_json::to_string_pretty(self)?; + file.write_all(&serialized.as_bytes()).await?; + + Ok(()) + } + + pub fn save(&mut self, guild_id: u64) -> IoResult<()> { + let path = PathBuf::from(format!("./data/guilds_options/{}.json", guild_id)); + + if !path.parent().unwrap().exists() { + std::fs::create_dir_all(path.parent().unwrap())?; + } + + let mut file = std::fs::File::create(path)?; + let serialized = serde_json::to_string_pretty(self)?; + file.write_all(&serialized.as_bytes())?; + + Ok(()) + } + + pub(crate) fn set_guild_id(mut self, id: GuildId) -> Self { + self.guild_id = Some(id); + self + } + + pub fn load_from_dir>(path: P) -> IoResult> { + let mut res = HashMap::new(); + for entry in fs::read_dir(path)?.filter_map(|e| e.ok()) { + match serde_json::from_reader::(fs::File::open(entry.path())?) { + Ok(options) => { + let id = GuildId::from( + entry + .file_name() + .to_string_lossy() + .split('.') + .next() + .unwrap() + .parse::() + .unwrap(), + ); + res.insert(id, options.set_guild_id(id)); + } + Err(e) => { + error!("While parsing guild option {}", e); + } + } + } + Ok(res) + } + + #[cfg(feature = "music")] + pub(crate) async fn get_mute_role( + &mut self, + cache: &C, + ) -> tokio::io::Result { + match self.mute_id { + Some(mute_id) => Ok(mute_id), + + None => { + if let Some(guild_id) = self.guild_id { + let guild: PartialGuild = cache + .http() + .get_guild(guild_id.0) + .await + .map_err(|e| tokio::io::Error::new(tokio::io::ErrorKind::Other, e))?; + + match guild.role_by_name("mute") { + Some(role) => { + self.mute_id = Some(role.id); + Ok(role.id) + } + None => Err(tokio::io::Error::new( + tokio::io::ErrorKind::Other, + "Unkown role", + )), + } + } else { + Err(tokio::io::Error::new( + tokio::io::ErrorKind::Other, + "Unkown guild id", + )) + } + } + } + } +} + +impl Default for GuildOptions { + fn default() -> Self { + Self { + roulette_options: RouletteOptions::default(), + guild_id: None, + mute_id: None, + } + } +} + +impl Drop for GuildOptions { + fn drop(&mut self) { + if let Some(id) = self.guild_id { + log::debug!("Saving {:?}", self); + if let Err(e) = self.save(id.0) { + log::error!("While saving {} : {}", id.0, e); + } + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct RouletteOptions { + pub(crate) kick_enabled: bool, +} + +impl Default for RouletteOptions { + fn default() -> Self { + Self { kick_enabled: true } + } +} + +pub(crate) struct GuildOptionsKey; + +impl TypeMapKey for GuildOptionsKey { + type Value = HashMap; +} diff --git a/src/data/mod.rs b/src/data/mod.rs new file mode 100644 index 0000000..94f83f6 --- /dev/null +++ b/src/data/mod.rs @@ -0,0 +1,20 @@ +use serenity::{ + client::bridge::gateway::ShardManager, + prelude::{Mutex as SerenityMutex, TypeMapKey}, +}; +use std::{collections::HashMap, sync::Arc}; + +mod guilds_options; +pub(crate) use guilds_options::{GuildOptions, GuildOptionsKey}; + +pub(crate) struct ShardManagerContainer; + +impl TypeMapKey for ShardManagerContainer { + type Value = Arc>; +} + +pub(crate) struct BulletsContainer; + +impl TypeMapKey for BulletsContainer { + type Value = HashMap; +} diff --git a/src/main.rs b/src/main.rs index 709a4f3..11af12c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,14 +1,13 @@ -use crate::commands::{ - admin::ADMIN_GROUP, - general::GENERAL_GROUP, - owner::OWNER_GROUP, - roulette::{BulletsContainer, NonKickGuildsContainer, ROULETTE_GROUP}, +use crate::{ + commands::{ + admin::ADMIN_GROUP, general::GENERAL_GROUP, owner::OWNER_GROUP, roulette::ROULETTE_GROUP, + }, + data::{BulletsContainer, GuildOptions, GuildOptionsKey, ShardManagerContainer}, }; use async_trait::async_trait; use log::{debug, error, info}; use serde_json::Value; use serenity::{ - client::bridge::gateway::ShardManager, framework::standard::{ help_commands, macros::{help, hook}, @@ -35,6 +34,7 @@ use crate::commands::music::{VoiceManager, MUSIC_GROUP}; mod api; mod commands; mod config; +mod data; mod macros; mod presence; @@ -42,13 +42,6 @@ const PREFIX: &str = "?"; static mut LOG_ATTACHMENTS: bool = false; pub(crate) static mut INVITE_URL: Option = None; -struct ShardManagerContainer; - -// TODO SAVE ON DROP -impl TypeMapKey for ShardManagerContainer { - type Value = Arc>; -} - // TODO CLAP FOR CLI #[tokio::main] async fn main() -> IoResult<()> { @@ -135,7 +128,7 @@ async fn main() -> IoResult<()> { .unrecognised_command(unknown_command) } - let mut client = Client::new(&token) + let mut client = Client::builder(&token) .event_handler(Messages {}) .framework(framework) .await @@ -150,15 +143,29 @@ async fn main() -> IoResult<()> { data.insert::(std::sync::Arc::clone(&client.voice_manager)); } - let non_kick_guilds = data_dir.join("nonkickguilds.json"); - data.insert::(if non_kick_guilds.exists() { - serde_json::from_reader(fs::File::open(non_kick_guilds)?)? - } else { - HashSet::new() + data.insert::({ + let options = GuildOptions::load_from_dir("./data/guilds_options").unwrap_or_default(); + log::debug!("Loaded {:?}", options); + options }) } - client.start_autosharded().await.unwrap(); + let current_runtime = tokio::runtime::Handle::current(); + let shard_manager = Arc::clone(&client.shard_manager); + ctrlc::set_handler(move || { + let shard_manager = Arc::clone(&shard_manager); + current_runtime.spawn(async move { + let mut c = shard_manager.lock().await; + c.shutdown_all().await; + }); + }) + .unwrap(); + + if let Err(e) = client.start_autosharded().await { + log::error!("Error while running bot : {}", e); + } + + log::info!("Stopping bot"); Ok(()) } @@ -212,32 +219,13 @@ impl EventHandler for Messages { .await; let ctx_clone = ctx.clone(); - tokio::spawn(async move { - let delay = Duration::from_secs(5); - let mut presence_generator = presence::Presences::new(); - while let Some(act) = presence_generator.next(&ctx_clone).await { - ctx_clone - .set_presence(Some(act), OnlineStatus::Online) - .await; - tokio::time::delay_for(delay).await; - } - }); - - tokio::signal::ctrl_c().await.unwrap(); - debugln!("ctrl-c"); - let data = ctx.data.read().await; - - info!("Saving data ..."); - if let Err(e) = save_data(&data).await { - error!("Error while saving data : {:?}", e); - } - info!("Data saved"); - - if let Some(manager) = data.get::() { - manager.lock().await.shutdown_all().await; - info!("Stopped"); - } else { - error!("There was a problem getting the shard manager"); + let delay = Duration::from_secs(5); + let mut presence_generator = presence::Presences::new(); + while let Some(act) = presence_generator.next(&ctx_clone).await { + ctx_clone + .set_presence(Some(act), OnlineStatus::Online) + .await; + tokio::time::delay_for(delay).await; } } @@ -256,21 +244,6 @@ impl EventHandler for Messages { } } -async fn save_data(data: &tokio::sync::RwLockReadGuard<'_, TypeMap>) -> commands::Result<()> { - let data_path = Path::new("data"); - - if let Some(data) = data.get::() { - let mut f = File::create(data_path.join("nonkickguilds.json")).await?; - let json = if cfg!(debug_assertions) { - serde_json::to_string_pretty(data)? - } else { - serde_json::to_string(data)? - }; - f.write_all(&json.as_bytes()).await?; - } - Ok(()) -} - async fn download_to_log(attachment: Attachment) -> commands::Result<()> { debugln!("Download_to_log : {:?}", attachment); let path = Path::new("logging").join(format!("{}-{}", attachment.id, attachment.filename));