use a database :D

This commit is contained in:
2024-12-07 00:21:52 -05:00
parent 9fd85e38fb
commit 500c6a6eee
8 changed files with 112 additions and 65 deletions

View File

@@ -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<serenity::User>) -> 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(())
}

View File

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

View File

@@ -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<UserId, usize>, 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<i32, Error> {
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(())
}

View File

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

View File

@@ -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<Mutex<HashMap<UserId, usize>>>,
pub database: Pool<Postgres>,
pub mentions: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
}

View File

@@ -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())),
})
})

View File

@@ -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<UserId> {
// 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<UserId> {
.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(())
}