Admin commands done (#1)

This commit is contained in:
natsukagami 2019-11-08 14:31:53 +00:00 committed by Gitea
parent 4f6a0a2316
commit 384b7be52c
9 changed files with 466 additions and 2 deletions

1
.gitignore vendored
View file

@ -1,2 +1,3 @@
target
.env
*.ron

41
Cargo.lock generated
View file

@ -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"

View file

@ -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"]

View file

@ -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]

View file

@ -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::<UserId>()?.to_user(&ctx)?;
let duration = if args.is_empty() {
None
} else {
Some(
args.single::<args::Duration>()
.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::<SoftBans>()
.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::<RoleId>()?;
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::<SoftBans>()
.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::<SoftBans>()
.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))
}
};
}

View file

@ -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<StdDuration, Error> {
// We reject the empty case
if s == "" {
return Err(Error::from("empty strings are not valid durations"));
}
struct ParseStep {
current_value: Option<u64>,
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<Self, Self::Err> {
parse_duration_string(s).map(|v| Duration(v))
}
}
impl From<Duration> 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
);
}
}
}
}

View file

@ -7,7 +7,8 @@ use serenity::{
};
use std::collections::HashSet;
mod admin;
pub mod admin;
mod args;
pub use admin::ADMIN_GROUP;

111
youmubot/src/db/mod.rs Normal file
View file

@ -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<V> = HashMap<GuildId, V>;
/// The generic DB type we will be using.
pub struct DB<T>(std::marker::PhantomData<T>);
impl<T: std::any::Any> serenity::prelude::TypeMapKey for DB<T> {
type Value = FileDatabase<T, Ron>;
}
impl<T: std::any::Any + Default + Send + Sync + Clone + Serialize + std::fmt::Debug> DB<T>
where
for<'de> T: Deserialize<'de>,
{
fn insert_into(data: &mut ShareMap, path: impl AsRef<Path>) -> Result<(), Error> {
let db = FileDatabase::<T, Ron>::from_path(path, T::default())?;
db.load().or_else(|_| db.save())?;
data.insert::<DB<T>>(db);
Ok(())
}
}
/// A list of SoftBans for all servers.
pub type SoftBans = DB<GuildMap<ServerSoftBans>>;
/// 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<T, Ron>)
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned;
impl<'a, T> From<&'a mut FileDatabase<T, Ron>> for DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
fn from(v: &'a mut FileDatabase<T, Ron>) -> Self {
DBWriteGuard(v)
}
}
impl<'a, T> DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
pub fn borrow(&self) -> Result<std::sync::RwLockReadGuard<T>, rustbreak::RustbreakError> {
(*self).0.borrow_data()
}
pub fn borrow_mut(
&mut self,
) -> Result<std::sync::RwLockWriteGuard<T>, 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<UserId, DateTime<Utc>>,
}

View file

@ -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)