Initial commit

.gitignore

@@ -0,0 +1,3 @@
+/target
+Cargo.lock
+.env

Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "tg-dev-srv-bot"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+dotenv = "0.15.0"
+exmex = "0.21.0"
+f128 = "0.2.9"
+poise = "0.6.2"
+rug = "1.30.0"
+serenity = "0.12.5"
+sqlx = { version = "0.9.0", features = ["postgres", "runtime-tokio", "macros"] }
+tokio = { version = "1.52.3", features = ["full"] }

docker-compose-dev.yml

@@ -0,0 +1,22 @@
+services:
+  postgres:
+    image: postgres:16-alpine
+    container_name: tg-dev-srv-bot-postgres
+    restart: always
+    environment:
+      POSTGRES_USER: ${POSTGRES_USER:-postgres}
+      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-postgres}
+      POSTGRES_DB: ${POSTGRES_DB:-tg_dev_srv_bot}
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER:-postgres} -d $${POSTGRES_DB:-tg_dev_srv_bot}"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+volumes:
+  postgres_data:
+    driver: local

sqlschema.txt

@@ -0,0 +1,23 @@
+CREATE TABLE IF NOT EXISTS auto_roles (
+  guild_id BIGINT NOT NULL,
+  role_id BIGINT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS reaction_role (
+    guild_id BIGINT NOT NULL,
+    channel_id BIGINT NOT NULL,
+    message_id BIGINT NOT NULL,
+    role_id BIGINT NOT NULL,
+    emoji_id BIGINT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS admin_users (
+    guild_id BIGINT NOT NULL,
+    user_id BIGINT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS admin_roles (
+    guild_id BIGINT NOT NULL,
+    role_id BIGINT NOT NULL
+);
+

src/commands/calc.rs

@@ -0,0 +1,30 @@
+use crate::{Context, Error};
+use poise::ChoiceParameter;
+
+#[derive(Debug, poise::ChoiceParameter)]
+pub enum Precision {
+    #[name = "32-bit"]
+    Bits32,
+    #[name = "64-bit"]
+    Bits64,
+}
+
+/// Calculator
+#[poise::command(slash_command, prefix_command)]
+pub async fn calc(ctx: Context<'_>,
+    #[description = "Expression"]
+    expression: String,
+    #[description = "Precision"]
+    precision: Option<Precision>
+) -> Result<(), Error> {
+    let start_time = std::time::Instant::now();
+    let value = match precision {
+        Some(Precision::Bits32) => exmex::eval_str::<f32>(&expression).map_err(|err| err.to_string())?.to_string(),
+        Some(Precision::Bits64) => exmex::eval_str::<f64>(&expression).map_err(|err| err.to_string())?.to_string(),
+        None => exmex::eval_str::<f64>(&expression).map_err(|err| err.to_string())?.to_string(),
+    };
+    let end_time = std::time::Instant::now();
+    let duration = end_time - start_time;
+    ctx.say(format!("{} = {} \n -# Precision: {} Compute time : {:?}", expression, value, precision.map(|p| p.name()).unwrap_or("64-bit"), duration)).await?;
+    Ok(())
+}

src/commands/manage_users.rs

@@ -0,0 +1,265 @@
+use crate::{Context, Error};
+use poise::serenity_prelude as sere;
+
+#[derive(Debug, poise::ChoiceParameter)]
+pub enum AdminOperationType {
+    #[name = "Add user"]
+    AddUser,
+    #[name = "Remove user"]
+    RemoveUser,
+    #[name = "List users"]
+    ListUsers,
+    #[name = "Add role"]
+    AddRole,
+    #[name = "Remove role"]
+    RemoveRole,
+    #[name = "List roles"]
+    ListRoles,
+}
+
+/// Admin command
+#[poise::command(slash_command, prefix_command)]
+pub async fn manage_users(ctx: Context<'_>,
+    #[description = "Operation to perform"]
+    operation: Option<AdminOperationType>,
+    #[description = "User"]
+    user: Option<sere::Member>,
+    #[description = "Role"]
+    role: Option<sere::Role>
+) -> Result<(), Error> {
+    let response_message = ctx.say("Running command...").await?;
+    
+    if !is_admin(ctx).await? {
+        response_message.edit(ctx, poise::CreateReply::default().content("You are not allowed to run this command.")).await?;
+        return Ok(());
+    }
+
+
+    
+    match operation {
+        Some(AdminOperationType::AddUser) => {
+            if let Some(s_user) = user {
+                add_admin_user(ctx, &response_message, s_user).await?;
+            } else {
+                edit_response_message(&response_message, ctx, "Please provide a user".to_string(), false).await?;
+            }
+        }
+        Some(AdminOperationType::RemoveUser) => {
+            if let Some(s_user) = user {
+                remove_admin_user(ctx, &response_message, s_user).await?;
+            } else {
+                edit_response_message(&response_message, ctx, "Please provide a user".to_string(), false).await?;
+            }
+        }
+        Some(AdminOperationType::ListUsers) => {
+            list_admin_users(ctx, &response_message).await?;
+        }
+        Some(AdminOperationType::AddRole) => {
+            if let Some(s_role) = role {
+                add_admin_role(ctx, &response_message, s_role).await?;
+            } else {
+                edit_response_message(&response_message, ctx, "Please provide a role".to_string(), false).await?;
+            }
+        }
+        Some(AdminOperationType::RemoveRole) => {
+            if let Some(s_role) = role {
+                remove_admin_role(ctx, &response_message, s_role).await?;
+            } else {
+                edit_response_message(&response_message, ctx, "Please provide a role".to_string(), false).await?;
+            }
+        }
+        Some(AdminOperationType::ListRoles) => {
+            list_admin_roles(ctx, &response_message).await?;
+        }
+        None => {
+            edit_response_message(&response_message, ctx, "Please provide an operation".to_string(), false).await?;
+        }
+    }
+    Ok(())
+}
+
+async fn get_all_admin_users(ctx: Context<'_>) -> Result<Vec<sere::User>, Error> {
+    let guild_id = ctx.guild_id().unwrap();
+    let pool = &ctx.data().pool;
+    let db_guild_id = guild_id.get() as i64;
+
+    let admin_users_records = sqlx::query!("SELECT user_id FROM admin_users WHERE guild_id = $1;", db_guild_id)
+        .fetch_all(pool)
+        .await?;
+
+    let admin_roles_records = sqlx::query!("SELECT role_id FROM admin_roles WHERE guild_id = $1;", db_guild_id)
+        .fetch_all(pool)
+        .await?;
+
+    let admin_role_ids: std::collections::HashSet<_> = admin_roles_records
+        .into_iter()
+        .map(|record| sere::RoleId::new(record.role_id as u64))
+        .collect();
+
+    let mut admin_users = Vec::new();
+    let mut admin_user_ids = std::collections::HashSet::new();
+
+    for record in admin_users_records {
+        let user_id = sere::UserId::new(record.user_id as u64);
+        if admin_user_ids.insert(user_id) {
+            if let Ok(user) = user_id.to_user(ctx.http()).await {
+                admin_users.push(user);
+            }
+        }
+    }
+
+    let mut last_member_id = None;
+    loop {
+        let members = guild_id.members(ctx.http(), Some(1000), last_member_id).await?;
+        if members.is_empty() {
+            break;
+        }
+        last_member_id = Some(members.last().unwrap().user.id);
+
+        for member in members {
+            let has_role = member.roles.iter().any(|role_id| admin_role_ids.contains(role_id));
+            let has_admin_perms = ctx.guild()
+                .map(|g| g.member_permissions(&member).administrator())
+                .unwrap_or(false);
+            
+            if (has_role || has_admin_perms) && admin_user_ids.insert(member.user.id) {
+                admin_users.push(member.user);
+            }
+        }
+    }
+
+    Ok(admin_users)
+}
+
+async fn is_admin(ctx: Context<'_>) -> Result<bool, Error> {
+    let user_running_command = ctx.author().id.get() as i64;
+    let admins = get_all_admin_users(ctx).await?;
+    
+    for user in admins {
+        if user.id.get() as i64 == user_running_command {
+            return Ok(true);
+        }
+    }
+    Ok(false)
+}
+
+async fn add_admin_user<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>, user: sere::Member) -> Result<(), Error> {
+    if is_target_admin_user(ctx, &user).await? {
+        edit_response_message(msg, ctx, format!("{} is already an admin", user.user.name), false).await?;
+        return Ok(());
+    }
+    let query = sqlx::query!("INSERT INTO admin_users (guild_id, user_id) VALUES ($1, $2);", ctx.guild_id().unwrap().get() as i64, user.user.id.get() as i64).execute(&ctx.data().pool).await;
+    match query {
+        Ok(_) => {
+            edit_response_message(msg, ctx, format!("Added {} as admin", user.user.name), false).await?;
+        }
+        Err(e) => {
+            edit_response_message(msg, ctx, format!("Failed to add {} as admin: {}", user.user.name, e), false).await?;
+        }
+    }
+    Ok(())
+}
+
+async fn edit_response_message<'a>(response_message: &poise::ReplyHandle<'a>, ctx: Context<'_>, content: String, silent: bool) -> Result<(), Error> {
+    if silent {
+        response_message.edit(ctx, poise::CreateReply::default().content(content).allowed_mentions(serenity::all::CreateAllowedMentions::new().empty_users())).await?;
+    } else {
+        response_message.edit(ctx, poise::CreateReply::default().content(content)).await?;
+    }
+    Ok(())
+}
+
+async fn is_target_admin_user(ctx: Context<'_>, user: &sere::Member) -> Result<bool, Error> {
+    let admins = get_all_admin_users(ctx).await?;
+    for admin in admins {
+        if admin.id.get() as i64 == user.user.id.get() as i64 {
+            return Ok(true);
+        }
+    }
+    Ok(false)
+}
+
+
+async fn remove_admin_user<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>, user: sere::Member) -> Result<(), Error> {
+    if !is_target_admin_user(ctx, &user).await? {
+        edit_response_message(msg, ctx, format!("{} is not an admin", user.user.name), false).await?;
+        return Ok(());
+    }
+    let query = sqlx::query!("DELETE FROM admin_users WHERE guild_id = $1 AND user_id = $2;", ctx.guild_id().unwrap().get() as i64, user.user.id.get() as i64).execute(&ctx.data().pool).await;
+    match query {
+        Ok(_) => {
+            edit_response_message(msg, ctx, format!("Removed {} as admin", user.user.name), false).await?;
+        }
+        Err(e) => {
+            edit_response_message(msg, ctx, format!("Failed to remove {} as admin: {}", user.user.name, e), false).await?;
+        }
+    }
+    Ok(())
+}
+
+async fn list_admin_users<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>) -> Result<(), Error> {
+    let admins = get_all_admin_users(ctx).await?;
+    let admin_strings: Vec<String> = admins.iter().map(|user| format!("<@{}>", user.id.get())).collect();
+    edit_response_message(msg, ctx, format!("Admin users: {}", admin_strings.join(", ")), true).await?;
+    Ok(())
+}
+
+async fn add_admin_role<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>, role: sere::Role) -> Result<(), Error> {
+    if is_target_admin_role(ctx, &role).await? {
+        edit_response_message(msg, ctx, format!("{} is already an admin role", role.name), false).await?;
+        return Ok(());
+    }
+    let query = sqlx::query!("INSERT INTO admin_roles (guild_id, role_id) VALUES ($1, $2);", ctx.guild_id().unwrap().get() as i64, role.id.get() as i64).execute(&ctx.data().pool).await;
+    match query {
+        Ok(_) => {
+            edit_response_message(msg, ctx, format!("Added {} as admin role", role.name), false).await?;
+        }
+        Err(e) => {
+            edit_response_message(msg, ctx, format!("Failed to add {} as admin role: {}", role.name, e), false).await?;
+        }
+    }
+    Ok(())
+}
+
+async fn remove_admin_role<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>, role: sere::Role) -> Result<(), Error> {
+    if !is_target_admin_role(ctx, &role).await? {
+        edit_response_message(msg, ctx, format!("{} is not an admin role", role.name), false).await?;
+        return Ok(());
+    }
+    let query = sqlx::query!("DELETE FROM admin_roles WHERE guild_id = $1 AND role_id = $2;", ctx.guild_id().unwrap().get() as i64, role.id.get() as i64).execute(&ctx.data().pool).await;
+    match query {
+        Ok(_) => {
+            edit_response_message(msg, ctx, format!("Removed {} as admin role", role.name), false).await?;
+        }
+        Err(e) => {
+            edit_response_message(msg, ctx, format!("Failed to remove {} as admin role: {}", role.name, e), false).await?;
+        }
+    }
+    Ok(())
+}
+
+async fn list_admin_roles<'a>(ctx: Context<'a>, msg: &poise::ReplyHandle<'a>) -> Result<(), Error> {
+    let roles = sqlx::query!("SELECT role_id FROM admin_roles WHERE guild_id = $1", ctx.guild_id().unwrap().get() as i64).fetch_all(&ctx.data().pool).await?;
+    let role_strings: Vec<String> = roles.iter().map(|role| format!("<@&{}>", role.role_id)).collect();
+    edit_response_message(msg, ctx, format!("Admin roles: {}", role_strings.join(", ")), true).await?;
+    Ok(())
+}
+
+async fn get_admin_roles(ctx: Context<'_>) -> Result<Vec<sere::Role>, Error> {
+    let roles = sqlx::query!("SELECT role_id FROM admin_roles WHERE guild_id = $1", ctx.guild_id().unwrap().get() as i64).fetch_all(&ctx.data().pool).await?;
+    let mut admin_roles = Vec::new();
+    for role in roles {
+        admin_roles.push(ctx.guild().unwrap().roles.get(&sere::RoleId::new(role.role_id as u64)).unwrap().clone());
+    }
+    Ok(admin_roles)
+}
+
+async fn is_target_admin_role<'a>(ctx: Context<'a>, role: &sere::Role) -> Result<bool, Error> {
+    let roles = get_admin_roles(ctx).await?;
+    for r in roles {
+        if r.id.get() as i64 == role.id.get() as i64 {
+            return Ok(true);
+        }
+    }
+    Ok(false)
+}

src/commands/mod.rs

@@ -0,0 +1,7 @@
+pub mod ping;
+pub mod calc;
+pub mod manage_users;
+
+pub use ping::ping;
+pub use calc::calc;
+pub use manage_users::manage_users;

src/commands/ping.rs

@@ -0,0 +1,14 @@
+use crate::{Context, Error};
+use poise::serenity_prelude as sere;
+
+/// Get the bot's current ping (back and forth)
+#[poise::command(slash_command, prefix_command)]
+pub async fn ping(ctx: Context<'_>) -> Result<(), Error> {
+    let command_time_in_microseconds = ctx.created_at().timestamp_micros();
+    let current_time_in_microseconds = sere::Timestamp::now().timestamp_micros();
+    let ping = (current_time_in_microseconds - command_time_in_microseconds) / 500;
+    // We divide by 1000 since we get the time in microseconds, and then multiply by 2 to get the roundtrip time
+    // This is arguably not the best way to calculate ping, since it assumes perfect clock accuracy, but I'm lazy
+    ctx.say(format!("Current ping: {} ms", ping)).await?;
+    Ok(())
+}

src/main.rs

@@ -0,0 +1,68 @@
+use std::env;
+use poise::serenity_prelude as sere;
+pub mod commands;
+
+pub struct Data {
+    pub pool: sqlx::PgPool,
+}
+
+// Define standard types for Poise context, commands, and errors.
+pub type Error = Box<dyn std::error::Error + Send + Sync>;
+pub type Context<'a> = poise::Context<'a, Data, Error>;
+
+
+
+#[tokio::main]
+async fn main() {
+    // Load env from file
+    dotenv::dotenv().ok();
+
+    // Grab envs
+    let database_url = env::var("DATABASE_URL").expect("Expected a database url in the environment");
+    let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");
+
+    // Initialize database
+    let pool = sqlx::postgres::PgPoolOptions::new().connect(&database_url).await.expect("Failed to connect to database");
+
+    let _res = sqlx::raw_sql(include_str!("../sqlschema.txt")).execute(&pool).await.expect("Failed to execute sqlschema");
+
+    let intents = sere::GatewayIntents::non_privileged()
+        | sere::GatewayIntents::MESSAGE_CONTENT;
+
+    let framework = poise::Framework::builder()
+        .options(poise::FrameworkOptions {
+            event_handler: |ctx, event, framework, user_data| {
+                Box::pin(event_handler(ctx, event, framework, user_data))
+            },
+            commands: vec![commands::ping(), commands::calc(), commands::manage_users()],
+            ..Default::default()
+        })
+        .setup(|ctx, _ready, framework| {
+            Box::pin(async move {
+                let guild_id = sere::GuildId::new(1522685547641045002);
+                //poise::builtins::register_globally(ctx, &framework.options().commands).await?;
+                poise::builtins::register_in_guild(ctx, &framework.options().commands, guild_id).await?;
+                Ok(Data { pool: pool.clone() })
+            })
+        })
+        .build();
+
+    let mut client = sere::ClientBuilder::new(&token, intents).framework(framework).await.expect("Err creating client");
+
+
+    if let Err(why) = client.start().await {
+        println!("Client error: {why:?}");
+    }
+}
+
+async fn event_handler(
+    ctx: &sere::Context,
+    event: &sere::FullEvent,
+    framework: poise::FrameworkContext<'_, Data, Error>,
+    user_data: &Data,
+) -> Result<(), Error> {
+    if let sere::FullEvent::ChannelCreate {channel} = event {
+        println!("Channel created: {channel:?}");
+    }
+    Ok(())
+}