diff --git a/.rustfmt.toml b/.rustfmt.toml new file mode 100644 index 0000000..b8ae8dd --- /dev/null +++ b/.rustfmt.toml @@ -0,0 +1,9 @@ +max_width = 150 +single_line_if_else_max_width = 100 +chain_width = 100 +fn_params_layout = "Compressed" +#control_brace_style = "ClosingNextLine" +#brace_style = "PreferSameLine" +struct_lit_width = 0 +tab_spaces = 2 +use_small_heuristics = "Max" \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index a20deee..6b3936a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,280 +1,292 @@ -use std::env; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; use dotenvy::dotenv; -use serenity::async_trait; -use serenity::model::gateway::{GatewayIntents, Ready}; -use serenity::model::guild::Member; -use serenity::model::id::GuildId; -use serenity::prelude::*; -use serenity::model::prelude::RoleId; +use serenity::{ + async_trait, + client::{Context, EventHandler}, + model::{ + gateway::{GatewayIntents, Ready}, + guild::Member, + id::GuildId, + prelude::RoleId, + }, + prelude::TypeMapKey, + Client, +}; +use std::{ + env, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Duration, +}; +use tokio::sync::RwLock; struct Handler { - is_loop_running: AtomicBool, + is_loop_running: AtomicBool, } #[async_trait] impl EventHandler for Handler { - async fn guild_member_addition(&self, ctx: Context, mut new_member: Member) { - let config_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Config in TypeMap.").clone() - }; - let config = config_lock.read().await; - - let members_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Members in TypeMap.").clone() - }; - let members = members_lock.read().await; + async fn guild_member_addition(&self, ctx: Context, mut new_member: Member) { + let config_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Config in TypeMap.").clone() + }; + let config = config_lock.read().await; - if members.contains(&new_member.user.name) { - let mut roles = vec![]; + let members_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Members in TypeMap.").clone() + }; + let members = members_lock.read().await; - if !new_member.roles.contains(&config.member_role_past){ - roles.push(config.member_role_past); - } - if !new_member.roles.contains(&config.member_role_current){ - roles.push(config.member_role_current); - } + if members.contains(&new_member.user.name) { + let mut roles = vec![]; - if let Err(e) = new_member.add_roles(&ctx, &roles).await { - println!("{:?}", e); - } - } + if !new_member.roles.contains(&config.member_role_past) { + roles.push(config.member_role_past); + } + if !new_member.roles.contains(&config.member_role_current) { + roles.push(config.member_role_current); + } + + if let Err(e) = new_member.add_roles(&ctx, &roles).await { + println!("{:?}", e); + } } + } - async fn ready(&self, ctx: Context, ready: Ready) { - let ctx = Arc::new(ctx); - println!("{} is connected!", ready.user.name); - - let config_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Config in TypeMap.").clone() - }; - let config = config_lock.read().await; - let timing_update = config.timing_update; - let timing_fetch = config.timing_fetch; + async fn ready(&self, ctx: Context, ready: Ready) { + let ctx = Arc::new(ctx); + println!("{} is connected!", ready.user.name); - if !self.is_loop_running.load(Ordering::Relaxed) { - // We have to clone the Arc, as it gets moved into the new thread. - let ctx1 = Arc::clone(&ctx); - // tokio::spawn creates a new green thread that can run in parallel with the rest of - // the application. - tokio::spawn(async move { - loop { - // We clone Context again here, because Arc is owned, so it moves to the - // new function. - bulk_check(Arc::clone(&ctx1)).await; - tokio::time::sleep(Duration::from_secs(timing_update)).await; - } - }); - - let ctx2 = Arc::clone(&ctx); - tokio::spawn(async move { - loop { - fetch_accounts(Arc::clone(&ctx2)).await; - tokio::time::sleep(Duration::from_secs(timing_fetch)).await; - } - }); + let config_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Config in TypeMap.").clone() + }; + let config = config_lock.read().await; + let timing_update = config.timing_update; + let timing_fetch = config.timing_fetch; - // Now that the loop is running, we set the bool to true - self.is_loop_running.swap(true, Ordering::Relaxed); + if !self.is_loop_running.load(Ordering::Relaxed) { + // We have to clone the Arc, as it gets moved into the new thread. + let ctx1 = Arc::clone(&ctx); + // tokio::spawn creates a new green thread that can run in parallel with the rest of + // the application. + tokio::spawn(async move { + loop { + // We clone Context again here, because Arc is owned, so it moves to the + // new function. + bulk_check(Arc::clone(&ctx1)).await; + tokio::time::sleep(Duration::from_secs(timing_update)).await; } + }); + + let ctx2 = Arc::clone(&ctx); + tokio::spawn(async move { + loop { + fetch_accounts(Arc::clone(&ctx2)).await; + tokio::time::sleep(Duration::from_secs(timing_fetch)).await; + } + }); + + // Now that the loop is running, we set the bool to true + self.is_loop_running.swap(true, Ordering::Relaxed); } + } } #[derive(Default, Debug)] struct MembersCount { - members: i32, - members_current: i32, + members: i32, + members_current: i32, } struct MemberCounter; impl TypeMapKey for MemberCounter { - type Value = Arc>; + type Value = Arc>; } struct Members; impl TypeMapKey for Members { - type Value = Arc>>; + type Value = Arc>>; } -async fn bulk_check(ctx: Arc){ - let config_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Config in TypeMap.").clone() - }; - let config = config_lock.read().await; - - let members_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Members in TypeMap.").clone() - }; - let members = members_lock.read().await; - - let mut roles_set = [0,0,0]; - let mut res = MembersCount { - members: 0, - members_current: 0, - }; - if let Ok(x) = config.server.members(&ctx, None, None).await { - for mut member in x { - if members.contains(&member.user.name) { - let mut roles = vec![]; +async fn bulk_check(ctx: Arc) { + let config_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Config in TypeMap.").clone() + }; + let config = config_lock.read().await; - if !member.roles.contains(&config.member_role_past){ - roles_set[0] += 1; - roles.push(config.member_role_past); - } - if !member.roles.contains(&config.member_role_current){ - roles_set[1] += 1; - roles.push(config.member_role_current); - } + let members_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Members in TypeMap.").clone() + }; + let members = members_lock.read().await; - if let Err(e) = member.add_roles(&ctx, &roles).await { - println!("{:?}", e); - } - } else if member.roles.contains(&config.member_role_current) { - roles_set[2] += 1; - // if theya re not a current member and have the role then remove it - if let Err(e) = member.remove_role(&ctx, &config.member_role_current).await { - println!("{:?}", e); - } - } + let mut roles_set = [0, 0, 0]; + let mut res = MembersCount { + members: 0, + members_current: 0, + }; + if let Ok(x) = config.server.members(&ctx, None, None).await { + for mut member in x { + if members.contains(&member.user.name) { + let mut roles = vec![]; - if member.roles.contains(&config.member_role_past){ - res.members += 1; - } - if member.roles.contains(&config.member_role_current){ - res.members_current += 1; - } + if !member.roles.contains(&config.member_role_past) { + roles_set[0] += 1; + roles.push(config.member_role_past); } + if !member.roles.contains(&config.member_role_current) { + roles_set[1] += 1; + roles.push(config.member_role_current); + } + + if let Err(e) = member.add_roles(&ctx, &roles).await { + println!("{:?}", e); + } + } else if member.roles.contains(&config.member_role_current) { + roles_set[2] += 1; + // if theya re not a current member and have the role then remove it + if let Err(e) = member.remove_role(&ctx, &config.member_role_current).await { + println!("{:?}", e); + } + } + + if member.roles.contains(&config.member_role_past) { + res.members += 1; + } + if member.roles.contains(&config.member_role_current) { + res.members_current += 1; + } } - // small bit of logging to note changes over time - println!("Changes: New: +{}, Current: +{}/-{}", roles_set[0], roles_set[1], roles_set[2]); + } + // small bit of logging to note changes over time + println!("Changes: New: +{}, Current: +{}/-{}", roles_set[0], roles_set[1], roles_set[2]); - { - let data_read = ctx.data.read().await; - let counter_lock = data_read.get::().expect("Expected MemberCounter in TypeMap.").clone(); - // The HashMap of CommandCounter is wrapped in an RwLock; since we want to write to it, we will - // open the lock in write mode. - let mut counter = counter_lock.write().await; + { + let data_read = ctx.data.read().await; + let counter_lock = data_read.get::().expect("Expected MemberCounter in TypeMap.").clone(); + // The HashMap of CommandCounter is wrapped in an RwLock; since we want to write to it, we will + // open the lock in write mode. + let mut counter = counter_lock.write().await; - // And we write the amount of times the command has been called to it. - counter.members_current = res.members_current; - counter.members = res.members; - }; + // And we write the amount of times the command has been called to it. + counter.members_current = res.members_current; + counter.members = res.members; + }; } -async fn fetch_accounts(ctx: Arc){ - let config_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Config in TypeMap.").clone() +async fn fetch_accounts(ctx: Arc) { + let config_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Config in TypeMap.").clone() + }; + let config = config_lock.read().await; + let auth = &config.auth; + let ldap_api = &config.ldap_api; + let url = format!("{}/ldap/discord?auth={}", ldap_api, auth); + if let Ok(result) = surf::get(url).recv_json::>().await { + let members_lock = { + let data_read = ctx.data.read().await; + data_read.get::().expect("Expected Members in TypeMap.").clone() }; - let config = config_lock.read().await; - let auth = &config.auth; - let ldap_api = &config.ldap_api; - let url = format!("{}/ldap/discord?auth={}", ldap_api, auth); - if let Ok(result) = surf::get(url).recv_json::>().await { - let members_lock = { - let data_read = ctx.data.read().await; - data_read.get::().expect("Expected Members in TypeMap.").clone() - }; - let mut accounts = members_lock.write().await; - *accounts = result; - } + let mut accounts = members_lock.write().await; + *accounts = result; + } } #[tokio::main] async fn main() { - let config = get_config(); + let config = get_config(); - // Intents are a bitflag, bitwise operations can be used to dictate which intents to use - let intents = GatewayIntents::GUILDS | GatewayIntents::GUILD_MESSAGES | GatewayIntents::MESSAGE_CONTENT | GatewayIntents::GUILD_MEMBERS; - // Build our client. - let mut client = Client::builder(&config.discord_token, intents) - .event_handler(Handler { is_loop_running: AtomicBool::new(false)}) - .await - .expect("Error creating client"); - - { - let mut data = client.data.write().await; - - // will keep track of how many past and current members we have - data.insert::(Arc::new(RwLock::new(MembersCount::default()))); - - // a list of all current members - data.insert::(Arc::new(RwLock::new(vec![]))); + // Intents are a bitflag, bitwise operations can be used to dictate which intents to use + let intents = GatewayIntents::GUILDS | GatewayIntents::GUILD_MESSAGES | GatewayIntents::MESSAGE_CONTENT | GatewayIntents::GUILD_MEMBERS; + // Build our client. + let mut client = Client::builder(&config.discord_token, intents) + .event_handler(Handler { + is_loop_running: AtomicBool::new(false), + }) + .await + .expect("Error creating client"); - data.insert::(Arc::new(RwLock::new(config))); - } - - - // Finally, start a single shard, and start listening to events. - // - // Shards will automatically attempt to reconnect, and will perform - // exponential backoff until it reconnects. - if let Err(why) = client.start().await { - println!("Client error: {:?}", why); - } + { + let mut data = client.data.write().await; + + // will keep track of how many past and current members we have + data.insert::(Arc::new(RwLock::new(MembersCount::default()))); + + // a list of all current members + data.insert::(Arc::new(RwLock::new(vec![]))); + + // make config available top all, strangely its easier to keep it in a shared lock state. + data.insert::(Arc::new(RwLock::new(config))); + } + + // Finally, start a single shard, and start listening to events. + // + // Shards will automatically attempt to reconnect, and will perform + // exponential backoff until it reconnects. + if let Err(why) = client.start().await { + println!("Client error: {:?}", why); + } } - struct Config { - server: GuildId, - member_role_current: RoleId, - member_role_past: RoleId, - ldap_api: String, - auth: String, - timing_update: u64, - timing_fetch: u64, - discord_token: String, + server: GuildId, + member_role_current: RoleId, + member_role_past: RoleId, + ldap_api: String, + auth: String, + timing_update: u64, + timing_fetch: u64, + discord_token: String, } impl TypeMapKey for Config { - type Value = Arc>; + type Value = Arc>; } fn get_config() -> Config { - dotenv().ok(); + dotenv().ok(); - // reasonable defaults - let mut config = Config { - server: Default::default(), - member_role_current: Default::default(), - member_role_past: Default::default(), - ldap_api: "https://api.account.skynet.ie".to_string(), - auth: "".to_string(), - timing_update: 0, - timing_fetch: 0, - discord_token: "".to_string(), - }; + // reasonable defaults + let mut config = Config { + server: Default::default(), + member_role_current: Default::default(), + member_role_past: Default::default(), + ldap_api: "https://api.account.skynet.ie".to_string(), + auth: "".to_string(), + timing_update: 0, + timing_fetch: 0, + discord_token: "".to_string(), + }; - if let Ok(x) = env::var("DISCORD_SERVER") { - config.server = GuildId::from(str_to_num::(&x)); - } - if let Ok(x) = env::var("DISCORD_ROLE_CURRENT") { - config.member_role_current = RoleId::from(str_to_num::(&x)); - } - if let Ok(x) = env::var("DISCORD_ROLE_PAST") { - config.member_role_past = RoleId::from(str_to_num::(&x)); - } - if let Ok(x) = env::var("LDAP_API") { - config.ldap_api = x.trim().to_string(); - } - if let Ok(x) = env::var("LDAP_DISCORD_AUTH") { - config.auth = x.trim().to_string(); - } - if let Ok(x) = env::var("DISCORD_TIMING_UPDATE") { - config.timing_update = str_to_num::(&x); - } - if let Ok(x) = env::var("DISCORD_TIMING_FETCH") { - config.timing_fetch = str_to_num::(&x); - } - if let Ok(x) = env::var("DISCORD_TOKEN") { - config.discord_token = x.trim().to_string(); - } - config + if let Ok(x) = env::var("DISCORD_SERVER") { + config.server = GuildId::from(str_to_num::(&x)); + } + if let Ok(x) = env::var("DISCORD_ROLE_CURRENT") { + config.member_role_current = RoleId::from(str_to_num::(&x)); + } + if let Ok(x) = env::var("DISCORD_ROLE_PAST") { + config.member_role_past = RoleId::from(str_to_num::(&x)); + } + if let Ok(x) = env::var("LDAP_API") { + config.ldap_api = x.trim().to_string(); + } + if let Ok(x) = env::var("LDAP_DISCORD_AUTH") { + config.auth = x.trim().to_string(); + } + if let Ok(x) = env::var("DISCORD_TIMING_UPDATE") { + config.timing_update = str_to_num::(&x); + } + if let Ok(x) = env::var("DISCORD_TIMING_FETCH") { + config.timing_fetch = str_to_num::(&x); + } + if let Ok(x) = env::var("DISCORD_TOKEN") { + config.discord_token = x.trim().to_string(); + } + config } fn str_to_num(x: &str) -> T { - x.trim().parse::().unwrap_or_default() -} \ No newline at end of file + x.trim().parse::().unwrap_or_default() +}