From 7d49edb1fe447ebf83761d116aae6d7dc89d83ff Mon Sep 17 00:00:00 2001 From: natsukagami Date: Sat, 18 Jan 2020 23:08:38 +0000 Subject: [PATCH] Announcer Trait (#9) --- youmubot/src/commands/announcer.rs | 84 ++++++++++++++++++++++++++++++ youmubot/src/commands/mod.rs | 3 ++ youmubot/src/db/mod.rs | 6 ++- 3 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 youmubot/src/commands/announcer.rs diff --git a/youmubot/src/commands/announcer.rs b/youmubot/src/commands/announcer.rs new file mode 100644 index 0000000..dd25ecd --- /dev/null +++ b/youmubot/src/commands/announcer.rs @@ -0,0 +1,84 @@ +use crate::db::{AnnouncerChannels, DBWriteGuard}; +use serenity::{ + framework::standard::{CommandError as Error, CommandResult}, + http::{CacheHttp, Http}, + model::id::{ChannelId, GuildId, UserId}, + prelude::ShareMap, +}; +use std::{ + collections::HashSet, + thread::{spawn, JoinHandle}, +}; + +pub trait Announcer { + fn announcer_key() -> &'static str; + fn send_messages( + c: &Http, + d: &mut ShareMap, + channels: impl Fn(UserId) -> Vec + Sync, + ) -> CommandResult; + + fn set_channel(d: &mut ShareMap, guild: GuildId, channel: ChannelId) -> CommandResult { + let mut data: DBWriteGuard<_> = d + .get_mut::() + .expect("DB initialized") + .into(); + let mut data = data.borrow_mut()?; + data.entry(Self::announcer_key().to_owned()) + .or_default() + .insert(guild, channel); + Ok(()) + } + + fn get_guilds(d: &mut ShareMap) -> Result, Error> { + let data = d + .get::() + .expect("DB initialized") + .read(|v| { + v.get(Self::announcer_key()) + .map(|m| m.iter().map(|(a, b)| (*a, *b)).collect()) + .unwrap_or_else(|| vec![]) + })?; + Ok(data) + } + + fn announce(c: &Http, d: &mut ShareMap) -> CommandResult { + let guilds: Vec<_> = Self::get_guilds(d)?; + let member_sets = { + let mut v = Vec::with_capacity(guilds.len()); + for (guild, channel) in guilds.into_iter() { + let mut s = HashSet::new(); + for user in guild + .members_iter(c.as_ref()) + .take_while(|u| u.is_ok()) + .filter_map(|u| u.ok()) + { + s.insert(user.user_id()); + } + v.push((s, channel)) + } + v + }; + Self::send_messages(c.as_ref(), d, |user_id| { + let mut v = Vec::new(); + for (members, channel) in member_sets.iter() { + if members.contains(&user_id) { + v.push(*channel); + } + } + v + })?; + Ok(()) + } + + fn scan(client: &serenity::Client, cooldown: std::time::Duration) -> JoinHandle<()> { + let c = client.cache_and_http.clone(); + let data = client.data.clone(); + spawn(move || loop { + if let Err(e) = Self::announce(c.http(), &mut *data.write()) { + dbg!(e); + } + std::thread::sleep(cooldown); + }) + } +} diff --git a/youmubot/src/commands/mod.rs b/youmubot/src/commands/mod.rs index b323cbb..d0eb82a 100644 --- a/youmubot/src/commands/mod.rs +++ b/youmubot/src/commands/mod.rs @@ -7,6 +7,7 @@ use serenity::{ }; use std::collections::HashSet; +mod announcer; mod args; pub mod admin; @@ -17,6 +18,8 @@ pub use admin::ADMIN_GROUP; pub use community::COMMUNITY_GROUP; pub use fun::FUN_GROUP; +pub use announcer::Announcer; + // A help command #[help] pub fn help( diff --git a/youmubot/src/db/mod.rs b/youmubot/src/db/mod.rs index 34a2fbe..1b8b055 100644 --- a/youmubot/src/db/mod.rs +++ b/youmubot/src/db/mod.rs @@ -5,7 +5,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use serenity::{ client::Client, framework::standard::CommandError as Error, - model::id::{GuildId, RoleId, UserId}, + model::id::{ChannelId, GuildId, RoleId, UserId}, prelude::*, }; use std::collections::HashMap; @@ -32,6 +32,9 @@ where } } +/// A map from announcer keys to guild IDs and to channels. +pub type AnnouncerChannels = DB>>; + /// A list of SoftBans for all servers. pub type SoftBans = DB>; @@ -43,6 +46,7 @@ pub fn setup_db(client: &mut Client) -> Result<(), Error> { }); let mut data = client.data.write(); SoftBans::insert_into(&mut *data, &path.join("soft_bans.ron"))?; + AnnouncerChannels::insert_into(&mut *data, &path.join("announcers.yaml"))?; Ok(()) }