peace/discord-bot
1
0
Fork 0
forked from Skynet/discord-bot
Code Pull requests Activity

feat: fmt and clippy

Browse source
This commit is contained in:
silver 2023-08-27 17:38:30 +01:00
parent 9b42187b22
commit 802f3fce59
2 changed files with 245 additions and 224 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
.rustfmt.toml Normal file
View file

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

478
src/main.rs
View file

@ -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 dotenvy::dotenv;
use serenity::async_trait; use serenity::{
use serenity::model::gateway::{GatewayIntents, Ready}; async_trait,
use serenity::model::guild::Member; client::{Context, EventHandler},
use serenity::model::id::GuildId; model::{
use serenity::prelude::*; gateway::{GatewayIntents, Ready},
use serenity::model::prelude::RoleId; 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 { struct Handler {
is_loop_running: AtomicBool, is_loop_running: AtomicBool,
} }
#[async_trait] #[async_trait]
impl EventHandler for Handler { impl EventHandler for Handler {
async fn guild_member_addition(&self, ctx: Context, mut new_member: Member) { 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::<Config>().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::<Members>().expect("Expected Members in TypeMap.").clone()
};
let members = members_lock.read().await;
if members.contains(&new_member.user.name) {
let mut roles = vec![];
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::<Config>().expect("Expected Config in TypeMap.").clone()
};
let config = config_lock.read().await;
let timing_update = config.timing_update;
let timing_fetch = config.timing_fetch;
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,
}
struct MemberCounter;
impl TypeMapKey for MemberCounter {
type Value = Arc<RwLock<MembersCount>>;
}
struct Members;
impl TypeMapKey for Members {
type Value = Arc<RwLock<Vec<String>>>;
}
async fn bulk_check(ctx: Arc<Context>){
let config_lock = { let config_lock = {
let data_read = ctx.data.read().await; let data_read = ctx.data.read().await;
data_read.get::<Config>().expect("Expected Config in TypeMap.").clone() data_read.get::<Config>().expect("Expected Config in TypeMap.").clone()
}; };
let config = config_lock.read().await; let config = config_lock.read().await;
let members_lock = { let members_lock = {
let data_read = ctx.data.read().await; let data_read = ctx.data.read().await;
data_read.get::<Members>().expect("Expected Members in TypeMap.").clone() data_read.get::<Members>().expect("Expected Members in TypeMap.").clone()
}; };
let members = members_lock.read().await; let members = members_lock.read().await;
let mut roles_set = [0,0,0]; if members.contains(&new_member.user.name) {
let mut res = MembersCount { let mut roles = vec![];
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){ if !new_member.roles.contains(&config.member_role_past) {
roles_set[0] += 1; roles.push(config.member_role_past);
roles.push(config.member_role_past); }
} if !new_member.roles.contains(&config.member_role_current) {
if !member.roles.contains(&config.member_role_current){ roles.push(config.member_role_current);
roles_set[1] += 1; }
roles.push(config.member_role_current);
}
if let Err(e) = member.add_roles(&ctx, &roles).await { if let Err(e) = new_member.add_roles(&ctx, &roles).await {
println!("{:?}", e); 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]);
{ async fn ready(&self, ctx: Context, ready: Ready) {
let data_read = ctx.data.read().await; let ctx = Arc::new(ctx);
let counter_lock = data_read.get::<MemberCounter>().expect("Expected MemberCounter in TypeMap.").clone(); println!("{} is connected!", ready.user.name);
// 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;
};
}
async fn fetch_accounts(ctx: Arc<Context>){
let config_lock = { let config_lock = {
let data_read = ctx.data.read().await; let data_read = ctx.data.read().await;
data_read.get::<Config>().expect("Expected Config in TypeMap.").clone() data_read.get::<Config>().expect("Expected Config in TypeMap.").clone()
}; };
let config = config_lock.read().await; let config = config_lock.read().await;
let auth = &config.auth; let timing_update = config.timing_update;
let ldap_api = &config.ldap_api; let timing_fetch = config.timing_fetch;
let url = format!("{}/ldap/discord?auth={}", ldap_api, auth);
if let Ok(result) = surf::get(url).recv_json::<Vec<String>>().await { if !self.is_loop_running.load(Ordering::Relaxed) {
let members_lock = { // We have to clone the Arc, as it gets moved into the new thread.
let data_read = ctx.data.read().await; let ctx1 = Arc::clone(&ctx);
data_read.get::<Members>().expect("Expected Members in TypeMap.").clone() // tokio::spawn creates a new green thread that can run in parallel with the rest of
}; // the application.
let mut accounts = members_lock.write().await; tokio::spawn(async move {
*accounts = result; 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,
}
struct MemberCounter;
impl TypeMapKey for MemberCounter {
type Value = Arc<RwLock<MembersCount>>;
}
struct Members;
impl TypeMapKey for Members {
type Value = Arc<RwLock<Vec<String>>>;
}
async fn bulk_check(ctx: Arc<Context>) {
let config_lock = {
let data_read = ctx.data.read().await;
data_read.get::<Config>().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::<Members>().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![];
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]);
{
let data_read = ctx.data.read().await;
let counter_lock = data_read.get::<MemberCounter>().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;
};
}
async fn fetch_accounts(ctx: Arc<Context>) {
let config_lock = {
let data_read = ctx.data.read().await;
data_read.get::<Config>().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::<Vec<String>>().await {
let members_lock = {
let data_read = ctx.data.read().await;
data_read.get::<Members>().expect("Expected Members in TypeMap.").clone()
};
let mut accounts = members_lock.write().await;
*accounts = result;
}
} }
#[tokio::main] #[tokio::main]
async fn 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 // 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; let intents = GatewayIntents::GUILDS | GatewayIntents::GUILD_MESSAGES | GatewayIntents::MESSAGE_CONTENT | GatewayIntents::GUILD_MEMBERS;
// Build our client. // Build our client.
let mut client = Client::builder(&config.discord_token, intents) let mut client = Client::builder(&config.discord_token, intents)
.event_handler(Handler { is_loop_running: AtomicBool::new(false)}) .event_handler(Handler {
.await is_loop_running: AtomicBool::new(false),
.expect("Error creating client"); })
.await
.expect("Error creating client");
{ {
let mut data = client.data.write().await; let mut data = client.data.write().await;
// will keep track of how many past and current members we have // will keep track of how many past and current members we have
data.insert::<MemberCounter>(Arc::new(RwLock::new(MembersCount::default()))); data.insert::<MemberCounter>(Arc::new(RwLock::new(MembersCount::default())));
// a list of all current members // a list of all current members
data.insert::<Members>(Arc::new(RwLock::new(vec![]))); data.insert::<Members>(Arc::new(RwLock::new(vec![])));
data.insert::<Config>(Arc::new(RwLock::new(config))); // make config available top all, strangely its easier to keep it in a shared lock state.
} data.insert::<Config>(Arc::new(RwLock::new(config)));
}
// Finally, start a single shard, and start listening to events.
// Finally, start a single shard, and start listening to events. //
// // Shards will automatically attempt to reconnect, and will perform
// Shards will automatically attempt to reconnect, and will perform // exponential backoff until it reconnects.
// exponential backoff until it reconnects. if let Err(why) = client.start().await {
if let Err(why) = client.start().await { println!("Client error: {:?}", why);
println!("Client error: {:?}", why); }
}
} }
struct Config { struct Config {
server: GuildId, server: GuildId,
member_role_current: RoleId, member_role_current: RoleId,
member_role_past: RoleId, member_role_past: RoleId,
ldap_api: String, ldap_api: String,
auth: String, auth: String,
timing_update: u64, timing_update: u64,
timing_fetch: u64, timing_fetch: u64,
discord_token: String, discord_token: String,
} }
impl TypeMapKey for Config { impl TypeMapKey for Config {
type Value = Arc<RwLock<Config>>; type Value = Arc<RwLock<Config>>;
} }
fn get_config() -> Config { fn get_config() -> Config {
dotenv().ok(); dotenv().ok();
// reasonable defaults // reasonable defaults
let mut config = Config { let mut config = Config {
server: Default::default(), server: Default::default(),
member_role_current: Default::default(), member_role_current: Default::default(),
member_role_past: Default::default(), member_role_past: Default::default(),
ldap_api: "https://api.account.skynet.ie".to_string(), ldap_api: "https://api.account.skynet.ie".to_string(),
auth: "".to_string(), auth: "".to_string(),
timing_update: 0, timing_update: 0,
timing_fetch: 0, timing_fetch: 0,
discord_token: "".to_string(), discord_token: "".to_string(),
}; };
if let Ok(x) = env::var("DISCORD_SERVER") { if let Ok(x) = env::var("DISCORD_SERVER") {
config.server = GuildId::from(str_to_num::<u64>(&x)); config.server = GuildId::from(str_to_num::<u64>(&x));
} }
if let Ok(x) = env::var("DISCORD_ROLE_CURRENT") { if let Ok(x) = env::var("DISCORD_ROLE_CURRENT") {
config.member_role_current = RoleId::from(str_to_num::<u64>(&x)); config.member_role_current = RoleId::from(str_to_num::<u64>(&x));
} }
if let Ok(x) = env::var("DISCORD_ROLE_PAST") { if let Ok(x) = env::var("DISCORD_ROLE_PAST") {
config.member_role_past = RoleId::from(str_to_num::<u64>(&x)); config.member_role_past = RoleId::from(str_to_num::<u64>(&x));
} }
if let Ok(x) = env::var("LDAP_API") { if let Ok(x) = env::var("LDAP_API") {
config.ldap_api = x.trim().to_string(); config.ldap_api = x.trim().to_string();
} }
if let Ok(x) = env::var("LDAP_DISCORD_AUTH") { if let Ok(x) = env::var("LDAP_DISCORD_AUTH") {
config.auth = x.trim().to_string(); config.auth = x.trim().to_string();
} }
if let Ok(x) = env::var("DISCORD_TIMING_UPDATE") { if let Ok(x) = env::var("DISCORD_TIMING_UPDATE") {
config.timing_update = str_to_num::<u64>(&x); config.timing_update = str_to_num::<u64>(&x);
} }
if let Ok(x) = env::var("DISCORD_TIMING_FETCH") { if let Ok(x) = env::var("DISCORD_TIMING_FETCH") {
config.timing_fetch = str_to_num::<u64>(&x); config.timing_fetch = str_to_num::<u64>(&x);
} }
if let Ok(x) = env::var("DISCORD_TOKEN") { if let Ok(x) = env::var("DISCORD_TOKEN") {
config.discord_token = x.trim().to_string(); config.discord_token = x.trim().to_string();
} }
config config
} }
fn str_to_num<T: std::str::FromStr + Default>(x: &str) -> T { fn str_to_num<T: std::str::FromStr + Default>(x: &str) -> T {
x.trim().parse::<T>().unwrap_or_default() x.trim().parse::<T>().unwrap_or_default()
} }