use a database :D
This commit is contained in:
@@ -11,3 +11,4 @@ rand = "0.8.5"
|
||||
clap = { version = "4.5.20", features = ["derive"] }
|
||||
lamm = { git = "https://github.com/minneelyyyy/lamm", branch = "dev" }
|
||||
regex = "1.11.1"
|
||||
sqlx = { version = "0.8", features = [ "runtime-tokio-native-tls", "postgres" ] }
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
users.get_mut(&id).unwrap()
|
||||
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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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>>>,
|
||||
}
|
||||
|
||||
|
||||
45
src/main.rs
45
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())),
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user