From 384b7be52cbd0469ce1ff12483df4fe77a5ea064 Mon Sep 17 00:00:00 2001 From: natsukagami Date: Fri, 8 Nov 2019 14:31:53 +0000 Subject: [PATCH] Admin commands done (#1) --- .gitignore | 1 + Cargo.lock | 41 ++++ youmubot/Cargo.toml | 6 + .../src/commands/{admin.rs => admin/mod.rs} | 6 +- youmubot/src/commands/admin/soft_ban.rs | 176 ++++++++++++++++++ youmubot/src/commands/args.rs | 117 ++++++++++++ youmubot/src/commands/mod.rs | 3 +- youmubot/src/db/mod.rs | 111 +++++++++++ youmubot/src/main.rs | 7 + 9 files changed, 466 insertions(+), 2 deletions(-) rename youmubot/src/commands/{admin.rs => admin/mod.rs} (93%) create mode 100644 youmubot/src/commands/admin/soft_ban.rs create mode 100644 youmubot/src/commands/args.rs create mode 100644 youmubot/src/db/mod.rs diff --git a/.gitignore b/.gitignore index c5dd462..3ca3667 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ target .env +*.ron diff --git a/Cargo.lock b/Cargo.lock index 68127bf..f36d014 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -518,6 +518,11 @@ name = "libc" version = "0.2.65" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "linked-hash-map" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "lock_api" version = "0.3.1" @@ -892,6 +897,16 @@ dependencies = [ "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "rustbreak" +version = "2.0.0-rc3" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "failure 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)", + "serde_yaml 0.7.5 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "rustc-demangle" version = "0.1.16" @@ -988,6 +1003,17 @@ dependencies = [ "url 1.7.2 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "serde_yaml" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "dtoa 0.4.4 (registry+https://github.com/rust-lang/crates.io-index)", + "linked-hash-map 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)", + "yaml-rust 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "serenity" version = "0.7.2" @@ -1531,11 +1557,22 @@ dependencies = [ "winapi-build 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "yaml-rust" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "linked-hash-map 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "youmubot" version = "0.1.0" dependencies = [ + "chrono 0.4.9 (registry+https://github.com/rust-lang/crates.io-index)", "dotenv 0.15.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rustbreak 2.0.0-rc3 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)", "serenity 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -1608,6 +1645,7 @@ dependencies = [ "checksum kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "7507624b29483431c0ba2d82aece8ca6cdba9382bff4ddd0f7490560c056098d" "checksum lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" "checksum libc 0.2.65 (registry+https://github.com/rust-lang/crates.io-index)" = "1a31a0627fdf1f6a39ec0dd577e101440b7db22672c0901fe00a9a6fbb5c24e8" +"checksum linked-hash-map 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "ae91b68aebc4ddb91978b11a1b02ddd8602a05ec19002801c5666000e05e0f83" "checksum lock_api 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)" = "f8912e782533a93a167888781b836336a6ca5da6175c05944c86cf28c31104dc" "checksum log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)" = "14b6052be84e6b71ab17edffc2eeabf5c2c3ae1fdb464aae35ac50c67a44e1f7" "checksum matches 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" @@ -1648,6 +1686,7 @@ dependencies = [ "checksum regex-syntax 0.6.12 (registry+https://github.com/rust-lang/crates.io-index)" = "11a7e20d1cce64ef2fed88b66d347f88bd9babb82845b2b858f3edbf59a4f716" "checksum reqwest 0.9.22 (registry+https://github.com/rust-lang/crates.io-index)" = "2c2064233e442ce85c77231ebd67d9eca395207dec2127fe0bbedde4bd29a650" "checksum ring 0.16.9 (registry+https://github.com/rust-lang/crates.io-index)" = "6747f8da1f2b1fabbee1aaa4eb8a11abf9adef0bf58a41cee45db5d59cecdfac" +"checksum rustbreak 2.0.0-rc3 (registry+https://github.com/rust-lang/crates.io-index)" = "b1c185a2ede13fcb28feb6864ee9412a20f57bd83b4be18dc81fde4d6e786982" "checksum rustc-demangle 0.1.16 (registry+https://github.com/rust-lang/crates.io-index)" = "4c691c0e608126e00913e33f0ccf3727d5fc84573623b8d65b2df340b5201783" "checksum rustc_version 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "138e3e0acb6c9fb258b19b67cb8abd63c00679d2851805ea151465464fe9030a" "checksum rustls 0.16.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b25a18b1bf7387f0145e7f8324e700805aade3842dd3db2e74e4cdeb4677c09e" @@ -1660,6 +1699,7 @@ dependencies = [ "checksum serde_derive 1.0.101 (registry+https://github.com/rust-lang/crates.io-index)" = "4b133a43a1ecd55d4086bd5b4dc6c1751c68b1bfbeba7a5040442022c7e7c02e" "checksum serde_json 1.0.41 (registry+https://github.com/rust-lang/crates.io-index)" = "2f72eb2a68a7dc3f9a691bfda9305a1c017a6215e5a4545c258500d2099a37c2" "checksum serde_urlencoded 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)" = "642dd69105886af2efd227f75a520ec9b44a820d65bc133a9131f7d229fd165a" +"checksum serde_yaml 0.7.5 (registry+https://github.com/rust-lang/crates.io-index)" = "ef8099d3df28273c99a1728190c7a9f19d444c941044f64adf986bee7ec53051" "checksum serenity 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)" = "644b296c0c732d33d1258f74a5862e25fa8b91aff16df90e2c98c594d0f019f5" "checksum sha-1 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "23962131a91661d643c98940b20fcaffe62d776a823247be80a48fcb8b6fce68" "checksum slab 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c111b5bd5695e56cffe5129854aa230b39c93a305372fdbb2668ca2394eea9f8" @@ -1720,3 +1760,4 @@ dependencies = [ "checksum winapi-x86_64-pc-windows-gnu 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" "checksum winreg 0.6.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b2986deb581c4fe11b621998a5e53361efe6b48a151178d0cd9eeffa4dc6acc9" "checksum ws2_32-sys 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "d59cefebd0c892fa2dd6de581e937301d8552cb44489cdff035c6187cb63fa5e" +"checksum yaml-rust 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "65923dd1784f44da1d2c3dbbc5e822045628c590ba72123e1c73d3c230c4434d" diff --git a/youmubot/Cargo.toml b/youmubot/Cargo.toml index d087142..06b8644 100644 --- a/youmubot/Cargo.toml +++ b/youmubot/Cargo.toml @@ -9,3 +9,9 @@ edition = "2018" [dependencies] serenity = "0.7" dotenv = "0.15" +serde = { version = "1.0", features = ["derive"] } +chrono = "0.4.9" + +[dependencies.rustbreak] +version = "2.0.0-rc3" +features = ["yaml_enc"] diff --git a/youmubot/src/commands/admin.rs b/youmubot/src/commands/admin/mod.rs similarity index 93% rename from youmubot/src/commands/admin.rs rename to youmubot/src/commands/admin/mod.rs index 9131252..3cc9adf 100644 --- a/youmubot/src/commands/admin.rs +++ b/youmubot/src/commands/admin/mod.rs @@ -6,8 +6,12 @@ use serenity::{ }, model::{channel::Message, id::UserId}, }; +use soft_ban::{SOFT_BAN_COMMAND, SOFT_BAN_INIT_COMMAND}; use std::{thread::sleep, time::Duration}; +mod soft_ban; +pub use soft_ban::watch_soft_bans; + group!({ name: "admin", options: { @@ -15,7 +19,7 @@ group!({ prefixes: ["admin", "a"], description: "Administrative commands for the server.", }, - commands: [clean, ban, kick], + commands: [clean, ban, kick, soft_ban, soft_ban_init], }); #[command] diff --git a/youmubot/src/commands/admin/soft_ban.rs b/youmubot/src/commands/admin/soft_ban.rs new file mode 100644 index 0000000..d997dea --- /dev/null +++ b/youmubot/src/commands/admin/soft_ban.rs @@ -0,0 +1,176 @@ +use crate::{ + commands::args, + db::{DBWriteGuard, ServerSoftBans, SoftBans}, +}; +use chrono::offset::Utc; +use serenity::prelude::*; +use serenity::{ + framework::standard::{macros::command, Args, CommandError as Error, CommandResult}, + model::{ + channel::Message, + id::{RoleId, UserId}, + }, +}; +use std::cmp::max; + +#[command] +#[required_permissions(ADMINISTRATOR)] +#[description = "Soft-ban an user, might be with a certain amount of time. Re-banning an user removes the ban itself."] +#[usage = "user#1234 [time]"] +#[example = "user#1234 5s"] +#[min_args(1)] +#[max_args(2)] +pub fn soft_ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult { + let user = args.single::()?.to_user(&ctx)?; + let duration = if args.is_empty() { + None + } else { + Some( + args.single::() + .map_err(|e| Error::from(&format!("{:?}", e)))?, + ) + }; + let guild = msg.guild_id.ok_or(Error::from("Command is guild only"))?; + + let mut data = ctx.data.write(); + let mut data = data + .get_mut::() + .ok_or(Error::from("DB initialized")) + .map(|v| DBWriteGuard::from(v))?; + let mut data = data.borrow_mut()?; + let mut server_ban = data.get_mut(&guild).and_then(|v| match v { + ServerSoftBans::Unimplemented => None, + ServerSoftBans::Implemented(ref mut v) => Some(v), + }); + + match server_ban { + None => { + println!("get here"); + msg.reply(&ctx, format!("⚠ This server has not enabled the soft-ban feature. Check out `y!a soft-ban-init`."))?; + } + Some(ref mut server_ban) => { + let mut member = guild.member(&ctx, &user)?; + match duration { + None if member.roles.contains(&server_ban.role) => { + msg.reply(&ctx, format!("⛓ Lifting soft-ban for user {}.", user.tag()))?; + member.remove_role(&ctx, server_ban.role)?; + return Ok(()); + } + None => { + msg.reply(&ctx, format!("⛓ Soft-banning user {}.", user.tag()))?; + } + Some(v) => { + let until = Utc::now() + v.0; + let until = server_ban + .periodical_bans + .entry(user.id) + .and_modify(|v| *v = max(*v, until)) + .or_insert(until); + msg.reply( + &ctx, + format!("⛓ Soft-banning user {} until {}.", user.tag(), until), + )?; + } + } + member.add_role(&ctx, server_ban.role)?; + } + } + + Ok(()) +} + +#[command] +#[required_permissions(ADMINISTRATOR)] +#[description = "Sets up the soft-ban command. This command can only be run once.\nThe soft-ban command assigns a role, temporarily, to a user."] +#[usage = "{soft_ban_role_id}"] +#[num_args(1)] +pub fn soft_ban_init(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult { + let role_id = args.single::()?; + let guild = msg.guild(&ctx).ok_or(Error::from("Guild-only command"))?; + let guild = guild.read(); + // Check whether the role_id is the one we wanted + if !guild.roles.contains_key(&role_id) { + return Err(Error::from(format!( + "{} is not a role in this server.", + role_id + ))); + } + // Check if we already set up + let mut data = ctx.data.write(); + let mut db: DBWriteGuard<_> = data + .get_mut::() + .ok_or(Error::from("DB uninitialized"))? + .into(); + let mut db = db.borrow_mut()?; + let server = db + .get_mut(&guild.id) + .map(|v| match v { + ServerSoftBans::Unimplemented => false, + _ => true, + }) + .unwrap_or(false); + + if !server { + db.insert(guild.id, ServerSoftBans::new_implemented(role_id)); + msg.react(&ctx, "👌")?; + Ok(()) + } else { + Err(Error::from("Server already set up soft-bans.")) + } +} + +// Watch the soft bans. +pub fn watch_soft_bans(client: &mut serenity::Client) -> impl FnOnce() -> () + 'static { + let cache_http = { + let cache_http = client.cache_and_http.clone(); + let cache: serenity::cache::CacheRwLock = cache_http.cache.clone().into(); + (cache, cache_http.http.clone()) + }; + let data = client.data.clone(); + return move || { + let cache_http = (&cache_http.0, &*cache_http.1); + loop { + // Scope so that locks are released + { + // Poll the data for any changes. + let mut data = data.write(); + let mut db: DBWriteGuard<_> = data + .get_mut::() + .expect("DB wrongly initialized") + .into(); + let mut db = db.borrow_mut().expect("cannot unpack DB"); + let now = Utc::now(); + for (server_id, soft_bans) in db.iter_mut() { + let server_name: String = match server_id.to_partial_guild(cache_http) { + Err(_) => continue, + Ok(v) => v.name, + }; + if let ServerSoftBans::Implemented(ref mut bans) = soft_bans { + let to_remove: Vec<_> = bans + .periodical_bans + .iter() + .filter_map(|(user, time)| if time <= &now { Some(user) } else { None }) + .cloned() + .collect(); + for user_id in to_remove { + server_id + .member(cache_http, user_id) + .and_then(|mut m| { + println!( + "Soft-ban for `{}` in server `{}` unlifted.", + m.user.read().name, + server_name + ); + m.remove_role(cache_http, bans.role) + }) + .unwrap_or(()); + bans.periodical_bans.remove(&user_id); + } + } + } + } + // Sleep the thread for a minute + std::thread::sleep(std::time::Duration::from_secs(60)) + } + }; +} diff --git a/youmubot/src/commands/args.rs b/youmubot/src/commands/args.rs new file mode 100644 index 0000000..4a1db68 --- /dev/null +++ b/youmubot/src/commands/args.rs @@ -0,0 +1,117 @@ +pub use duration::Duration; + +mod duration { + use chrono::Duration as StdDuration; + use serenity::framework::standard::CommandError as Error; + // Parse a single duration unit + fn parse_duration_string(s: &str) -> Result { + // We reject the empty case + if s == "" { + return Err(Error::from("empty strings are not valid durations")); + } + struct ParseStep { + current_value: Option, + current_duration: StdDuration, + } + s.chars() + .try_fold( + ParseStep { + current_value: None, + current_duration: StdDuration::zero(), + }, + |s, item| match (item, s.current_value) { + ('0'..='9', v) => Ok(ParseStep { + current_value: Some(v.unwrap_or(0) * 10 + ((item as u64) - ('0' as u64))), + ..s + }), + (_, None) => Err(Error::from("Not a valid duration")), + (item, Some(v)) => Ok(ParseStep { + current_value: None, + current_duration: s.current_duration + + match item { + 's' => StdDuration::seconds, + 'm' => StdDuration::minutes, + 'h' => StdDuration::hours, + 'D' => StdDuration::days, + 'W' => StdDuration::weeks, + _ => return Err(Error::from("Not a valid duration")), + }(v as i64), + }), + }, + ) + .and_then(|v| match v.current_value { + // All values should be consumed + None => Ok(v), + _ => Err(Error::from("Not a valid duration")), + }) + .map(|v| v.current_duration) + } + + // Our new-orphan-type of duration. + #[derive(Copy, Clone, Debug)] + pub struct Duration(pub StdDuration); + + impl std::str::FromStr for Duration { + type Err = Error; + fn from_str(s: &str) -> Result { + parse_duration_string(s).map(|v| Duration(v)) + } + } + + impl From for StdDuration { + fn from(d: Duration) -> Self { + d.0 + } + } + + #[cfg(test)] + mod tests { + use super::*; + use chrono::Duration as StdDuration; + #[test] + fn test_parse_success() { + let tests = [ + ( + "2D2h1m", + StdDuration::seconds(2 * 60 * 60 * 24 + 2 * 60 * 60 + 1 * 60), + ), + ( + "1W2D3h4m5s", + StdDuration::seconds( + 1 * 7 * 24 * 60 * 60 + // 1W + 2 * 24 * 60 * 60 + // 2D + 3 * 60 * 60 + // 3h + 4 * 60 + // 4m + 5, // 5s + ), + ), + ( + "1W2D3h4m5s6W", + StdDuration::seconds( + 1 * 7 * 24 * 60 * 60 + // 1W + 2 * 24 * 60 * 60 + // 2D + 3 * 60 * 60 + // 3h + 4 * 60 + // 4m + 5 + // 5s + 6 * 7 * 24 * 60 * 60, + ), // 6W + ), + ]; + for (input, output) in &tests { + assert_eq!(parse_duration_string(input).unwrap(), *output); + } + } + + #[test] + fn test_parse_fail() { + let tests = ["", "1w", "-1W", "1"]; + for input in &tests { + assert!( + parse_duration_string(input).is_err(), + "parsing {} succeeded", + input + ); + } + } + } +} diff --git a/youmubot/src/commands/mod.rs b/youmubot/src/commands/mod.rs index 64fcc87..dfca885 100644 --- a/youmubot/src/commands/mod.rs +++ b/youmubot/src/commands/mod.rs @@ -7,7 +7,8 @@ use serenity::{ }; use std::collections::HashSet; -mod admin; +pub mod admin; +mod args; pub use admin::ADMIN_GROUP; diff --git a/youmubot/src/db/mod.rs b/youmubot/src/db/mod.rs new file mode 100644 index 0000000..34a2fbe --- /dev/null +++ b/youmubot/src/db/mod.rs @@ -0,0 +1,111 @@ +use chrono::{DateTime, Utc}; +use dotenv::var; +use rustbreak::{deser::Yaml as Ron, FileDatabase}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use serenity::{ + client::Client, + framework::standard::CommandError as Error, + model::id::{GuildId, RoleId, UserId}, + prelude::*, +}; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +/// GuildMap defines the guild-map type. +/// It is basically a HashMap from a GuildId to a data structure. +pub type GuildMap = HashMap; +/// The generic DB type we will be using. +pub struct DB(std::marker::PhantomData); +impl serenity::prelude::TypeMapKey for DB { + type Value = FileDatabase; +} + +impl DB +where + for<'de> T: Deserialize<'de>, +{ + fn insert_into(data: &mut ShareMap, path: impl AsRef) -> Result<(), Error> { + let db = FileDatabase::::from_path(path, T::default())?; + db.load().or_else(|_| db.save())?; + data.insert::>(db); + Ok(()) + } +} + +/// A list of SoftBans for all servers. +pub type SoftBans = DB>; + +/// Sets up all databases in the client. +pub fn setup_db(client: &mut Client) -> Result<(), Error> { + let path: PathBuf = var("DBPATH").map(|v| PathBuf::from(v)).unwrap_or_else(|e| { + println!("No DBPATH set up ({:?}), using `/data`", e); + PathBuf::from("data") + }); + let mut data = client.data.write(); + SoftBans::insert_into(&mut *data, &path.join("soft_bans.ron"))?; + + Ok(()) +} + +pub struct DBWriteGuard<'a, T>(&'a mut FileDatabase) +where + T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned; + +impl<'a, T> From<&'a mut FileDatabase> for DBWriteGuard<'a, T> +where + T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned, +{ + fn from(v: &'a mut FileDatabase) -> Self { + DBWriteGuard(v) + } +} + +impl<'a, T> DBWriteGuard<'a, T> +where + T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned, +{ + pub fn borrow(&self) -> Result, rustbreak::RustbreakError> { + (*self).0.borrow_data() + } + pub fn borrow_mut( + &mut self, + ) -> Result, rustbreak::RustbreakError> { + (*self).0.borrow_data_mut() + } +} + +impl<'a, T> Drop for DBWriteGuard<'a, T> +where + T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned, +{ + fn drop(&mut self) { + self.0.save().expect("Save succeed") + } +} + +/// For the admin commands: +/// - Each server might have a `soft ban` role implemented. +/// - We allow periodical `soft ban` applications. +#[derive(Serialize, Deserialize, Debug, Clone)] +pub enum ServerSoftBans { + Implemented(ImplementedSoftBans), + Unimplemented, +} + +impl ServerSoftBans { + // Create a new, implemented role. + pub fn new_implemented(role: RoleId) -> ServerSoftBans { + ServerSoftBans::Implemented(ImplementedSoftBans { + role: role, + periodical_bans: HashMap::new(), + }) + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct ImplementedSoftBans { + /// The soft-ban role. + pub role: RoleId, + /// List of all to-unban people. + pub periodical_bans: HashMap>, +} diff --git a/youmubot/src/main.rs b/youmubot/src/main.rs index dd22ef4..56eec78 100644 --- a/youmubot/src/main.rs +++ b/youmubot/src/main.rs @@ -7,6 +7,7 @@ use serenity::{ }; mod commands; +mod db; struct Handler; @@ -30,6 +31,12 @@ fn main() { setup_framework(Client::new(token, Handler).expect("Cannot connect...")) }; + // Setup initial data + db::setup_db(&mut client).expect("Setup db should succeed"); + + // Create handler threads + std::thread::spawn(commands::admin::watch_soft_bans(&mut client)); + println!("Starting..."); if let Err(v) = client.start() { panic!(v)