diff --git a/Cargo.toml b/Cargo.toml index 4896e94..2d8ced7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,4 +10,5 @@ dotenv = "0.15.0" rand = "0.8.5" clap = { version = "4.5.20", features = ["derive"] } lamm = { git = "https://github.com/minneelyyyy/lamm", branch = "dev" } -regex = "1.11.1" \ No newline at end of file +regex = "1.11.1" +sqlx = { version = "0.8", features = [ "runtime-tokio-native-tls", "postgres" ] } \ No newline at end of file diff --git a/src/commands/gambling/balance.rs b/src/commands/gambling/balance.rs index 92a4a66..651398c 100644 --- a/src/commands/gambling/balance.rs +++ b/src/commands/gambling/balance.rs @@ -1,4 +1,3 @@ -use super::get_user_wealth_mut; use crate::common::{Context, Error}; use poise::serenity_prelude as serenity; @@ -6,16 +5,15 @@ use poise::serenity_prelude as serenity; #[poise::command(slash_command, prefix_command)] pub async fn balance(ctx: Context<'_>, user: Option) -> Result<(), Error> { let user = user.as_ref().unwrap_or(ctx.author()); - let mut users = ctx.data().users.lock().await; - let wealth = get_user_wealth_mut(&mut users, user.id); + let wealth = super::get_balance(user.id, &ctx.data()).await?; ctx.reply(format!("{} **{}** token(s).", if user.id == ctx.author().id { "You have".to_string() } else { format!("{} has", user.name) - }, *wealth)).await?; + }, wealth)).await?; Ok(()) } diff --git a/src/commands/gambling/give.rs b/src/commands/gambling/give.rs index 63cc1d3..04acef0 100644 --- a/src/commands/gambling/give.rs +++ b/src/commands/gambling/give.rs @@ -1,30 +1,29 @@ use crate::{Context, Error}; -use super::get_user_wealth_mut; use poise::serenity_prelude as serenity; /// Generously donate your tokens to someone else #[poise::command(slash_command, prefix_command)] -pub async fn give(ctx: Context<'_>, user: serenity::User, amount: usize) -> Result<(), Error> { +pub async fn give(ctx: Context<'_>, user: serenity::User, #[min = 1] amount: i32) -> Result<(), Error> { if user.bot { - ctx.reply("Don't waste your token(s) by giving them to a bot!").await?; + ctx.reply("Don't waste your tokens by giving them to a bot!").await?; return Ok(()); } - let mut users = ctx.data().users.lock().await; - let author_wealth = get_user_wealth_mut(&mut users, ctx.author().id); + let data = ctx.data(); - if *author_wealth < amount { - ctx.reply(format!("You only have **{}** token(s) and cannot give away **{}**.", - *author_wealth, amount)).await?; - return Ok(()); + let author_balance = super::get_balance(ctx.author().id, &data).await?; + + if author_balance < amount { + ctx.reply(format!("You do not have a high enough balance (**{author_balance}**) to complete this transaction.")).await?; + } else { + let author_new_balance = author_balance - amount; + let reciever_new_balance = super::get_balance(user.id, &data).await? + amount; + + super::change_balance(user.id, reciever_new_balance, &data).await?; + super::change_balance(ctx.author().id, author_new_balance, data).await?; + + ctx.reply(format!("You've given **{}** **{}** tokens!", user.display_name(), amount)).await?; } - *author_wealth -= amount; - - let receiver_wealth = get_user_wealth_mut(&mut users, user.id); - *receiver_wealth += amount; - - ctx.reply(format!("You've given **{}** **{}** token(s).", user.name, amount)).await?; - Ok(()) } \ No newline at end of file diff --git a/src/commands/gambling/mod.rs b/src/commands/gambling/mod.rs index 5d024e0..75f4596 100644 --- a/src/commands/gambling/mod.rs +++ b/src/commands/gambling/mod.rs @@ -1,14 +1,35 @@ -use std::collections::HashMap; -use poise::serenity_prelude::UserId; pub mod balance; pub mod give; pub mod wager; -pub(self) fn get_user_wealth_mut(users: &mut HashMap, id: UserId) -> &mut usize { - if users.get(&id).is_none() { - users.insert(id, 100); - } +use crate::common::{Data, Error}; +use poise::serenity_prelude::UserId; +use sqlx::Row; - users.get_mut(&id).unwrap() +pub async fn get_balance(id: UserId, data: &Data) -> Result { + let db = &data.database; + + let row = sqlx::query("SELECT balance FROM bank WHERE id = $1") + .bind(id.get() as i64) + .fetch_one(db).await.ok(); + + let balance = if let Some(row) = row { + row.try_get("balance")? + } else { + 100 + }; + + Ok(balance) } + +pub async fn change_balance(id: UserId, balance: i32, data: &Data) -> Result<(), Error> { + let db = &data.database; + + sqlx::query("INSERT INTO bank (id, balance) VALUES ($1, $2) ON CONFLICT (id) DO UPDATE SET balance = EXCLUDED.balance") + .bind(id.get() as i64) + .bind(balance) + .execute(db).await?; + + Ok(()) +} \ No newline at end of file diff --git a/src/commands/gambling/wager.rs b/src/commands/gambling/wager.rs index 6dc1e98..fd7530c 100644 --- a/src/commands/gambling/wager.rs +++ b/src/commands/gambling/wager.rs @@ -1,32 +1,32 @@ use crate::common::{Context, Error}; -use super::get_user_wealth_mut; /// Put forward an amount of tokens to either lose or earn #[poise::command(slash_command, prefix_command)] pub async fn wager( ctx: Context<'_>, #[min = 1] - amount: usize) -> Result<(), Error> + amount: i32) -> Result<(), Error> { - let mut users = ctx.data().users.lock().await; + let data = ctx.data(); + let mut wealth = super::get_balance(ctx.author().id, &data).await?; - let wealth = get_user_wealth_mut(&mut users, ctx.author().id); - - if *wealth < amount { + if wealth < amount { ctx.reply(format!("You do not have enough tokens (**{}**) to wager this amount.", - *wealth)).await?; + wealth)).await?; return Ok(()); } if rand::random() { - *wealth += amount; + wealth += amount; ctx.reply(format!("You just gained **{}** token(s)! You now have **{}**.", - amount, *wealth)).await?; + amount, wealth)).await?; } else { - *wealth -= amount; + wealth -= amount; ctx.reply(format!("You've lost **{}** token(s), you now have **{}**.", - amount, *wealth)).await?; + amount, wealth)).await?; } + super::change_balance(ctx.author().id, wealth, &data).await?; + Ok(()) } \ No newline at end of file diff --git a/src/common.rs b/src/common.rs index 1950350..8a60d90 100644 --- a/src/common.rs +++ b/src/common.rs @@ -2,9 +2,10 @@ use std::sync::Arc; use tokio::sync::Mutex; use std::collections::HashMap; use poise::serenity_prelude::UserId; +use sqlx::{Pool, Postgres}; pub struct Data { - pub users: Arc>>, + pub database: Pool, pub mentions: Arc>>, } diff --git a/src/main.rs b/src/main.rs index 2db2cfa..45384d1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,13 +7,14 @@ use crate::common::{Context, Error, Data}; use std::collections::HashMap; use std::env; use std::sync::Arc; -use std::time::{Instant, Duration}; -use poise::serenity_prelude as serenity; +use poise::serenity_prelude::{self as serenity}; use tokio::sync::Mutex; use clap::Parser; +use sqlx::postgres::PgPoolOptions; + #[derive(Parser, Debug)] struct BotArgs { /// Prefix for the bot (if unspecified, the bot will not have one) @@ -29,28 +30,9 @@ async fn event_handler( ) -> Result<(), Error> { match event { serenity::FullEvent::Message { new_message: message } => { - let mentions = ping_limit::extract_mentions(&message.content); - let mut cooldowns = data.mentions.lock().await; + if message.author.bot { return Ok(()) } - if mentions.iter() - .filter(|&&id| id != message.author.id) - .any(|mention| cooldowns.get(mention).map(|t| Instant::now().duration_since(*t) < Duration::from_secs(20)).unwrap_or(false)) - { - message.reply(ctx, "stop spamming!").await?; - - let guild = match message.guild_id { - Some(g) => g, - None => return Ok(()), - }; - - let mut member = guild.member(ctx, message.author.id).await.unwrap(); - member.disable_communication_until_datetime(ctx, - serenity::Timestamp::from_unix_timestamp(serenity::Timestamp::now().unix_timestamp() + 60 as i64).unwrap()).await?; - } - - for mention in mentions { - cooldowns.insert(mention, Instant::now()); - } + ping_limit::ping_spam_yeller(ctx, &message, data).await?; } _ => (), } @@ -64,6 +46,7 @@ async fn main() -> Result<(), Error> { let args = BotArgs::parse(); let token = env::var("DISCORD_BOT_TOKEN")?; + let database_url = env::var("DATABASE_URL")?; let intents = serenity::GatewayIntents::all(); let framework = poise::Framework::builder() @@ -84,8 +67,22 @@ async fn main() -> Result<(), Error> { .setup(|ctx, _ready, framework| { Box::pin(async move { poise::builtins::register_globally(ctx, &framework.options().commands).await?; + + let database = PgPoolOptions::new() + .max_connections(4) + .connect(&database_url).await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS bank ( + id BIGINT PRIMARY KEY, + balance INT + ) + "#, + ).execute(&database).await?; + Ok(Data { - users: Arc::new(Mutex::new(HashMap::new())), + database, mentions: Arc::new(Mutex::new(HashMap::new())), }) }) diff --git a/src/ping_limit.rs b/src/ping_limit.rs index 07c33b6..ddd99b7 100644 --- a/src/ping_limit.rs +++ b/src/ping_limit.rs @@ -1,7 +1,10 @@ +use crate::common::{Error, Data}; -use poise::serenity_prelude::*; +use poise::serenity_prelude::{self as serenity, Message, UserId}; use regex::Regex; +use std::time::{Instant, Duration}; + pub fn extract_mentions(content: &str) -> Vec { // Define the regex pattern for user mentions let re = Regex::new(r"<@(\d+)>").unwrap(); @@ -11,3 +14,30 @@ pub fn extract_mentions(content: &str) -> Vec { .filter_map(|cap| cap.get(1).map(|id| id.as_str().parse().unwrap())) .collect() } + +pub async fn ping_spam_yeller(ctx: &serenity::Context, message: &Message, data: &Data) -> Result<(), Error> { + let mentions = extract_mentions(&message.content); + let mut cooldowns = data.mentions.lock().await; + + if mentions.iter() + .filter(|&&id| id != message.author.id) + .any(|mention| cooldowns.get(mention).map(|t| Instant::now().duration_since(*t) < Duration::from_secs(20)).unwrap_or(false)) + { + message.reply_ping(ctx, "stop spamming!").await?; + + let guild = match message.guild_id { + Some(g) => g, + None => return Ok(()), + }; + + let mut member = guild.member(ctx, message.author.id).await.unwrap(); + member.disable_communication_until_datetime(ctx, + serenity::Timestamp::from_unix_timestamp(serenity::Timestamp::now().unix_timestamp() + 60i64).unwrap()).await?; + } + + for mention in mentions { + cooldowns.insert(mention, Instant::now()); + } + + Ok(()) +} \ No newline at end of file