feat: now have two scripts that can run ona timer
This commit is contained in:
parent
cb363e6558
commit
ca6ae993c5
7 changed files with 538 additions and 239 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
@ -5,3 +5,5 @@
|
||||||
|
|
||||||
result
|
result
|
||||||
/result
|
/result
|
||||||
|
|
||||||
|
*.db
|
||||||
|
|
9
Cargo.lock
generated
9
Cargo.lock
generated
|
@ -2198,9 +2198,9 @@ checksum = "388a1df253eca08550bef6c72392cfe7c30914bf41df5269b68cbd6ff8f570a3"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "serde"
|
name = "serde"
|
||||||
version = "1.0.187"
|
version = "1.0.188"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "30a7fe14252655bd1e578af19f5fa00fe02fd0013b100ca6b49fde31c41bae4c"
|
checksum = "cf9e0fcba69a370eed61bcf2b728575f726b50b55cba78064753d708ddc7549e"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"serde_derive",
|
"serde_derive",
|
||||||
]
|
]
|
||||||
|
@ -2217,9 +2217,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "serde_derive"
|
name = "serde_derive"
|
||||||
version = "1.0.187"
|
version = "1.0.188"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "e46b2a6ca578b3f1d4501b12f78ed4692006d79d82a1a7c561c12dbc3d625eb8"
|
checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"proc-macro2",
|
"proc-macro2",
|
||||||
"quote",
|
"quote",
|
||||||
|
@ -2369,6 +2369,7 @@ dependencies = [
|
||||||
"dotenvy",
|
"dotenvy",
|
||||||
"lettre",
|
"lettre",
|
||||||
"maud",
|
"maud",
|
||||||
|
"serde",
|
||||||
"serenity",
|
"serenity",
|
||||||
"sqlx",
|
"sqlx",
|
||||||
"surf",
|
"surf",
|
||||||
|
|
11
Cargo.toml
11
Cargo.toml
|
@ -5,6 +5,13 @@ edition = "2021"
|
||||||
|
|
||||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "update_data"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "update_users"
|
||||||
|
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
serenity = { version = "0.11.6", default-features = false, features = ["client", "gateway", "rustls_backend", "model", "cache"] }
|
serenity = { version = "0.11.6", default-features = false, features = ["client", "gateway", "rustls_backend", "model", "cache"] }
|
||||||
tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
|
tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
|
||||||
|
@ -22,4 +29,6 @@ csv = "1.2"
|
||||||
|
|
||||||
# for email
|
# for email
|
||||||
lettre = "0.10.4"
|
lettre = "0.10.4"
|
||||||
maud = "0.25.0"
|
maud = "0.25.0"
|
||||||
|
|
||||||
|
serde = "1.0.188"
|
111
src/bin/update_data.rs
Normal file
111
src/bin/update_data.rs
Normal file
|
@ -0,0 +1,111 @@
|
||||||
|
use skynet_discord_bot::{db_init, get_config, Accounts, Config, DataBase};
|
||||||
|
use std::process;
|
||||||
|
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serenity::model::id::GuildId;
|
||||||
|
use serenity::{
|
||||||
|
async_trait,
|
||||||
|
client::{Context, EventHandler},
|
||||||
|
model::gateway::{GatewayIntents, Ready},
|
||||||
|
Client,
|
||||||
|
};
|
||||||
|
use sqlx::{Pool, Sqlite};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() {
|
||||||
|
let config = get_config();
|
||||||
|
let db = match db_init(&config).await {
|
||||||
|
Ok(x) => x,
|
||||||
|
Err(_) => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 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 {})
|
||||||
|
.await
|
||||||
|
.expect("Error creating client");
|
||||||
|
|
||||||
|
{
|
||||||
|
let mut data = client.data.write().await;
|
||||||
|
|
||||||
|
data.insert::<Config>(Arc::new(RwLock::new(config)));
|
||||||
|
data.insert::<DataBase>(Arc::new(RwLock::new(db)));
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 Handler;
|
||||||
|
#[async_trait]
|
||||||
|
impl EventHandler for Handler {
|
||||||
|
async fn ready(&self, ctx: Context, ready: Ready) {
|
||||||
|
println!("{} is connected!", ready.user.name);
|
||||||
|
|
||||||
|
fetch_accounts(&ctx).await;
|
||||||
|
|
||||||
|
process::exit(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct DiscordResult {
|
||||||
|
discord: String,
|
||||||
|
wolves_id: String,
|
||||||
|
}
|
||||||
|
async fn fetch_accounts(ctx: &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 db_lock = {
|
||||||
|
let data_read = ctx.data.read().await;
|
||||||
|
data_read.get::<DataBase>().expect("Expected Config in TypeMap.").clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
let db = db_lock.read().await;
|
||||||
|
|
||||||
|
// handle wolves api here
|
||||||
|
|
||||||
|
// get from skynet for the compsoc server only
|
||||||
|
let url = format!("{}/ldap/discord?auth={}", &config.ldap_api, &config.auth);
|
||||||
|
if let Ok(result) = surf::get(url).recv_json::<Vec<DiscordResult>>().await {
|
||||||
|
for user in result {
|
||||||
|
add_users_skynet(&db, &config.skynet_server, &user).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn add_users_skynet(db: &Pool<Sqlite>, server: &GuildId, user: &DiscordResult) {
|
||||||
|
match sqlx::query_as::<_, Accounts>(
|
||||||
|
"
|
||||||
|
UPDATE accounts
|
||||||
|
SET discord = ?
|
||||||
|
WHERE server = ? AND wolves_id = ?
|
||||||
|
",
|
||||||
|
)
|
||||||
|
.bind(&user.discord)
|
||||||
|
.bind(*server.as_u64() as i64)
|
||||||
|
.bind(&user.wolves_id)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(e) => {
|
||||||
|
println!("Failure to insert into {} {:?}", server.as_u64(), user);
|
||||||
|
println!("{:?}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
186
src/bin/update_users.rs
Normal file
186
src/bin/update_users.rs
Normal file
|
@ -0,0 +1,186 @@
|
||||||
|
use serenity::model::id::GuildId;
|
||||||
|
use serenity::{
|
||||||
|
async_trait,
|
||||||
|
client::{Context, EventHandler},
|
||||||
|
model::{
|
||||||
|
gateway::{GatewayIntents, Ready},
|
||||||
|
prelude::RoleId,
|
||||||
|
},
|
||||||
|
Client,
|
||||||
|
};
|
||||||
|
use skynet_discord_bot::{db_init, get_config, Accounts, Config, DataBase, Servers};
|
||||||
|
use sqlx::{Pool, Sqlite};
|
||||||
|
use std::process;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() {
|
||||||
|
let config = get_config();
|
||||||
|
let db = match db_init(&config).await {
|
||||||
|
Ok(x) => x,
|
||||||
|
Err(_) => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 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 {})
|
||||||
|
.await
|
||||||
|
.expect("Error creating client");
|
||||||
|
|
||||||
|
{
|
||||||
|
let mut data = client.data.write().await;
|
||||||
|
|
||||||
|
data.insert::<Config>(Arc::new(RwLock::new(config)));
|
||||||
|
data.insert::<DataBase>(Arc::new(RwLock::new(db)));
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(why) = client.start().await {
|
||||||
|
println!("Client error: {:?}", why);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Handler;
|
||||||
|
#[async_trait]
|
||||||
|
impl EventHandler for Handler {
|
||||||
|
async fn ready(&self, ctx: Context, ready: Ready) {
|
||||||
|
let ctx = Arc::new(ctx);
|
||||||
|
println!("{} is connected!", ready.user.name);
|
||||||
|
|
||||||
|
bulk_check(Arc::clone(&ctx)).await;
|
||||||
|
|
||||||
|
// finish up
|
||||||
|
process::exit(0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn bulk_check(ctx: Arc<Context>) {
|
||||||
|
let db_lock = {
|
||||||
|
let data_read = ctx.data.read().await;
|
||||||
|
data_read.get::<DataBase>().expect("Expected Config in TypeMap.").clone()
|
||||||
|
};
|
||||||
|
|
||||||
|
let db = db_lock.read().await;
|
||||||
|
|
||||||
|
for server_config in get_server_config_bulk(&db).await {
|
||||||
|
let Servers {
|
||||||
|
server,
|
||||||
|
role_past,
|
||||||
|
role_current,
|
||||||
|
..
|
||||||
|
} = server_config;
|
||||||
|
|
||||||
|
let mut roles_set = [0, 0, 0];
|
||||||
|
let mut members = vec![];
|
||||||
|
|
||||||
|
let server = GuildId::from(server as u64);
|
||||||
|
for member in get_server_member_bulk(&db, &server).await {
|
||||||
|
if let Some(x) = member.discord {
|
||||||
|
members.push(x);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut members_all = members.len();
|
||||||
|
|
||||||
|
if let Ok(x) = server.members(&ctx, None, None).await {
|
||||||
|
for mut member in x {
|
||||||
|
if members.contains(&member.user.name) {
|
||||||
|
let mut roles = vec![];
|
||||||
|
|
||||||
|
if let Some(role) = &role_past {
|
||||||
|
let role = RoleId::from(*role as u64);
|
||||||
|
if !member.roles.contains(&role) {
|
||||||
|
roles_set[0] += 1;
|
||||||
|
roles.push(role.to_owned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(role) = &role_current {
|
||||||
|
let role = RoleId::from(*role as u64);
|
||||||
|
if !member.roles.contains(&role) {
|
||||||
|
roles_set[1] += 1;
|
||||||
|
roles.push(role.to_owned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = member.add_roles(&ctx, &roles).await {
|
||||||
|
println!("{:?}", e);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// old and never
|
||||||
|
|
||||||
|
if let Some(role) = &role_past {
|
||||||
|
let role = RoleId::from(*role as u64);
|
||||||
|
if member.roles.contains(&role) {
|
||||||
|
members_all += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(role) = &role_current {
|
||||||
|
let role = RoleId::from(*role as u64);
|
||||||
|
if member.roles.contains(&role) {
|
||||||
|
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, role).await {
|
||||||
|
println!("{:?}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
set_server_numbers(&db, &server, members_all as i64, members.len() as i64).await;
|
||||||
|
|
||||||
|
// small bit of logging to note changes over time
|
||||||
|
println!("{:?} Changes: New: +{}, Current: +{}/-{}", server.as_u64(), roles_set[0], roles_set[1], roles_set[2]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
async fn get_server_config_bulk(db: &Pool<Sqlite>) -> Vec<Servers> {
|
||||||
|
sqlx::query_as::<_, Servers>(
|
||||||
|
r#"
|
||||||
|
SELECT *
|
||||||
|
FROM servers
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.fetch_all(db)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_server_member_bulk(db: &Pool<Sqlite>, server: &GuildId) -> Vec<Accounts> {
|
||||||
|
sqlx::query_as::<_, Accounts>(
|
||||||
|
r#"
|
||||||
|
SELECT *
|
||||||
|
FROM accounts
|
||||||
|
WHERE server = ? AND discord IS NOT NULL
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(*server.as_u64() as i64)
|
||||||
|
.fetch_all(db)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn set_server_numbers(db: &Pool<Sqlite>, server: &GuildId, past: i64, current: i64) {
|
||||||
|
match sqlx::query_as::<_, Accounts>(
|
||||||
|
"
|
||||||
|
UPDATE servers
|
||||||
|
SET member_past = ?, member_current = ?
|
||||||
|
WHERE server = ?
|
||||||
|
",
|
||||||
|
)
|
||||||
|
.bind(past)
|
||||||
|
.bind(current)
|
||||||
|
.bind(*server.as_u64() as i64)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(e) => {
|
||||||
|
println!("Failure to insert into {}", server.as_u64());
|
||||||
|
println!("{:?}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
192
src/lib.rs
Normal file
192
src/lib.rs
Normal file
|
@ -0,0 +1,192 @@
|
||||||
|
use dotenvy::dotenv;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serenity::{
|
||||||
|
model::{guild, id::GuildId},
|
||||||
|
prelude::TypeMapKey,
|
||||||
|
};
|
||||||
|
use sqlx::{
|
||||||
|
sqlite::{SqliteConnectOptions, SqlitePoolOptions},
|
||||||
|
Error, Pool, Sqlite,
|
||||||
|
};
|
||||||
|
use std::str::FromStr;
|
||||||
|
use std::{env, sync::Arc};
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
pub struct Config {
|
||||||
|
pub skynet_server: GuildId,
|
||||||
|
pub ldap_api: String,
|
||||||
|
pub auth: String,
|
||||||
|
pub timing_update: u64,
|
||||||
|
pub timing_fetch: u64,
|
||||||
|
pub discord_token: String,
|
||||||
|
|
||||||
|
pub home: String,
|
||||||
|
pub database: String,
|
||||||
|
|
||||||
|
pub mail_smtp: String,
|
||||||
|
pub mail_user: String,
|
||||||
|
pub mail_pass: String,
|
||||||
|
}
|
||||||
|
impl TypeMapKey for Config {
|
||||||
|
type Value = Arc<RwLock<Config>>;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct DataBase;
|
||||||
|
impl TypeMapKey for DataBase {
|
||||||
|
type Value = Arc<RwLock<Pool<Sqlite>>>;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_config() -> Config {
|
||||||
|
dotenv().ok();
|
||||||
|
|
||||||
|
// reasonable defaults
|
||||||
|
let mut config = Config {
|
||||||
|
skynet_server: Default::default(),
|
||||||
|
ldap_api: "https://api.account.skynet.ie".to_string(),
|
||||||
|
auth: "".to_string(),
|
||||||
|
timing_update: 0,
|
||||||
|
timing_fetch: 0,
|
||||||
|
discord_token: "".to_string(),
|
||||||
|
|
||||||
|
home: ".".to_string(),
|
||||||
|
database: "database.db".to_string(),
|
||||||
|
|
||||||
|
mail_smtp: "".to_string(),
|
||||||
|
mail_user: "".to_string(),
|
||||||
|
mail_pass: "".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Ok(x) = env::var("SKYNET_SERVER") {
|
||||||
|
config.skynet_server = GuildId::from(str_to_num::<u64>(&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::<u64>(&x);
|
||||||
|
}
|
||||||
|
if let Ok(x) = env::var("DISCORD_TIMING_FETCH") {
|
||||||
|
config.timing_fetch = str_to_num::<u64>(&x);
|
||||||
|
}
|
||||||
|
if let Ok(x) = env::var("DISCORD_TOKEN") {
|
||||||
|
config.discord_token = x.trim().to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Ok(x) = env::var("HOME") {
|
||||||
|
config.home = x.trim().to_string();
|
||||||
|
}
|
||||||
|
if let Ok(x) = env::var("DATABASE") {
|
||||||
|
config.database = x.trim().to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Ok(x) = env::var("EMAIL_SMTP") {
|
||||||
|
config.mail_smtp = x.trim().to_string();
|
||||||
|
}
|
||||||
|
if let Ok(x) = env::var("EMAIL_USER") {
|
||||||
|
config.mail_user = x.trim().to_string();
|
||||||
|
}
|
||||||
|
if let Ok(x) = env::var("EMAIL_PASS") {
|
||||||
|
config.mail_pass = x.trim().to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
config
|
||||||
|
}
|
||||||
|
|
||||||
|
fn str_to_num<T: FromStr + Default>(x: &str) -> T {
|
||||||
|
x.trim().parse::<T>().unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize, Serialize, sqlx::FromRow)]
|
||||||
|
pub struct Accounts {
|
||||||
|
pub server: i64,
|
||||||
|
pub wolves_id: String,
|
||||||
|
pub email: String,
|
||||||
|
pub expiry: String,
|
||||||
|
pub discord: Option<String>,
|
||||||
|
pub minecraft: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize, Serialize, sqlx::FromRow)]
|
||||||
|
pub struct Servers {
|
||||||
|
pub server: i64,
|
||||||
|
pub wolves_api: String,
|
||||||
|
pub role_past: Option<i64>,
|
||||||
|
pub role_current: Option<i64>,
|
||||||
|
pub member_past: i64,
|
||||||
|
pub member_current: i64,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn db_init(config: &Config) -> Result<Pool<Sqlite>, Error> {
|
||||||
|
let database = format!("{}/{}", &config.home, &config.database);
|
||||||
|
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(5)
|
||||||
|
.connect_with(SqliteConnectOptions::from_str(&format!("sqlite://{}", database))?.create_if_missing(true))
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE IF NOT EXISTS accounts (
|
||||||
|
server integer not null,
|
||||||
|
wolves_id text not null,
|
||||||
|
email text not null,
|
||||||
|
expiry text not null,
|
||||||
|
discord text,
|
||||||
|
minecraft text,
|
||||||
|
CONSTRAINT con_server_wolves PRIMARY KEY(server,wolves_id)
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
sqlx::query("CREATE INDEX IF NOT EXISTS index_server ON accounts (server)").execute(&pool).await?;
|
||||||
|
sqlx::query("CREATE INDEX IF NOT EXISTS index_wolves_id ON accounts (wolves_id)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE IF NOT EXISTS servers (
|
||||||
|
server integer key,
|
||||||
|
wolves_api text not null,
|
||||||
|
role_past integer,
|
||||||
|
role_current integer,
|
||||||
|
member_past integer DEFAULT 0,
|
||||||
|
member_current integer DEFAULT 0
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(pool)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_server_config(db: &Pool<Sqlite>, server: &GuildId) -> Option<Servers> {
|
||||||
|
sqlx::query_as::<_, Servers>(
|
||||||
|
r#"
|
||||||
|
SELECT *
|
||||||
|
FROM servers
|
||||||
|
WHERE server = ?
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(*server.as_u64() as i64)
|
||||||
|
.fetch_one(db)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_server_member(db: &Pool<Sqlite>, server: &GuildId, member: &guild::Member) -> Option<Accounts> {
|
||||||
|
sqlx::query_as::<_, Accounts>(
|
||||||
|
r#"
|
||||||
|
SELECT *
|
||||||
|
FROM accounts
|
||||||
|
WHERE server = ? AND discord = ?
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(*server.as_u64() as i64)
|
||||||
|
.bind(&member.user.name)
|
||||||
|
.fetch_one(db)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
}
|
266
src/main.rs
266
src/main.rs
|
@ -1,53 +1,49 @@
|
||||||
use dotenvy::dotenv;
|
|
||||||
use serenity::{
|
use serenity::{
|
||||||
async_trait,
|
async_trait,
|
||||||
client::{Context, EventHandler},
|
client::{Context, EventHandler},
|
||||||
model::{
|
model::{
|
||||||
gateway::{GatewayIntents, Ready},
|
gateway::{GatewayIntents, Ready},
|
||||||
guild::Member,
|
guild,
|
||||||
id::GuildId,
|
|
||||||
prelude::RoleId,
|
prelude::RoleId,
|
||||||
},
|
},
|
||||||
prelude::TypeMapKey,
|
|
||||||
Client,
|
Client,
|
||||||
};
|
};
|
||||||
use std::{
|
use std::sync::Arc;
|
||||||
env,
|
|
||||||
sync::{
|
use skynet_discord_bot::{db_init, get_config, get_server_config, get_server_member, Config, DataBase};
|
||||||
atomic::{AtomicBool, Ordering},
|
|
||||||
Arc,
|
|
||||||
},
|
|
||||||
time::Duration,
|
|
||||||
};
|
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
struct Handler {
|
struct Handler;
|
||||||
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: guild::Member) {
|
||||||
let config_lock = {
|
let db_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::<DataBase>().expect("Expected Config in TypeMap.").clone()
|
||||||
};
|
};
|
||||||
let config = config_lock.read().await;
|
|
||||||
|
|
||||||
let members_lock = {
|
let db = db_lock.read().await;
|
||||||
let data_read = ctx.data.read().await;
|
let config = match get_server_config(&db, &new_member.guild_id).await {
|
||||||
data_read.get::<Members>().expect("Expected Members in TypeMap.").clone()
|
None => return,
|
||||||
|
Some(x) => x,
|
||||||
};
|
};
|
||||||
let members = members_lock.read().await;
|
|
||||||
|
|
||||||
if members.contains(&new_member.user.name) {
|
if get_server_member(&db, &new_member.guild_id, &new_member).await.is_some() {
|
||||||
let mut roles = vec![];
|
let mut roles = vec![];
|
||||||
|
|
||||||
if !new_member.roles.contains(&config.member_role_past) {
|
if let Some(role) = &config.role_past {
|
||||||
roles.push(config.member_role_past);
|
let role = RoleId::from(*role as u64);
|
||||||
|
if !new_member.roles.contains(&role) {
|
||||||
|
roles.push(role.to_owned());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if !new_member.roles.contains(&config.member_role_current) {
|
|
||||||
roles.push(config.member_role_current);
|
if let Some(role) = &config.role_current {
|
||||||
|
let role = RoleId::from(*role as u64);
|
||||||
|
if !new_member.roles.contains(&role) {
|
||||||
|
roles.push(role.to_owned());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Err(e) = new_member.add_roles(&ctx, &roles).await {
|
if let Err(e) = new_member.add_roles(&ctx, &roles).await {
|
||||||
|
@ -56,171 +52,32 @@ impl EventHandler for Handler {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn ready(&self, ctx: Context, ready: Ready) {
|
async fn ready(&self, _ctx: Context, ready: Ready) {
|
||||||
let ctx = Arc::new(ctx);
|
println!("[Main] {} is connected!", ready.user.name);
|
||||||
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 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();
|
||||||
|
let db = match db_init(&config).await {
|
||||||
|
Ok(x) => x,
|
||||||
|
Err(_) => return,
|
||||||
|
};
|
||||||
|
|
||||||
// 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 {
|
.event_handler(Handler {})
|
||||||
is_loop_running: AtomicBool::new(false),
|
|
||||||
})
|
|
||||||
.await
|
.await
|
||||||
.expect("Error creating client");
|
.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
|
|
||||||
data.insert::<MemberCounter>(Arc::new(RwLock::new(MembersCount::default())));
|
|
||||||
|
|
||||||
// a list of all current members
|
|
||||||
data.insert::<Members>(Arc::new(RwLock::new(vec![])));
|
|
||||||
|
|
||||||
// make config available top all, strangely its easier to keep it in a shared lock state.
|
|
||||||
data.insert::<Config>(Arc::new(RwLock::new(config)));
|
data.insert::<Config>(Arc::new(RwLock::new(config)));
|
||||||
|
data.insert::<DataBase>(Arc::new(RwLock::new(db)));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Finally, start a single shard, and start listening to events.
|
// Finally, start a single shard, and start listening to events.
|
||||||
|
@ -231,62 +88,3 @@ async fn main() {
|
||||||
println!("Client error: {:?}", why);
|
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,
|
|
||||||
}
|
|
||||||
impl TypeMapKey for Config {
|
|
||||||
type Value = Arc<RwLock<Config>>;
|
|
||||||
}
|
|
||||||
fn get_config() -> Config {
|
|
||||||
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(),
|
|
||||||
};
|
|
||||||
|
|
||||||
if let Ok(x) = env::var("DISCORD_SERVER") {
|
|
||||||
config.server = GuildId::from(str_to_num::<u64>(&x));
|
|
||||||
}
|
|
||||||
if let Ok(x) = env::var("DISCORD_ROLE_CURRENT") {
|
|
||||||
config.member_role_current = RoleId::from(str_to_num::<u64>(&x));
|
|
||||||
}
|
|
||||||
if let Ok(x) = env::var("DISCORD_ROLE_PAST") {
|
|
||||||
config.member_role_past = RoleId::from(str_to_num::<u64>(&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::<u64>(&x);
|
|
||||||
}
|
|
||||||
if let Ok(x) = env::var("DISCORD_TIMING_FETCH") {
|
|
||||||
config.timing_fetch = str_to_num::<u64>(&x);
|
|
||||||
}
|
|
||||||
if let Ok(x) = env::var("DISCORD_TOKEN") {
|
|
||||||
config.discord_token = x.trim().to_string();
|
|
||||||
}
|
|
||||||
config
|
|
||||||
}
|
|
||||||
|
|
||||||
fn str_to_num<T: std::str::FromStr + Default>(x: &str) -> T {
|
|
||||||
x.trim().parse::<T>().unwrap_or_default()
|
|
||||||
}
|
|
||||||
|
|
Loading…
Reference in a new issue