Test

.vs/settings.json

@@ -1,5 +1,7 @@
 {
-    "rust.features": [
+    "rust-analyzer.cargo.features": [
         "music"
-    ]
+    ],
+    "rust-analyzer.cargo.loadOutDirsFromCheck": true,
+    "rust-analyzer.procMacro.enable": true
 }

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"
+log4rs = "0.13"
+ctrlc = "3.1"

log4rs.yaml

@@ -5,7 +5,10 @@ appenders:
     kind: file
     path: "log/debug.log"
 root:
-  level: debug
+  level: warn
   appenders: 
     - stdout
     - file-debug
+loggers:
+  rusty_bot:
+    level: trace

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?;
         }

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<Message>) {
         error!("Error sending message: {:?}", why);
     }
 }
+
+async fn is_mute(
+    ctx: &Context,
+    member: &PartialMember,
+    guild_id: GuildId,
+) -> tokio::io::Result<bool> {
+    let mut data = ctx.data.write().await;
+    let data = data
+        .get_mut::<GuildOptionsKey>()
+        .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?))
+}

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<u64, (u8, u8)>;
-}
-
-pub(crate) struct NonKickGuildsContainer;
-
-impl TypeMapKey for NonKickGuildsContainer {
-    type Value = HashSet<u64>;
-}
 
 #[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::<NonKickGuildsContainer>()
+        let guilds_options = data
+            .get_mut::<GuildOptionsKey>()
             .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::<bool>()?,
     };
     let mut data = ctx.data.write().await;
-    let non_kick_guilds = data
-        .get_mut::<NonKickGuildsContainer>()
+    let guilds_options = data
+        .get_mut::<GuildOptionsKey>()
         .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(())
 }

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<GuildId>,
+    #[serde(skip_serializing)]
+    pub(crate) mute_id: Option<RoleId>,
+    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<P: AsRef<Path>>(path: P) -> IoResult<HashMap<GuildId, GuildOptions>> {
+        let mut res = HashMap::new();
+        for entry in fs::read_dir(path)?.filter_map(|e| e.ok()) {
+            match serde_json::from_reader::<fs::File, GuildOptions>(fs::File::open(entry.path())?) {
+                Ok(options) => {
+                    let id = GuildId::from(
+                        entry
+                            .file_name()
+                            .to_string_lossy()
+                            .split('.')
+                            .next()
+                            .unwrap()
+                            .parse::<u64>()
+                            .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<C: CacheHttp>(
+        &mut self,
+        cache: &C,
+    ) -> tokio::io::Result<RoleId> {
+        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<GuildId, GuildOptions>;
+}

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<SerenityMutex<ShardManager>>;
+}
+
+pub(crate) struct BulletsContainer;
+
+impl TypeMapKey for BulletsContainer {
+    type Value = HashMap<u64, (u8, u8)>;
+}

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<String> = None;
 
-struct ShardManagerContainer;
-
-// TODO SAVE ON DROP
-impl TypeMapKey for ShardManagerContainer {
-    type Value = Arc<Mutex<ShardManager>>;
-}
-
 // 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::<VoiceManager>(std::sync::Arc::clone(&client.voice_manager));
         }
 
-        let non_kick_guilds = data_dir.join("nonkickguilds.json");
-        data.insert::<NonKickGuildsContainer>(if non_kick_guilds.exists() {
-            serde_json::from_reader(fs::File::open(non_kick_guilds)?)?
-        } else {
-            HashSet::new()
+        data.insert::<GuildOptionsKey>({
+            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::<ShardManagerContainer>() {
-            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::<NonKickGuildsContainer>() {
-        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));