prison

src/handler.rs

@@ -18,7 +18,8 @@ use crate::{
 };
 
 fn slash_commands(commands: &mut CreateApplicationCommands) -> &mut CreateApplicationCommands {
-    commands.create_application_command(|command| {
+    commands
+        .create_application_command(|command| {
             command
                 .name("lawsuit")
                 .description("Einen Gerichtsprozess starten")
@@ -103,6 +104,50 @@ fn slash_commands(commands: &mut CreateApplicationCommands) -> &mut CreateApplic
                         .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<Response> {
+    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<Self::Get<'_>> {
+        if let ApplicationCommandInteractionDataOptionValue::Role(role) = command {
+            Ok(role)
+        } else {
+            Err(eyre!("Expected string!"))
+        }
+    }
+}

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")?;
 

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<Lawsuit>,
     pub court_category: Option<SnowflakeId>,
     pub court_rooms: Vec<CourtRoom>,
+    pub prison_role: Option<SnowflakeId>,
 }
 
 #[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>("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! { "guild_id": &guild_id  },
-            doc! {"$set": { "court_category": category }},
+            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(
+    pub async fn delete_guild(&self, guild_id: SnowflakeId) -> Result<()> {
-        &self,
-        guild_id: SnowflakeId,
-    ) -> Result<()> {
         let coll = self.state_coll();
 
-        coll.delete_one(
+        coll.delete_one(doc! { "guild_id": &guild_id }, None)
-            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("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<Option<PrisonEntry>> {
+        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<State> {
         self.db.collection("state")
     }
+
+    fn prison_coll(&self) -> Collection<PrisonEntry> {
+        self.db.collection("prison")
+    }
 }