From c910fa4b97b776ae257ae379b41c43748bbf5204 Mon Sep 17 00:00:00 2001 From: Kevin Trocolli Date: Mon, 28 Aug 2023 01:36:26 -0400 Subject: [PATCH] add message sending --- .gitignore | 3 +- Cargo.lock | 109 ++++++++++++++++++ Cargo.toml | 4 +- diesel.toml | 9 ++ example_config.yaml | 7 +- migrations/.keep | 0 migrations/2023-08-28-015858_initial/down.sql | 1 + migrations/2023-08-28-015858_initial/up.sql | 17 +++ src/config.rs | 14 ++- src/db/mod.rs | 37 ++++++ src/db/schema.rs | 28 +++++ src/handlers.rs | 41 +++++++ src/main.rs | 22 +++- 13 files changed, 287 insertions(+), 5 deletions(-) create mode 100644 diesel.toml create mode 100644 migrations/.keep create mode 100644 migrations/2023-08-28-015858_initial/down.sql create mode 100644 migrations/2023-08-28-015858_initial/up.sql create mode 100644 src/db/mod.rs create mode 100644 src/db/schema.rs create mode 100644 src/handlers.rs diff --git a/.gitignore b/.gitignore index 8f1a83e..ebca782 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ *.yaml !example_config.yaml /data -*.log \ No newline at end of file +*.log +*.lib \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index c0d3066..c2893f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -274,8 +274,11 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98235fdc2f355d330a8244184ab6b4b33c28679c0b4158f63138e51d6cf7e88" dependencies = [ + "chrono", "diesel_derives", "libsqlite3-sys", + "r2d2", + "serde_json", "time 0.3.27", ] @@ -291,6 +294,17 @@ dependencies = [ "syn 2.0.29", ] +[[package]] +name = "diesel_migrations" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6036b3f0120c5961381b570ee20a02432d7e2d27ea60de9578799cf9156914ac" +dependencies = [ + "diesel", + "migrations_internals", + "migrations_macros", +] + [[package]] name = "diesel_table_macro_syntax" version = "0.1.0" @@ -681,6 +695,27 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" +[[package]] +name = "migrations_internals" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f23f71580015254b020e856feac3df5878c2c7a8812297edd6c0a485ac9dada" +dependencies = [ + "serde", + "toml", +] + +[[package]] +name = "migrations_macros" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cce3325ac70e67bbab5bd837a31cae01f1a6db64e0e744a33cb03a543469ef08" +dependencies = [ + "migrations_internals", + "proc-macro2", + "quote", +] + [[package]] name = "mime" version = "0.3.17" @@ -862,6 +897,17 @@ 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" @@ -1048,6 +1094,15 @@ 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" @@ -1105,6 +1160,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96426c9936fd7a0124915f9185ea1d20aa9445cc9821142f0a73bc9207a2e186" +dependencies = [ + "serde", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1379,6 +1443,40 @@ dependencies = [ "tracing", ] +[[package]] +name = "toml" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c17e963a819c331dcacd7ab957d80bc2b9a9c1e71c804826d2f283dd65306542" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.19.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8123f27e969974a3dfba720fdb560be359f57b44302d280ba72e76a74480e8a" +dependencies = [ + "indexmap 2.0.0", + "serde", + "serde_spanned", + "toml_datetime", + "winnow", +] + [[package]] name = "tower-service" version = "0.3.2" @@ -1670,9 +1768,11 @@ version = "0.1.0" dependencies = [ "chrono", "diesel", + "diesel_migrations", "once_cell", "poise", "serde", + "serde_json", "serde_yaml", "serenity", "tokio", @@ -1775,6 +1875,15 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +[[package]] +name = "winnow" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c2e3184b9c4e92ad5167ca73039d0c42476302ab603e2fec4487511f38ccefc" +dependencies = [ + "memchr", +] + [[package]] name = "winreg" version = "0.50.0" diff --git a/Cargo.toml b/Cargo.toml index c6d2d24..a604cac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,4 +13,6 @@ chrono = "0.4" once_cell = "1.18" serde = "1.0" serde_yaml = "0.9" -diesel = { version = "2.1", features = ["sqlite"] } \ No newline at end of file +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 diff --git a/diesel.toml b/diesel.toml new file mode 100644 index 0000000..fa8cc69 --- /dev/null +++ b/diesel.toml @@ -0,0 +1,9 @@ +# For documentation on how to configure this file, +# see https://diesel.rs/guides/configuring-diesel-cli + +[print_schema] +file = "src/db/schema.rs" +custom_type_derives = ["diesel::query_builder::QueryId"] + +[migrations_directory] +dir = "migrations" diff --git a/example_config.yaml b/example_config.yaml index 2601023..725524b 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -1,6 +1,11 @@ token: "" +log_level: "info" log_dir: "logs" +sqlite_file: "data/db.sqlite" hash_ids: true channel: 0 admins: - - 0 \ No newline at end of file + - 0 +picture_dir: "data" +pictures: + - "" \ No newline at end of file diff --git a/migrations/.keep b/migrations/.keep new file mode 100644 index 0000000..e69de29 diff --git a/migrations/2023-08-28-015858_initial/down.sql b/migrations/2023-08-28-015858_initial/down.sql new file mode 100644 index 0000000..d9a93fe --- /dev/null +++ b/migrations/2023-08-28-015858_initial/down.sql @@ -0,0 +1 @@ +-- This file should undo anything in `up.sql` diff --git a/migrations/2023-08-28-015858_initial/up.sql b/migrations/2023-08-28-015858_initial/up.sql new file mode 100644 index 0000000..d219a00 --- /dev/null +++ b/migrations/2023-08-28-015858_initial/up.sql @@ -0,0 +1,17 @@ +-- Your SQL goes here +CREATE TABLE guild ( + guild_id BIGINT PRIMARY KEY NOT NULL, + channel_id BIGINT NOT NULL +); + +CREATE TABLE user ( + discord_id TEXT PRIMARY KEY NOT NULL, + server BIGINT, + fake_id BIGINT 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, + is_banned TINYINT NOT NULL DEFAULT 0, + suspend_expires DATETIME, + FOREIGN KEY(server) REFERENCES guild(guild_id) +); diff --git a/src/config.rs b/src/config.rs index 886c241..63c7f32 100644 --- a/src/config.rs +++ b/src/config.rs @@ -7,24 +7,36 @@ use serde_yaml::{self}; pub struct CoreConfig { /// Discord bot token. pub token: String, + /// Log Level + pub log_level: String, /// Folder to store program logs. pub log_dir: String, + /// Location for where the sqlite database should be stored + pub sqlite_file: String, /// If user IDs should be hased in the database, disables @mentioning. pub hash_ids: bool, /// Channel where messages should be posted. pub channel: u64, /// List of user IDs of users who have admin perms with the bot. pub admins: Vec, + /// Directory where pictures are kept. + pub picture_dir: String, + /// List of pictures to acompany users. + pub pictures: Vec, } impl Default for CoreConfig { fn default() -> Self { Self { token: Default::default(), + log_level: "info".to_string(), log_dir: "logs".to_string(), + sqlite_file: "data/db.sqlite".to_string(), hash_ids: true, channel: 0, - admins: vec![0] + admins: vec![0], + picture_dir: "data".to_string(), + pictures: Vec::new(), } } } diff --git a/src/db/mod.rs b/src/db/mod.rs new file mode 100644 index 0000000..f17c275 --- /dev/null +++ b/src/db/mod.rs @@ -0,0 +1,37 @@ +use diesel::{Connection, sqlite::SqliteConnection, r2d2::*}; +use once_cell::sync::OnceCell; + +use crate::{config::CORE_CFG, logging::{LogLevel, log}}; + +pub mod schema; + +pub static DB_CON_POOL: OnceCell>> = OnceCell::new(); + +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:"); + return false; + } + + let conn_pool: Pool> = conn_pool_.unwrap(); + let test_res = conn_pool.get().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 +} \ No newline at end of file diff --git a/src/db/schema.rs b/src/db/schema.rs new file mode 100644 index 0000000..7826f59 --- /dev/null +++ b/src/db/schema.rs @@ -0,0 +1,28 @@ +// @generated automatically by Diesel CLI. + +diesel::table! { + guild (guild_id) { + guild_id -> BigInt, + channel_id -> BigInt, + } +} + +diesel::table! { + user (discord_id) { + discord_id -> Text, + server -> Nullable, + fake_id -> BigInt, + when_generated -> Timestamp, + should_notify -> Bool, + is_guild_admin -> Bool, + is_banned -> Bool, + suspend_expires -> Nullable, + } +} + +diesel::joinable!(user -> guild (server)); + +diesel::allow_tables_to_appear_in_same_query!( + guild, + user, +); diff --git a/src/handlers.rs b/src/handlers.rs new file mode 100644 index 0000000..97dbce2 --- /dev/null +++ b/src/handlers.rs @@ -0,0 +1,41 @@ +use serde::ser::StdError; +use poise::serenity_prelude::ChannelId; +use poise::serenity_prelude::{self as serenity}; + +use crate::logging::{LogLevel, log}; +use crate::config::CORE_CFG; + +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 { + return Ok(()); + } + + 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 msg: Result = new_channel + .send_message(&ctx.http, |m| { + m.content("") + .embed(|e| { + e.title("") + .description(&msg.content) + .color(0x9834eb) + .author(|a| { + a.icon_url(attachment_uri); + a.name("85422818") + }) + }) + .add_file(icon_file_uri.as_str()) + }) + .await; + 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)>> { + Ok(()) +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index d875fb4..0d3a9b9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ use poise::serenity_prelude::{self as serenity}; +use std::str::FromStr; mod config; use config::{load_core_cfg, CORE_CFG}; @@ -9,15 +10,34 @@ use logging::{LogLevel, init_logger, log}; mod commands; use commands::{notify::*, Data}; +mod db; +use db::init_db; + +mod handlers; +use handlers::*; + #[tokio::main] async fn main() { load_core_cfg(&"config.yaml".to_string()); let cfg = CORE_CFG.get().unwrap(); - init_logger(&cfg.log_dir, LogLevel::Info); + init_logger(&cfg.log_dir, LogLevel::from_str(cfg.log_level.as_str()).unwrap()); + if !init_db() { + return; + } let framework = poise::Framework::builder() .options(poise::FrameworkOptions { + event_handler: |ctx: &serenity::Context, event: &poise::Event<'_>, _framework, _data: &Data| { + Box::pin(async move { + log(LogLevel::Trace, "main", "event_handler", format!("Event Fired: {:?}", event.name()).as_str()); + match event { + poise::Event::GuildCreate { guild, is_new } => handle_guild_join(ctx, guild, is_new).await, + poise::Event::Message { new_message } => handle_msg(ctx, new_message).await, + _ => Ok(()), + } + }) + }, commands: vec![notify()], ..Default::default() })