add database interactions

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",

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"
+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"

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;

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,

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

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<Pool<ConnectionManager<SqliteConnection>>> = OnceCell::new();
+pub fn connect() -> Option<SqliteConnection> {
+    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::<SqliteConnection>::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<ConnectionManager<SqliteConnection>> = 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();
+    }
 }

src/db/schema.rs

@@ -12,6 +12,7 @@ diesel::table! {
         discord_id -> Text,
         server -> Nullable<BigInt>,
         fake_id -> BigInt,
+        image_idx -> Integer,
         when_generated -> Timestamp,
         should_notify -> Bool,
         is_guild_admin -> Bool,

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<i64>,
+    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<NaiveDateTime>,
+}
+
+pub fn create_user(user_id: &u64, srv: Option<i64>, fid: i64, img_idx: i32) -> Option<User> {
+    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<User> {
+    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::<User>(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,
+    }
+}

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<serenity::Message, ::serenity::Error> = 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<serenity::Message, ::serenity::Error> = 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(()) 
 }

src/main.rs

@@ -16,6 +16,17 @@ use db::init_db;
 mod handlers;
 use handlers::*;
 
+
+
+pub fn to_hex(data: Vec<u8>) -> 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());