From 5b029f3f673a4d9a2db508cf849f05d3e490397b Mon Sep 17 00:00:00 2001 From: Kevin Trocolli Date: Wed, 15 Nov 2023 22:38:33 -0500 Subject: [PATCH] add database interactions --- Cargo.lock | 29 ++---- Cargo.toml | 6 +- migrations/2023-08-28-015858_initial/down.sql | 2 + migrations/2023-08-28-015858_initial/up.sql | 1 + src/config.rs | 4 +- src/db/mod.rs | 50 ++++++---- src/db/schema.rs | 1 + src/db/user.rs | 96 +++++++++++++++++++ src/handlers.rs | 42 ++++++-- src/main.rs | 11 +++ 10 files changed, 186 insertions(+), 56 deletions(-) create mode 100644 src/db/user.rs diff --git a/Cargo.lock b/Cargo.lock index c2893f0..25861a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -277,7 +277,6 @@ dependencies = [ "chrono", "diesel_derives", "libsqlite3-sys", - "r2d2", "serde_json", "time 0.3.27", ] @@ -505,6 +504,12 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" +[[package]] +name = "hmac-sha256" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3688e69b38018fec1557254f64c8dc2cc8ec502890182f395dbb0aa997aa5735" + [[package]] name = "http" version = "0.2.9" @@ -897,17 +902,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "r2d2" -version = "0.8.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" -dependencies = [ - "log", - "parking_lot", - "scheduled-thread-pool", -] - [[package]] name = "rand" version = "0.8.5" @@ -1094,15 +1088,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" -[[package]] -name = "scheduled-thread-pool" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" -dependencies = [ - "parking_lot", -] - [[package]] name = "scopeguard" version = "1.2.0" @@ -1769,8 +1754,10 @@ dependencies = [ "chrono", "diesel", "diesel_migrations", + "hmac-sha256", "once_cell", "poise", + "rand", "serde", "serde_json", "serde_yaml", diff --git a/Cargo.toml b/Cargo.toml index a604cac..997b877 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,5 +14,7 @@ once_cell = "1.18" serde = "1.0" serde_yaml = "0.9" serde_json = "1.0" -diesel = { version = "2.1", features = ["sqlite", "serde_json", "chrono", "r2d2", "32-column-tables", "64-column-tables", "128-column-tables"] } -diesel_migrations = "2.0" \ No newline at end of file +diesel = { version = "2.1", features = ["sqlite", "serde_json", "chrono", "returning_clauses_for_sqlite_3_35"] } +diesel_migrations = "2.0" +rand = "0.8" +hmac-sha256 = "1.1" diff --git a/migrations/2023-08-28-015858_initial/down.sql b/migrations/2023-08-28-015858_initial/down.sql index d9a93fe..c0c121f 100644 --- a/migrations/2023-08-28-015858_initial/down.sql +++ b/migrations/2023-08-28-015858_initial/down.sql @@ -1 +1,3 @@ -- This file should undo anything in `up.sql` +DROP TABLE user; +DROP TABLE guild; diff --git a/migrations/2023-08-28-015858_initial/up.sql b/migrations/2023-08-28-015858_initial/up.sql index d219a00..76b8baf 100644 --- a/migrations/2023-08-28-015858_initial/up.sql +++ b/migrations/2023-08-28-015858_initial/up.sql @@ -8,6 +8,7 @@ CREATE TABLE user ( discord_id TEXT PRIMARY KEY NOT NULL, server BIGINT, fake_id BIGINT NOT NULL, + image_idx INT NOT NULL, when_generated DATETIME default CURRENT_TIMESTAMP NOT NULL, should_notify TINYINT NOT NULL DEFAULT 1, is_guild_admin TINYINT NOT NULL DEFAULT 0, diff --git a/src/config.rs b/src/config.rs index 63c7f32..1a129a3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -33,8 +33,8 @@ impl Default for CoreConfig { log_dir: "logs".to_string(), sqlite_file: "data/db.sqlite".to_string(), hash_ids: true, - channel: 0, - admins: vec![0], + channel: 0, + admins: Vec::new(), picture_dir: "data".to_string(), pictures: Vec::new(), } diff --git a/src/db/mod.rs b/src/db/mod.rs index f17c275..c403ecb 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,37 +1,45 @@ -use diesel::{Connection, sqlite::SqliteConnection, r2d2::*}; -use once_cell::sync::OnceCell; +use diesel::{Connection, sqlite::SqliteConnection}; +use hmac_sha256::Hash; -use crate::{config::CORE_CFG, logging::{LogLevel, log}}; +use crate::{config::CORE_CFG, logging::{LogLevel, log}, to_hex}; pub mod schema; +pub mod user; -pub static DB_CON_POOL: OnceCell>> = OnceCell::new(); +pub fn connect() -> Option { + let cfg = CORE_CFG.get().unwrap(); + let t = SqliteConnection::establish(&cfg.sqlite_file); + + if t.is_err() { + log(LogLevel::Fatal, "Database", "connect", format!("Failed to connect to database: {:?}", t.err()).as_str()); + return None; + } + + return Some(t.unwrap()) +} pub fn init_db() -> bool { - let cfg = CORE_CFG.get().unwrap(); - - let manager = ConnectionManager::::new(&cfg.sqlite_file); - let conn_pool_ = Pool::builder() - .test_on_check_out(true) - .build(manager); - - if conn_pool_.is_err() { - log(LogLevel::Fatal, "Database", "init_db", "Failed to connect to database:"); + let conn = connect(); + if conn.is_none() { return false; } - let conn_pool: Pool> = conn_pool_.unwrap(); - let test_res = conn_pool.get().unwrap().begin_test_transaction(); + let test_res = conn.unwrap().begin_test_transaction(); if test_res.is_err() { log(LogLevel::Fatal, "Database", "init_db", format!("Test Transaction failed: {}", test_res.err().unwrap()).as_str()); return false; } - let r_ = DB_CON_POOL.set(conn_pool); - if r_.is_err() { - log(LogLevel::Fatal, "Database", "init_db", "Failed to set connection pool"); - return false; - } - true +} + +pub fn get_hashed_uid(user_id: &u64) -> String { + if CORE_CFG.get().unwrap().hash_ids { + let mut h: Hash = Hash::new(); + h.update(user_id.to_ne_bytes()); + let h_final: [u8; 32] = h.finalize(); + return to_hex(h_final.to_vec()); + } else { + return user_id.to_string(); + } } \ No newline at end of file diff --git a/src/db/schema.rs b/src/db/schema.rs index 7826f59..9970571 100644 --- a/src/db/schema.rs +++ b/src/db/schema.rs @@ -12,6 +12,7 @@ diesel::table! { discord_id -> Text, server -> Nullable, fake_id -> BigInt, + image_idx -> Integer, when_generated -> Timestamp, should_notify -> Bool, is_guild_admin -> Bool, diff --git a/src/db/user.rs b/src/db/user.rs new file mode 100644 index 0000000..8cafddd --- /dev/null +++ b/src/db/user.rs @@ -0,0 +1,96 @@ +use diesel::{prelude::*, insert_into}; +use chrono::NaiveDateTime; +use crate::{logging::{LogLevel, log}, db::{connect, get_hashed_uid}}; + +#[derive(Insertable, Selectable, Queryable, Debug, Clone, Default)] +#[diesel(table_name = crate::db::schema::guild)] +pub struct Guild { + pub guild_id: i64, + pub channel_id: i64, +} + +#[derive(Insertable, Selectable, Queryable, Debug, Clone, Default)] +#[diesel(belongs_to(Guild))] +#[diesel(table_name = crate::db::schema::user)] +pub struct User { + pub discord_id: String, + pub server: Option, + pub fake_id: i64, + pub image_idx: i32, + pub when_generated: NaiveDateTime, + pub should_notify: bool, + pub is_guild_admin: bool, + pub is_banned: bool, + pub suspend_expires: Option, +} + +pub fn create_user(user_id: &u64, srv: Option, fid: i64, img_idx: i32) -> Option { + use crate::db::schema::user::dsl::*; + let conn = connect(); + if conn.is_none() { + return None + } + let mut conn = conn.unwrap(); + let uid: String = get_hashed_uid(user_id); + + let res = conn.transaction(|c| { + let sql = insert_into(user) + .values(( + discord_id.eq(uid), + server.eq(srv), + fake_id.eq(fid), + image_idx.eq(img_idx) + )) + .execute(c); + + match sql { + Err(e) => { + log(LogLevel::Error, "db", "create_user", format!("SQL Error: {:?}", e).as_str()); + return Err(e); + } + Ok(usr) => { + return Ok(usr) + } + } + }); + + match res { + Ok(_) => return get_user_by_id(user_id), + Err(_) => return None, + } +} + +pub fn get_user_by_id(user_id: &u64) -> Option { + use crate::db::schema::user::dsl::*; + let conn = connect(); + if conn.is_none() { + return None + } + let mut conn = conn.unwrap(); + let uid: String = get_hashed_uid(user_id); + + let res = conn.transaction(|c| { + let sql = user + .filter(discord_id.eq(uid)) + .select(User::as_select()) + .load::(c); + + match sql { + Err(e) => { + log(LogLevel::Error, "db", "get_user_by_id", format!("SQL Error: {:?}", e).as_str()); + return Err(e); + } + Ok(usr) => { + if usr.len() > 0 { + return Ok(Some(usr[0].clone())) + } + return Ok(None) + } + } + }); + + match res { + Ok(u) => return u, + Err(_) => return None, + } +} \ No newline at end of file diff --git a/src/handlers.rs b/src/handlers.rs index 97dbce2..ffe4bad 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -1,9 +1,10 @@ +use rand::Rng; use serde::ser::StdError; use poise::serenity_prelude::ChannelId; use poise::serenity_prelude::{self as serenity}; +use rand::seq::SliceRandom; -use crate::logging::{LogLevel, log}; -use crate::config::CORE_CFG; +use crate::{logging::{LogLevel, log}, config::CORE_CFG, db::user::*}; pub async fn handle_msg(ctx: &serenity::Context, msg: &serenity::Message) -> Result<(), Box<(dyn StdError + std::marker::Send + Sync + 'static)>> { if !msg.is_private() || msg.author.bot { @@ -12,13 +13,34 @@ pub async fn handle_msg(ctx: &serenity::Context, msg: &serenity::Message) -> Res log(LogLevel::Debug, "event_handlers", "handle_msg", msg.content.as_str()); - let cfg = CORE_CFG.get().unwrap(); - let new_channel = ChannelId::from(cfg.channel); - let picture = cfg.pictures.first().unwrap(); - let attachment_uri = format!("attachment://{}", picture.clone()); - let icon_file_uri = format!("{}/{}", cfg.picture_dir, picture.clone()); + let cfg: &crate::config::CoreConfig = CORE_CFG.get().unwrap(); + + let mut uusr = get_user_by_id(&msg.author.id.0); + if uusr.is_none() { + let fid = rand::thread_rng().gen_range(10000000..=99999999); + let idx = rand::thread_rng().gen_range(0..cfg.pictures.len()); + uusr = create_user(&msg.author.id.0, None, fid, idx as i32); + log(LogLevel::Info, "event_handlers", "handle_msg", format!("Create new user with fake ID {}", fid).as_str()); + + + if uusr.is_none() { + let _ = msg.reply(&ctx.http, "Failed to create ID for you, please contact a sysadmin."); + return Ok(()) + } + } + let usr = uusr.unwrap(); - let msg: Result = new_channel + let new_channel = ChannelId::from(cfg.channel); + let mut picture = cfg.pictures.choose(&mut rand::thread_rng()).unwrap(); + let idx_usize = usr.image_idx as usize; + if idx_usize < cfg.pictures.len() { + picture = &cfg.pictures[idx_usize]; + } + + let attachment_uri = format!("attachment://{}", picture); + let icon_file_uri = format!("{}/{}", cfg.picture_dir, picture); + + let _msg: Result = new_channel .send_message(&ctx.http, |m| { m.content("") .embed(|e| { @@ -27,7 +49,7 @@ pub async fn handle_msg(ctx: &serenity::Context, msg: &serenity::Message) -> Res .color(0x9834eb) .author(|a| { a.icon_url(attachment_uri); - a.name("85422818") + a.name(usr.fake_id.to_string()) }) }) .add_file(icon_file_uri.as_str()) @@ -36,6 +58,6 @@ pub async fn handle_msg(ctx: &serenity::Context, msg: &serenity::Message) -> Res Ok(()) } -pub async fn handle_guild_join(ctx: &serenity::Context, guild: &serenity::Guild, is_new: &bool) -> Result<(), Box<(dyn StdError + std::marker::Send + Sync + 'static)>> { +pub async fn handle_guild_join(_ctx: &serenity::Context, _guild: &serenity::Guild, _is_new: &bool) -> Result<(), Box<(dyn StdError + std::marker::Send + Sync + 'static)>> { Ok(()) } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 0d3a9b9..f881f46 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,17 @@ use db::init_db; mod handlers; use handlers::*; + + +pub fn to_hex(data: Vec) -> String { + let mut ret: String = "".to_string(); + + for x in data { + ret.push_str(format!("{:02X}", x).as_str()); + } + ret +} + #[tokio::main] async fn main() { load_core_cfg(&"config.yaml".to_string());