From a248066e9e0d1a05225a635f60ac8460c45ef26b Mon Sep 17 00:00:00 2001 From: minneelyyyy Date: Mon, 17 Feb 2025 17:21:40 -0500 Subject: [PATCH] update v0.1.2, update self roles code --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/commands/mod.rs | 2 +- src/commands/self_roles/color.rs | 63 ++++++++++++----------------- src/commands/self_roles/disown.rs | 38 ++++++++--------- src/commands/self_roles/mod.rs | 40 +++++++++++++++++- src/commands/self_roles/name.rs | 60 ++++++++++++--------------- src/commands/self_roles/register.rs | 46 +++++++++------------ src/commands/self_roles/remove.rs | 46 +++++++++------------ src/commands/self_roles/whois.rs | 8 +--- src/main.rs | 25 ++++++++++-- 11 files changed, 172 insertions(+), 160 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa4b04b..7b53930 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -202,7 +202,7 @@ dependencies = [ [[package]] name = "bot" -version = "0.1.1" +version = "0.1.2" dependencies = [ "clap", "dotenv", diff --git a/Cargo.toml b/Cargo.toml index dc5dc47..2176875 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "bot" -version = "0.1.1" +version = "0.1.2" edition = "2021" [dependencies] diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 9e126e5..54e3cd2 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -5,7 +5,7 @@ mod dox; mod yeehaw; mod gambling; mod eval; -mod self_roles; +pub mod self_roles; mod settings; use crate::common::{Data, Error, Context}; diff --git a/src/commands/self_roles/color.rs b/src/commands/self_roles/color.rs index 395e9ce..ed230fd 100644 --- a/src/commands/self_roles/color.rs +++ b/src/commands/self_roles/color.rs @@ -1,11 +1,11 @@ -use crate::common::{Context, Error}; +use crate::common::{self, Context, Error}; use once_cell::sync::Lazy; use std::collections::HashMap; use hex_color::HexColor; -use poise::serenity_prelude::{colours, Color, EditRole}; +use poise::serenity_prelude::{colours, Color, EditRole, Permissions}; static COLORS: Lazy> = Lazy::new(|| { HashMap::from([ @@ -69,42 +69,29 @@ pub async fn color(ctx: Context<'_>, #[autocomplete = "autocomplete_colors"] col Color::from_rgb(rgb.r, rgb.g, rgb.b) }; - if let Some(guild) = ctx.guild_id() { - let mut tx = ctx.data().database.begin().await?; - - match super::get_user_role(ctx.author().id, guild, &mut *tx).await? { - Some(role) => { - guild.edit_role(ctx, role, EditRole::new().colour(color)).await?; - let role = guild.role(ctx, role).await?; - - ctx.reply(format!("{}'s color has been updated!", role)).await?; - - Ok(()) - }, - None => { - let role = guild.create_role(ctx, - EditRole::new() - .name(format!("#{:06x}", color.0)) - .colour(color)).await?; - - sqlx::query("INSERT INTO selfroles (userid, roleid, guildid) VALUES ($1, $2, $3)") - .bind(ctx.author().id.get() as i64) - .bind(role.id.get() as i64) - .bind(guild.get() as i64) - .execute(&mut *tx).await?; - - tx.commit().await?; - - let member = guild.member(ctx, ctx.author().id).await?; - member.add_role(ctx, role.clone()).await?; - - ctx.reply(format!("You have been given the {} role!", role)).await?; - - return Ok(()); - } - } + let guild = if let Some(guild) = ctx.guild_id() { + guild } else { - ctx.reply("This command must be run within a server.").await?; - Ok(()) + ctx.reply("This command can only be run inside of a guild.").await?; + return Ok(()); + }; + + let user = ctx.author(); + + let mut tx = ctx.data().database.begin().await?; + + if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? { + let role = guild.role(ctx, role).await?; + guild.edit_role(ctx, role.id, EditRole::new().colour(color)).await?; + common::no_ping_reply(&ctx, format!("{}'s color has been updated.", role)).await?; + } else { + let role = guild.create_role(ctx, EditRole::new().colour(color).name(user.name.clone()).permissions(Permissions::empty())).await?; + super::update_user_role(user.id, guild, role.id, &mut *tx).await?; + let member = guild.member(ctx, user).await?; + member.add_role(ctx, role.id).await?; + tx.commit().await?; + common::no_ping_reply(&ctx, format!("{} has been given the new role {}.", user, role)).await?; } + + Ok(()) } diff --git a/src/commands/self_roles/disown.rs b/src/commands/self_roles/disown.rs index c388c65..5f2a818 100644 --- a/src/commands/self_roles/disown.rs +++ b/src/commands/self_roles/disown.rs @@ -4,25 +4,25 @@ use crate::common::{Context, Error}; /// Remove and delete your personal role #[poise::command(slash_command, prefix_command)] pub async fn disown(ctx: Context<'_>) -> Result<(), Error> { - let db = &ctx.data().database; - - if let Some(guild) = ctx.guild_id() { - if let Some(role) = super::get_user_role(ctx.author().id, guild, db).await? { - guild.delete_role(ctx, role).await?; - - sqlx::query("DELETE FROM selfroles WHERE roleid = $1") - .bind(role.get() as i64) - .execute(db).await?; - - ctx.reply("Your role has been successfully removed.").await?; - - Ok(()) - } else { - ctx.reply("You do not currently have a personal role to remove.").await?; - Ok(()) - } + let guild = if let Some(guild) = ctx.guild_id() { + guild } else { - ctx.reply("This command must be called within a server.").await?; - Ok(()) + ctx.reply("This command can only be run inside of a guild.").await?; + return Ok(()); + }; + + let user = ctx.author(); + + let mut tx = ctx.data().database.begin().await?; + + if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? { + guild.delete_role(ctx, role).await?; + super::remove_user_role(user.id, guild, &mut *tx).await?; + tx.commit().await?; + ctx.reply("Your self role has been deleted.").await?; + } else { + ctx.reply("You currently have no self role to delete.").await?; } + + Ok(()) } diff --git a/src/commands/self_roles/mod.rs b/src/commands/self_roles/mod.rs index 2e20ce0..f9489b3 100644 --- a/src/commands/self_roles/mod.rs +++ b/src/commands/self_roles/mod.rs @@ -26,9 +26,45 @@ pub async fn role(_ctx: Context<'_>) -> Result<(), Error> { Ok(()) } +/// Remove a row concerning a user's self role from the database +pub async fn remove_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result<(), Error> + where E: PgExecutor<'a>, +{ + sqlx::query("DELETE FROM selfroles WHERE userid = $1 AND guildid = $2") + .bind(user.get() as i64) + .bind(guild.get() as i64) + .execute(db).await?; + + Ok(()) +} + +pub async fn remove_role<'a, E>(role: RoleId, guild: GuildId, db: E) -> Result<(), Error> +where E: PgExecutor<'a> +{ + sqlx::query("DELETE FROM selfroles WHERE roleid = $1 AND guildid = $2") + .bind(role.get() as i64) + .bind(guild.get() as i64) + .execute(db).await?; + + Ok(()) +} + +/// Replace a user's custom role with a new one +pub async fn update_user_role<'a, E>(user: UserId, guild: GuildId, role: RoleId, db: E) -> Result<(), Error> + where E: PgExecutor<'a>, +{ + sqlx::query("INSERT INTO selfroles (userid, guildid, roleid) VALUES($1, $2, $3) ON CONFLICT (userid, guildid) DO UPDATE SET roleid = EXCLUDED.roleid") + .bind(user.get() as i64) + .bind(guild.get() as i64) + .bind(role.get() as i64) + .execute(db).await?; + + Ok(()) +} + +/// Get a user's personal role id from the database pub async fn get_user_role<'a, E>(user: UserId, guild: GuildId, db: E) -> Result, Error> -where - E: PgExecutor<'a>, + where E: PgExecutor<'a>, { match sqlx::query("SELECT roleid FROM selfroles WHERE userid = $1 AND guildid = $2") .bind(user.get() as i64) diff --git a/src/commands/self_roles/name.rs b/src/commands/self_roles/name.rs index b9fa4a5..58760c1 100644 --- a/src/commands/self_roles/name.rs +++ b/src/commands/self_roles/name.rs @@ -1,44 +1,34 @@ -use crate::common::{Context, Error}; +use crate::common::{self, Context, Error}; -use poise::serenity_prelude::EditRole; +use poise::serenity_prelude::{EditRole, Permissions}; /// Change the name of your personal role #[poise::command(slash_command, prefix_command)] pub async fn name(ctx: Context<'_>, name: String) -> Result<(), Error> { - if let Some(guild) = ctx.guild_id() { - let mut tx = ctx.data().database.begin().await?; - - let role = match super::get_user_role(ctx.author().id, guild, &mut *tx).await? { - Some(role) => role, - None => { - let role = guild.create_role(ctx, EditRole::new().name(name)).await?; - - sqlx::query("INSERT INTO selfroles (userid, roleid, guildid) VALUES ($1, $2, $3)") - .bind(ctx.author().id.get() as i64) - .bind(role.id.get() as i64) - .bind(guild.get() as i64) - .execute(&mut *tx).await?; - - tx.commit().await?; - - let member = guild.member(ctx, ctx.author().id).await?; - member.add_role(ctx, role.clone()).await?; - - ctx.reply(format!("You've been given the {} role!", role)).await?; - - return Ok(()); - } - }; - - guild.edit_role(ctx, role, EditRole::new().name(name)).await?; - let role = guild.role(ctx, role).await?; - - ctx.reply(format!("{} has been updated.", role)).await?; - - Ok(()) + let guild = if let Some(guild) = ctx.guild_id() { + guild } else { - ctx.reply("This command must be run within a server.").await?; - Ok(()) + ctx.reply("This command can only be run inside of a guild.").await?; + return Ok(()); + }; + + let user = ctx.author(); + + let mut tx = ctx.data().database.begin().await?; + + if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? { + let role = guild.role(ctx, role).await?; + guild.edit_role(ctx, role.id, EditRole::new().name(name)).await?; + common::no_ping_reply(&ctx, format!("{} has been updated.", role)).await?; + } else { + let role = guild.create_role(ctx, EditRole::new().name(name).permissions(Permissions::empty())).await?; + super::update_user_role(user.id, guild, role.id, &mut *tx).await?; + let member = guild.member(ctx, user).await?; + member.add_role(ctx, role.id).await?; + tx.commit().await?; + common::no_ping_reply(&ctx, format!("{} has been given the new role {}.", user, role)).await?; } + + Ok(()) } diff --git a/src/commands/self_roles/register.rs b/src/commands/self_roles/register.rs index 2538db1..92ebb02 100644 --- a/src/commands/self_roles/register.rs +++ b/src/commands/self_roles/register.rs @@ -3,35 +3,29 @@ use crate::common::{self, Context, Error}; use poise::serenity_prelude as serenity; -/// Register an existing role as a user's custom role +/// Register an existing role as a user's custom role. This deletes their current self role. #[poise::command(slash_command, prefix_command, required_permissions = "MANAGE_ROLES")] -pub async fn register(ctx: Context<'_>, user: serenity::User, role: serenity::Role) -> Result<(), Error> { - if let Some(guild) = ctx.guild_id() { - let mut tx = ctx.data().database.begin().await?; - - match super::get_user_role(user.id, guild, &mut *tx).await? { - Some(role) => { - let role = guild.role(ctx, role).await?; - common::no_ping_reply(&ctx, format!("{} already has the role {} registered.", user, role)).await?; - }, - None => { - sqlx::query("INSERT INTO selfroles (userid, guildid, roleid) VALUES ($1, $2, $3)") - .bind(user.id.get() as i64) - .bind(guild.get() as i64) - .bind(role.id.get() as i64) - .execute(&mut *tx).await?; - - tx.commit().await?; - - let member = guild.member(ctx, user.id).await?; - member.add_role(ctx, role.id).await?; - - common::no_ping_reply(&ctx, format!("{} now has {} set as their personal role.", user, role)).await?; - } - } +pub async fn register(ctx: Context<'_>, user: serenity::User, role: serenity::Role) -> Result<(), Error> { + let guild = if let Some(guild) = ctx.guild_id() { + guild } else { - ctx.reply("This command can only be run in a guild!").await?; + ctx.reply("This command can only be run inside of a guild.").await?; + return Ok(()); + }; + + let mut tx = ctx.data().database.begin().await?; + + if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? { + guild.delete_role(ctx, role).await?; } + let member = guild.member(ctx, user).await?; + member.add_role(ctx, role.id).await?; + + super::update_user_role(member.user.id, guild, role.id, &mut *tx).await?; + tx.commit().await?; + + common::no_ping_reply(&ctx, format!("{} has been set as {}'s self role.", role, member.user)).await?; + Ok(()) } diff --git a/src/commands/self_roles/remove.rs b/src/commands/self_roles/remove.rs index f842c7d..e145c12 100644 --- a/src/commands/self_roles/remove.rs +++ b/src/commands/self_roles/remove.rs @@ -3,35 +3,25 @@ use crate::common::{self, Context, Error}; use poise::serenity_prelude as serenity; -/// force remove someone's role (this does not delete the role) #[poise::command(slash_command, prefix_command, required_permissions = "MANAGE_ROLES")] pub async fn remove(ctx: Context<'_>, user: serenity::User) -> Result<(), Error> { - let db = &ctx.data().database; - - if let Some(guild) = ctx.guild_id() { - match super::get_user_role(user.id, guild, db).await? { - Some(role) => { - sqlx::query("DELETE FROM selfroles WHERE userid = $1 AND guildid = $2") - .bind(user.id.get() as i64) - .bind(guild.get() as i64) - .execute(db).await?; - - let member = guild.member(ctx, user.id).await?; - - member.remove_role(ctx, role).await?; - let role = guild.role(ctx, role).await?; - - common::no_ping_reply(&ctx, format!("{} has had {} removed.", user, role)).await?; - - Ok(()) - }, - None => { - common::no_ping_reply(&ctx, format!("{} does not have a personal role to remove.", user)).await?; - Ok(()) - } - } + let guild = if let Some(guild) = ctx.guild_id() { + guild } else { - ctx.reply("This command can only be run in a guild!").await?; - Ok(()) + ctx.reply("This command can only be run inside of a guild.").await?; + return Ok(()); + }; + + let mut tx = ctx.data().database.begin().await?; + + if let Some(role) = super::get_user_role(user.id, guild, &mut *tx).await? { + guild.delete_role(ctx, role).await?; + super::remove_role(role, guild, &mut *tx).await?; + tx.commit().await?; + common::no_ping_reply(&ctx, format!("{}'s self role has been deleted.", user)).await?; + } else { + common::no_ping_reply(&ctx, format!("{} has no self role.", user)).await?; } -} + + Ok(()) +} \ No newline at end of file diff --git a/src/commands/self_roles/whois.rs b/src/commands/self_roles/whois.rs index 1bd6417..de1d177 100644 --- a/src/commands/self_roles/whois.rs +++ b/src/commands/self_roles/whois.rs @@ -11,11 +11,11 @@ pub async fn whois(ctx: Context<'_>, role: serenity::Role) -> Result<(), Error> let db = &ctx.data().database; if let Some(guild) = ctx.guild_id() { - let row = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1") + let user = match sqlx::query("SELECT userid FROM selfroles WHERE roleid = $1") .bind(role.id.get() as i64) .fetch_one(db).await { - Ok(row) => row, + Ok(row) => UserId::new(row.try_get::(0)? as u64), Err(sqlx::Error::RowNotFound) => { ctx.reply("This role is not owned by anyone.").await?; return Ok(()); @@ -23,10 +23,6 @@ pub async fn whois(ctx: Context<'_>, role: serenity::Role) -> Result<(), Error> Err(e) => return Err(Box::new(e)), }; - let user: i64 = row.try_get(0)?; - - let user = UserId::new(user as u64); - let member = guild.member(ctx, user).await?; common::no_ping_reply(&ctx, format!("{} owns this role.", member)).await?; diff --git a/src/main.rs b/src/main.rs index 72ae44a..9d43278 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ mod commands; pub mod common; pub mod inventory; +use crate::commands::self_roles; use crate::common::{Context, Error, Data}; use std::env; @@ -22,15 +23,31 @@ struct BotArgs { } async fn event_handler( - _ctx: &serenity::Context, + ctx: &serenity::Context, event: &serenity::FullEvent, _framework: poise::FrameworkContext<'_, Data, Error>, - _data: &Data, + data: &Data, ) -> Result<(), Error> { match event { serenity::FullEvent::Message { new_message: message } => { if message.author.bot { return Ok(()) } } + serenity::FullEvent::GuildMemberRemoval { guild_id, user, .. } => { + let mut tx = data.database.begin().await?; + + if let Some(role) = self_roles::get_user_role(user.id, *guild_id, &mut *tx).await? { + guild_id.delete_role(ctx, role).await?; + self_roles::remove_role(role, *guild_id, &mut *tx).await?; + tx.commit().await?; + } + }, + serenity::FullEvent::GuildRoleDelete { guild_id, removed_role_id, .. } => { + let mut tx = data.database.begin().await?; + + self_roles::remove_role(*removed_role_id, *guild_id, &mut *tx).await?; + + tx.commit().await?; + } _ => (), } @@ -64,7 +81,9 @@ async fn main() -> Result<(), Error> { serenity::GatewayIntents::GUILD_MESSAGES | serenity::GatewayIntents::DIRECT_MESSAGES | serenity::GatewayIntents::MESSAGE_CONTENT) - .unwrap_or(serenity::GatewayIntents::empty()); + .unwrap_or(serenity::GatewayIntents::empty()) + | serenity::GatewayIntents::GUILD_MEMBERS + | serenity::GatewayIntents::GUILDS; let framework = poise::Framework::builder() .options(poise::FrameworkOptions {