From d6e6ab9699616d9c18ab51d9089f0a5b27897db5 Mon Sep 17 00:00:00 2001 From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com> Date: Sun, 19 Jun 2022 13:36:06 +0200 Subject: [PATCH] prison --- src/handler.rs | 372 +++++++++++++++++++++++++++++++++++++------------ src/main.rs | 1 + src/model.rs | 89 ++++++++++-- 3 files changed, 365 insertions(+), 97 deletions(-) diff --git a/src/handler.rs b/src/handler.rs index f418c64..fb89d80 100644 --- a/src/handler.rs +++ b/src/handler.rs @@ -18,91 +18,136 @@ use crate::{ }; fn slash_commands(commands: &mut CreateApplicationCommands) -> &mut CreateApplicationCommands { - commands.create_application_command(|command| { - command - .name("lawsuit") - .description("Einen Gerichtsprozess starten") - .create_option(|option| { - option - .name("create") - .description("Einen neuen Gerichtsprozess anfangen") - .kind(ApplicationCommandOptionType::SubCommand) - .create_sub_option(|option| { - option - .name("plaintiff") - .description("Der Kläger") - .kind(ApplicationCommandOptionType::User) - .required(true) - }) - .create_sub_option(|option| { - option - .name("accused") - .description("Der Angeklagte") - .kind(ApplicationCommandOptionType::User) - .required(true) - }) - .create_sub_option(|option| { - option - .name("judge") - .description("Der Richter") - .kind(ApplicationCommandOptionType::User) - .required(true) - }) - .create_sub_option(|option| { - option - .name("reason") - .description("Der Grund für die Klage") - .kind(ApplicationCommandOptionType::String) - .required(true) - }) - .create_sub_option(|option| { - option - .name("plaintiff_lawyer") - .description("Der Anwalt des Klägers") - .kind(ApplicationCommandOptionType::User) - .required(false) - }) - .create_sub_option(|option| { - option - .name("accused_lawyer") - .description("Der Anwalt des Angeklagten") - .kind(ApplicationCommandOptionType::User) - .required(false) - }) - }) - .create_option(|option| { - option - .name("set_category") - .description("Die Gerichtskategorie setzen") - .kind(ApplicationCommandOptionType::SubCommand) - .create_sub_option(|option| { - option - .name("category") - .description("Die Kategorie") - .kind(ApplicationCommandOptionType::Channel) - .required(true) - }) - }) - .create_option(|option| { - option - .name("close") - .description("Den Prozess abschliessen") - .kind(ApplicationCommandOptionType::SubCommand) - .create_sub_option(|option| { - option - .name("verdict") - .description("Das Urteil") - .kind(ApplicationCommandOptionType::String) - .required(true) - }) - }) - .create_option(|option| { - option - .name("clear") - .description("Alle Rechtsprozessdaten löschen") - .kind(ApplicationCommandOptionType::SubCommand) - }) - }) + commands + .create_application_command(|command| { + command + .name("lawsuit") + .description("Einen Gerichtsprozess starten") + .create_option(|option| { + option + .name("create") + .description("Einen neuen Gerichtsprozess anfangen") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("plaintiff") + .description("Der Kläger") + .kind(ApplicationCommandOptionType::User) + .required(true) + }) + .create_sub_option(|option| { + option + .name("accused") + .description("Der Angeklagte") + .kind(ApplicationCommandOptionType::User) + .required(true) + }) + .create_sub_option(|option| { + option + .name("judge") + .description("Der Richter") + .kind(ApplicationCommandOptionType::User) + .required(true) + }) + .create_sub_option(|option| { + option + .name("reason") + .description("Der Grund für die Klage") + .kind(ApplicationCommandOptionType::String) + .required(true) + }) + .create_sub_option(|option| { + option + .name("plaintiff_lawyer") + .description("Der Anwalt des Klägers") + .kind(ApplicationCommandOptionType::User) + .required(false) + }) + .create_sub_option(|option| { + option + .name("accused_lawyer") + .description("Der Anwalt des Angeklagten") + .kind(ApplicationCommandOptionType::User) + .required(false) + }) + }) + .create_option(|option| { + option + .name("set_category") + .description("Die Gerichtskategorie setzen") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("category") + .description("Die Kategorie") + .kind(ApplicationCommandOptionType::Channel) + .required(true) + }) + }) + .create_option(|option| { + option + .name("close") + .description("Den Prozess abschliessen") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("verdict") + .description("Das Urteil") + .kind(ApplicationCommandOptionType::String) + .required(true) + }) + }) + .create_option(|option| { + option + .name("clear") + .description("Alle Rechtsprozessdaten löschen") + .kind(ApplicationCommandOptionType::SubCommand) + }) + }) + .create_application_command(|command| { + command + .name("prison") + .description("Leute im Gefängnis einsperren") + .create_option(|option| { + option + .name("arrest") + .description("Jemanden einsperren") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("user") + .description("Die Person zum einsperren") + .kind(ApplicationCommandOptionType::User) + .required(true) + }) + }) + .create_option(|option| { + option + .name("release") + .description("Jemanden freilassen") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("user") + .description("Die Person zum freilassen") + .kind(ApplicationCommandOptionType::User) + .required(true) + }) + }) + .create_option(|option| { + option + .name("set_role") + .description("Die Rolle für Gefangene setzen") + .kind(ApplicationCommandOptionType::SubCommand) + .create_sub_option(|option| { + option + .name("role") + .description("Die Rolle") + .kind(ApplicationCommandOptionType::Role) + .required(true) + }) + }) + }) } pub struct Handler { @@ -119,6 +164,12 @@ pub enum Response { #[async_trait] impl EventHandler for Handler { + async fn guild_member_addition(&self, ctx: Context, new_member: Member) { + if let Err(err) = self.handle_guild_member_join(ctx, new_member).await { + error!(?err, "An error occurred in guild_member_addition handler"); + } + } + async fn ready(&self, ctx: Context, ready: Ready) { info!(name = %ready.user.name, "Bot is connected!"); @@ -146,7 +197,7 @@ impl EventHandler for Handler { async fn interaction_create(&self, ctx: Context, interaction: Interaction) { if let Interaction::ApplicationCommand(command) = interaction { if let Err(err) = self.handle_interaction(ctx, command).await { - error!(?err, "An error occurred"); + error!(?err, "An error occurred in interaction_create handler"); } } } @@ -161,6 +212,7 @@ impl Handler { let response = match command.data.name.as_str() { "lawsuit" => lawsuit_command_handler(&command, &ctx, &self.mongo).await, + "prison" => prison_command_handler(&command, &ctx, &self.mongo).await, _ => Ok(Response::EphemeralStr("not implemented :(")), }; @@ -178,6 +230,36 @@ impl Handler { } } + async fn handle_guild_member_join( + &self, + ctx: Context, + mut member: Member, + ) -> color_eyre::Result<()> { + let guild_id = member.guild_id; + let user_id = member.user.id; + let state = self.mongo.find_or_insert_state(guild_id.into()).await?; + + debug!(member = ?member.user.id, "New member joined"); + + if let Some(role_id) = state.prison_role { + if self + .mongo + .find_prison_entry(guild_id.into(), user_id.into()) + .await? + .is_some() + { + info!("New member was in prison, giving them the prison role"); + + member + .add_role(&ctx.http, role_id) + .await + .wrap_err("add role to member in prison")?; + } + } + + Ok(()) + } + async fn send_response( &self, ctx: Context, @@ -369,6 +451,107 @@ async fn lawsuit_command_handler( } } +async fn prison_command_handler( + command: &ApplicationCommandInteraction, + ctx: &Context, + mongo_client: &Mongo, +) -> color_eyre::Result { + let options = &command.data.options; + let subcommand = options.get(0).wrap_err("needs subcommand")?; + + let options = &subcommand.options; + let guild_id = command.guild_id.wrap_err("guild_id not found")?; + + let member = command + .member + .as_ref() + .wrap_err("command must be used my member")?; + let permissions = member.permissions.wrap_err("must be in interaction")?; + + match subcommand.name.as_str() { + "set_role" => { + if !permissions.contains(Permissions::MANAGE_GUILD) { + return Ok(Response::NoPermissions); + } + + let role = RoleOption::get(options.get(0))?; + + mongo_client + .set_prison_role(guild_id.into(), role.id.into()) + .await?; + + Ok(Response::EphemeralStr("isch gsetzt")) + } + "arrest" => { + if !permissions.contains(Permissions::MANAGE_GUILD) { + return Ok(Response::NoPermissions); + } + + let (user, _) = UserOption::get(options.get(0))?; + + let state = mongo_client.find_or_insert_state(guild_id.into()).await?; + let role = state.prison_role; + + let role = match role { + Some(role) => role, + None => { + return Ok(Response::EphemeralStr( + "du mosch zerst e rolle setze mit /prison set_role", + )) + } + }; + + mongo_client + .add_to_prison(guild_id.into(), user.id.into()) + .await?; + + guild_id + .member(&ctx.http, user.id) + .await + .wrap_err("fetching guild member")? + .add_role(&ctx.http, role) + .await + .wrap_err("add guild member role")?; + + Ok(Response::EphemeralStr("hani igsperrt")) + } + "release" => { + if !permissions.contains(Permissions::MANAGE_GUILD) { + return Ok(Response::NoPermissions); + } + + let (user, _) = UserOption::get(options.get(0))?; + + let state = mongo_client.find_or_insert_state(guild_id.into()).await?; + let role = state.prison_role; + + let role = match role { + Some(role) => role, + None => { + return Ok(Response::EphemeralStr( + "du mosch zerst e rolle setze mit /prison set_role", + )) + } + }; + + mongo_client + .remove_from_prison(guild_id.into(), user.id.into()) + .await?; + + guild_id + .member(&ctx.http, user.id) + .await + .wrap_err("fetching guild member")? + .remove_role(&ctx.http, role) + .await + .wrap_err("remove guild member role")?; + + Ok(Response::EphemeralStr("d'freiheit wartet")) + } + _ => Err(eyre!("Unknown subcommand")), + } +} + #[nougat::gat] trait GetOption { type Get<'a>; @@ -452,3 +635,20 @@ impl GetOption for ChannelOption { } } } + +struct RoleOption; + +#[nougat::gat] +impl GetOption for RoleOption { + type Get<'a> = &'a Role; + + fn extract( + command: &ApplicationCommandInteractionDataOptionValue, + ) -> crate::Result> { + if let ApplicationCommandInteractionDataOptionValue::Role(role) = command { + Ok(role) + } else { + Err(eyre!("Expected string!")) + } + } +} diff --git a/src/main.rs b/src/main.rs index caa8558..89366df 100644 --- a/src/main.rs +++ b/src/main.rs @@ -57,6 +57,7 @@ async fn main() -> Result<()> { set_global_commands, mongo, }) + .intents(GatewayIntents::GUILD_MEMBERS) .await .wrap_err("failed to create discord client")?; diff --git a/src/model.rs b/src/model.rs index 84937df..3132786 100644 --- a/src/model.rs +++ b/src/model.rs @@ -8,7 +8,7 @@ use color_eyre::Result; use mongodb::{ bson, bson::{doc, Bson, Uuid}, - options::{ClientOptions, Credential}, + options::{ClientOptions, Credential, UpdateOptions}, Client, Collection, Database, }; use serde::{Deserialize, Serialize}; @@ -92,6 +92,7 @@ pub struct State { pub lawsuits: Vec, pub court_category: Option, pub court_rooms: Vec, + pub prison_role: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -101,6 +102,12 @@ pub struct CourtRoom { pub role_id: SnowflakeId, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PrisonEntry { + pub guild_id: SnowflakeId, + pub user_id: SnowflakeId, +} + #[derive(Clone)] pub struct Mongo { db: Database, @@ -154,6 +161,7 @@ impl Mongo { lawsuits: vec![], court_category: None, court_rooms: vec![], + prison_role: None, }; let coll = self.db.collection::("state"); @@ -171,8 +179,8 @@ impl Mongo { let _ = self.find_or_insert_state(guild_id).await?; let coll = self.state_coll(); coll.update_one( - doc! {"guild_id": &guild_id }, - doc! {"$set": { "court_category": category }}, + doc! { "guild_id": &guild_id }, + doc! { "$set": { "court_category": category } }, None, ) .await @@ -180,6 +188,23 @@ impl Mongo { Ok(()) } + pub async fn set_prison_role( + &self, + guild_id: SnowflakeId, + prison_role: SnowflakeId, + ) -> Result<()> { + let _ = self.find_or_insert_state(guild_id).await?; + let coll = self.state_coll(); + coll.update_one( + doc! { "guild_id": &guild_id }, + doc! { "$set": { "prison_role": prison_role } }, + None, + ) + .await + .wrap_err("update prison role")?; + Ok(()) + } + pub async fn add_court_room(&self, guild_id: SnowflakeId, room: &CourtRoom) -> Result<()> { let _ = self.find_or_insert_state(guild_id).await?; let coll = self.state_coll(); @@ -246,22 +271,64 @@ impl Mongo { Ok(()) } - pub async fn delete_guild( - &self, - guild_id: SnowflakeId, - ) -> Result<()> { + pub async fn delete_guild(&self, guild_id: SnowflakeId) -> Result<()> { let coll = self.state_coll(); - coll.delete_one( - doc! { "guild_id": &guild_id }, - None, + coll.delete_one(doc! { "guild_id": &guild_id }, None) + .await + .wrap_err("delete guild")?; + Ok(()) + } + + pub async fn add_to_prison(&self, guild_id: SnowflakeId, user_id: SnowflakeId) -> Result<()> { + let coll = self.prison_coll(); + + coll.update_one( + doc! { "guild_id": guild_id, "user_id": user_id }, + doc! { + "$setOnInsert": { + "guild_id": guild_id, "user_id": user_id, + } + }, + UpdateOptions::builder().upsert(true).build(), ) .await - .wrap_err("delete guild")?; + .wrap_err("add to prison collection")?; + Ok(()) } + pub async fn remove_from_prison( + &self, + guild_id: SnowflakeId, + user_id: SnowflakeId, + ) -> Result<()> { + let coll = self.prison_coll(); + + coll.delete_one(doc! { "guild_id": guild_id, "user_id": user_id }, None) + .await + .wrap_err("remove from prison")?; + + Ok(()) + } + + pub async fn find_prison_entry( + &self, + guild_id: SnowflakeId, + user_id: SnowflakeId, + ) -> Result> { + let coll = self.prison_coll(); + + coll.find_one(doc! { "guild_id": guild_id, "user_id": user_id }, None) + .await + .wrap_err("remove from prison") + } + fn state_coll(&self) -> Collection { self.db.collection("state") } + + fn prison_coll(&self) -> Collection { + self.db.collection("prison") + } }