forked from Skynet/discord-bot
Merge branch '#3-wolves-api' into 'main'
#3 wolves api See merge request compsoc1/skynet/discord-bot!1
This commit is contained in:
commit
2162c4a824
8 changed files with 1820 additions and 261 deletions
6
.gitignore
vendored
6
.gitignore
vendored
|
@ -5,3 +5,9 @@
|
|||
|
||||
result
|
||||
/result
|
||||
|
||||
*.db
|
||||
|
||||
tmp/
|
||||
tmp.*
|
||||
*.csv
|
||||
|
|
1066
Cargo.lock
generated
1066
Cargo.lock
generated
File diff suppressed because it is too large
Load diff
24
Cargo.toml
24
Cargo.toml
|
@ -5,6 +5,13 @@ edition = "2021"
|
|||
|
||||
# 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]
|
||||
serenity = { version = "0.11.6", default-features = false, features = ["client", "gateway", "rustls_backend", "model", "cache"] }
|
||||
tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
|
||||
|
@ -12,4 +19,19 @@ tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
|
|||
# to make the http requests
|
||||
surf = "2.3.2"
|
||||
|
||||
dotenvy = "0.15.7"
|
||||
dotenvy = "0.15.7"
|
||||
|
||||
# For sqlite
|
||||
sqlx = { version = "0.7.1", features = [ "runtime-tokio", "sqlite" ] }
|
||||
|
||||
# fancy time stuff
|
||||
chrono = "0.4.26"
|
||||
|
||||
# handlign teh csv export from wolves
|
||||
csv = "1.2"
|
||||
|
||||
# for email
|
||||
lettre = "0.10.4"
|
||||
maud = "0.25.0"
|
||||
|
||||
serde = "1.0.188"
|
40
flake.nix
40
flake.nix
|
@ -52,6 +52,10 @@
|
|||
LDAP_API = cfg.ldap;
|
||||
DISCORD_TIMING_UPDATE = cfg.discord.timing.update;
|
||||
DISCORD_TIMING_FETCH = cfg.discord.timing.fetch;
|
||||
|
||||
# local details
|
||||
HOME = cfg.home;
|
||||
DATABASE = "database.db";
|
||||
};
|
||||
in {
|
||||
options.services."${package_name}" = {
|
||||
|
@ -66,6 +70,10 @@
|
|||
type = types.str;
|
||||
description = "ENV file with DISCORD_TOKEN";
|
||||
};
|
||||
mail = mkOption rec {
|
||||
type = types.str;
|
||||
description = "ENV file with EMAIL_SMTP, EMAIL_USER, EMAIL_PASS";
|
||||
};
|
||||
};
|
||||
|
||||
discord = {
|
||||
|
@ -102,9 +110,37 @@
|
|||
default = "https://api.account.skynet.ie";
|
||||
description = "Location of the ldap api";
|
||||
};
|
||||
|
||||
user = mkOption rec {
|
||||
type = types.str;
|
||||
default = "${package_name}";
|
||||
description = "The user to run the service";
|
||||
};
|
||||
|
||||
home = mkOption rec {
|
||||
type = types.str;
|
||||
default = "/etc/${cfg.prefix}${package_name}";
|
||||
description = "The home for the user";
|
||||
};
|
||||
|
||||
prefix = mkOption rec {
|
||||
type = types.str;
|
||||
default = "skynet_";
|
||||
example = default;
|
||||
description = "The prefix used to name service/folders";
|
||||
};
|
||||
};
|
||||
|
||||
config = mkIf cfg.enable {
|
||||
|
||||
users.groups."${cfg.user}" = { };
|
||||
|
||||
users.users."${cfg.user}" = {
|
||||
createHome = true;
|
||||
isSystemUser = true;
|
||||
home = "${cfg.home}";
|
||||
group = "${cfg.user}";
|
||||
};
|
||||
|
||||
systemd.services = {
|
||||
# main service
|
||||
|
@ -116,13 +152,15 @@
|
|||
environment = environment_config;
|
||||
|
||||
serviceConfig = {
|
||||
DynamicUser = "yes";
|
||||
User = "${cfg.user}";
|
||||
Group = "${cfg.user}";
|
||||
Restart = "always";
|
||||
ExecStart = "${self.defaultPackage."${system}"}/bin/${package_name}";
|
||||
# can have multiple env files
|
||||
EnvironmentFile = [
|
||||
"${cfg.env.ldap}"
|
||||
"${cfg.env.discord}"
|
||||
"${cfg.env.mail}"
|
||||
];
|
||||
};
|
||||
};
|
||||
|
|
241
src/bin/update_data.rs
Normal file
241
src/bin/update_data.rs
Normal file
|
@ -0,0 +1,241 @@
|
|||
use skynet_discord_bot::{db_init, get_config, get_server_config_bulk, Accounts, Config, Servers};
|
||||
|
||||
use serde::Deserialize;
|
||||
use serenity::model::id::GuildId;
|
||||
use sqlx::{Pool, Sqlite};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let config = get_config();
|
||||
let db = match db_init(&config).await {
|
||||
Ok(x) => x,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
// handle wolves api here
|
||||
get_wolves_csv(&db, &config).await;
|
||||
// handle wolves api here
|
||||
get_wolves(&db).await;
|
||||
// get from skynet for the compsoc server only
|
||||
get_skynet(&db, &config).await;
|
||||
}
|
||||
|
||||
async fn get_wolves_csv(db: &Pool<Sqlite>, config: &Config) {
|
||||
if let Ok(accounts) = get_csv(config) {
|
||||
for account in accounts {
|
||||
add_users_wolves_csv(db, &config.skynet_server, &account).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct RecordCSV {
|
||||
#[serde(rename = "MemID")]
|
||||
mem_id: String,
|
||||
#[serde(rename = "Contact Email")]
|
||||
email: String,
|
||||
#[serde(rename = "Expiry")]
|
||||
expiry: String,
|
||||
}
|
||||
|
||||
impl From<RecordCSV> for Accounts {
|
||||
fn from(input: RecordCSV) -> Self {
|
||||
Self {
|
||||
server: Default::default(),
|
||||
id_wolves: "".to_string(),
|
||||
id_member: input.mem_id,
|
||||
email: input.email,
|
||||
expiry: input.expiry,
|
||||
discord: None,
|
||||
minecraft: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_csv(config: &Config) -> Result<Vec<Accounts>, Box<dyn std::error::Error>> {
|
||||
let mut records: Vec<Accounts> = vec![];
|
||||
|
||||
let csv = format!("{}/{}", &config.home, &config.csv);
|
||||
if let Ok(mut rdr) = csv::Reader::from_path(csv) {
|
||||
for result in rdr.deserialize() {
|
||||
// Notice that we need to provide a type hint for automatic
|
||||
// deserialization.
|
||||
let record: RecordCSV = result?;
|
||||
if record.mem_id.is_empty() {
|
||||
continue;
|
||||
}
|
||||
records.push(Accounts::from(record));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(records)
|
||||
}
|
||||
async fn add_users_wolves_csv(db: &Pool<Sqlite>, server: &GuildId, user: &Accounts) {
|
||||
let existing = match sqlx::query_as::<_, Accounts>(
|
||||
r#"
|
||||
SELECT *
|
||||
FROM accounts
|
||||
WHERE server = ? AND id_member = ?
|
||||
"#,
|
||||
)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&user.id_member)
|
||||
.fetch_one(db)
|
||||
.await
|
||||
{
|
||||
Ok(acc) => acc.id_wolves,
|
||||
Err(_) => String::new(),
|
||||
};
|
||||
|
||||
match sqlx::query_as::<_, Accounts>(
|
||||
"
|
||||
INSERT OR REPLACE INTO accounts (server, id_wolves, id_member, email, expiry)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)
|
||||
",
|
||||
)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&existing)
|
||||
.bind(&user.id_member)
|
||||
.bind(&user.email)
|
||||
.bind(&user.expiry)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failure to insert into {} {:?}", server.as_u64(), user);
|
||||
println!("{:?}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SkynetResult {
|
||||
discord: String,
|
||||
id_wolves: String,
|
||||
id_member: String,
|
||||
}
|
||||
async fn get_skynet(db: &Pool<Sqlite>, config: &Config) {
|
||||
let url = format!("{}/ldap/discord?auth={}", &config.ldap_api, &config.auth);
|
||||
if let Ok(result) = surf::get(url).recv_json::<Vec<SkynetResult>>().await {
|
||||
for user in result {
|
||||
add_users_skynet(db, &config.skynet_server, &user).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn add_users_skynet(db: &Pool<Sqlite>, server: &GuildId, user: &SkynetResult) {
|
||||
if !user.id_wolves.is_empty() {
|
||||
match sqlx::query_as::<_, Accounts>(
|
||||
"
|
||||
UPDATE accounts
|
||||
SET discord = ?
|
||||
WHERE server = ? AND id_wolves = ?
|
||||
",
|
||||
)
|
||||
.bind(&user.discord)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&user.id_wolves)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failure to insert into {} {:?}", server.as_u64(), user);
|
||||
println!("{:?}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !user.id_member.is_empty() {
|
||||
match sqlx::query_as::<_, Accounts>(
|
||||
"
|
||||
UPDATE accounts
|
||||
SET discord = ?
|
||||
WHERE server = ? AND id_member = ?
|
||||
",
|
||||
)
|
||||
.bind(&user.discord)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&user.id_member)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failure to insert into {} {:?}", server.as_u64(), user);
|
||||
println!("{:?}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WolvesResult {
|
||||
pub id_wolves: String,
|
||||
pub id_member: String,
|
||||
pub email: String,
|
||||
pub expiry: String,
|
||||
}
|
||||
async fn get_wolves(db: &Pool<Sqlite>) {
|
||||
for server_config in get_server_config_bulk(db).await {
|
||||
let Servers {
|
||||
server,
|
||||
//wolves_api,
|
||||
..
|
||||
} = server_config;
|
||||
|
||||
// get the data here
|
||||
let result: Vec<WolvesResult> = vec![WolvesResult {
|
||||
id_wolves: "12345".to_string(),
|
||||
id_member: "166425".to_string(),
|
||||
email: "ul.wolves2@brendan.ie".to_string(),
|
||||
expiry: "2024-08-31".to_string(),
|
||||
}];
|
||||
|
||||
for user in result {
|
||||
add_users_wolves(db, &server, &user).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn add_users_wolves(db: &Pool<Sqlite>, server: &GuildId, user: &WolvesResult) {
|
||||
match sqlx::query_as::<_, Accounts>(
|
||||
"
|
||||
UPDATE accounts
|
||||
SET id_wolves = ?
|
||||
WHERE server = ? AND id_member = ?
|
||||
",
|
||||
)
|
||||
.bind(&user.id_wolves)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&user.id_member)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failure to update into {} {:?}", server.as_u64(), user);
|
||||
println!("{:?}", e);
|
||||
}
|
||||
}
|
||||
|
||||
match sqlx::query_as::<_, Accounts>(
|
||||
"
|
||||
INSERT OR REPLACE INTO accounts (server, id_wolves, id_member, email, expiry)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)
|
||||
",
|
||||
)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(&user.id_wolves)
|
||||
.bind(&user.id_member)
|
||||
.bind(&user.email)
|
||||
.bind(&user.expiry)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("Failure to insert into {} {:?}", server.as_u64(), user);
|
||||
println!("{:?}", e);
|
||||
}
|
||||
}
|
||||
}
|
169
src/bin/update_users.rs
Normal file
169
src/bin/update_users.rs
Normal file
|
@ -0,0 +1,169 @@
|
|||
use serenity::{
|
||||
async_trait,
|
||||
client::{Context, EventHandler},
|
||||
model::{
|
||||
gateway::{GatewayIntents, Ready},
|
||||
id::GuildId,
|
||||
},
|
||||
Client,
|
||||
};
|
||||
use skynet_discord_bot::{db_init, get_config, get_now_iso, get_server_config_bulk, Accounts, Config, DataBase, Servers};
|
||||
use sqlx::{Pool, Sqlite};
|
||||
use std::{process, 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![];
|
||||
|
||||
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 {
|
||||
if !member.roles.contains(role) {
|
||||
roles_set[0] += 1;
|
||||
roles.push(role.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(role) = &role_current {
|
||||
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 {
|
||||
if member.roles.contains(role) {
|
||||
members_all += 1;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(role) = &role_current {
|
||||
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_member_bulk(db: &Pool<Sqlite>, server: &GuildId) -> Vec<Accounts> {
|
||||
sqlx::query_as::<_, Accounts>(
|
||||
r#"
|
||||
SELECT *
|
||||
FROM accounts
|
||||
WHERE server = ? AND discord IS NOT NULL AND expiry > ?
|
||||
"#,
|
||||
)
|
||||
.bind(*server.as_u64() as i64)
|
||||
.bind(get_now_iso(true))
|
||||
.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);
|
||||
}
|
||||
}
|
||||
}
|
270
src/lib.rs
Normal file
270
src/lib.rs
Normal file
|
@ -0,0 +1,270 @@
|
|||
use dotenvy::dotenv;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serenity::{
|
||||
model::{
|
||||
guild,
|
||||
id::{GuildId, RoleId},
|
||||
},
|
||||
prelude::TypeMapKey,
|
||||
};
|
||||
|
||||
use chrono::{Datelike, SecondsFormat, Utc};
|
||||
use sqlx::{
|
||||
sqlite::{SqliteConnectOptions, SqlitePoolOptions, SqliteRow},
|
||||
Error, FromRow, Pool, Row, Sqlite,
|
||||
};
|
||||
use std::{env, str::FromStr, 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 csv: 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(),
|
||||
csv: "wolves.csv".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("CSV") {
|
||||
config.csv = 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)]
|
||||
pub struct Accounts {
|
||||
pub server: GuildId,
|
||||
pub id_wolves: String,
|
||||
pub id_member: String,
|
||||
pub email: String,
|
||||
pub expiry: String,
|
||||
pub discord: Option<String>,
|
||||
pub minecraft: Option<String>,
|
||||
}
|
||||
|
||||
impl<'r> FromRow<'r, SqliteRow> for Accounts {
|
||||
fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
|
||||
let server_tmp: i64 = row.try_get("server")?;
|
||||
let server = GuildId::from(server_tmp as u64);
|
||||
|
||||
Ok(Self {
|
||||
server,
|
||||
id_wolves: row.try_get("id_wolves")?,
|
||||
id_member: row.try_get("id_member")?,
|
||||
email: row.try_get("email")?,
|
||||
expiry: row.try_get("expiry")?,
|
||||
discord: row.try_get("discord")?,
|
||||
minecraft: row.try_get("minecraft")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Servers {
|
||||
pub server: GuildId,
|
||||
pub wolves_api: String,
|
||||
pub role_past: Option<RoleId>,
|
||||
pub role_current: Option<RoleId>,
|
||||
pub member_past: i64,
|
||||
pub member_current: i64,
|
||||
}
|
||||
impl<'r> FromRow<'r, SqliteRow> for Servers {
|
||||
fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
|
||||
let server_tmp: i64 = row.try_get("server")?;
|
||||
let server = GuildId::from(server_tmp as u64);
|
||||
let role_past = match row.try_get("role_past") {
|
||||
Ok(x) => {
|
||||
let tmp: i64 = x;
|
||||
Some(RoleId::from(tmp as u64))
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let role_current = match row.try_get("role_current") {
|
||||
Ok(x) => {
|
||||
let tmp: i64 = x;
|
||||
Some(RoleId::from(tmp as u64))
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
server,
|
||||
wolves_api: row.try_get("wolves_api")?,
|
||||
role_past,
|
||||
role_current,
|
||||
member_past: row.try_get("member_past")?,
|
||||
member_current: row.try_get("member_current")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
id_wolves text DEFAULT '',
|
||||
id_member text DEFAULT '',
|
||||
email text not null,
|
||||
expiry text not null,
|
||||
discord text,
|
||||
minecraft text,
|
||||
PRIMARY KEY(server,id_wolves,id_member)
|
||||
)",
|
||||
)
|
||||
.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_id_wolves ON accounts (id_wolves)")
|
||||
.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()
|
||||
}
|
||||
|
||||
pub 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()
|
||||
}
|
||||
|
||||
pub fn get_now_iso(short: bool) -> String {
|
||||
let now = Utc::now();
|
||||
if short {
|
||||
format!("{}-{:02}-{:02}", now.year(), now.month(), now.day())
|
||||
} else {
|
||||
now.to_rfc3339_opts(SecondsFormat::Millis, true)
|
||||
}
|
||||
}
|
265
src/main.rs
265
src/main.rs
|
@ -1,53 +1,46 @@
|
|||
use dotenvy::dotenv;
|
||||
use serenity::{
|
||||
async_trait,
|
||||
client::{Context, EventHandler},
|
||||
model::{
|
||||
gateway::{GatewayIntents, Ready},
|
||||
guild::Member,
|
||||
id::GuildId,
|
||||
prelude::RoleId,
|
||||
guild,
|
||||
},
|
||||
prelude::TypeMapKey,
|
||||
Client,
|
||||
};
|
||||
use std::{
|
||||
env,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
use skynet_discord_bot::{db_init, get_config, get_server_config, get_server_member, Config, DataBase};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
struct Handler {
|
||||
is_loop_running: AtomicBool,
|
||||
}
|
||||
struct Handler;
|
||||
|
||||
#[async_trait]
|
||||
impl EventHandler for Handler {
|
||||
async fn guild_member_addition(&self, ctx: Context, mut new_member: Member) {
|
||||
let config_lock = {
|
||||
async fn guild_member_addition(&self, ctx: Context, mut new_member: guild::Member) {
|
||||
let db_lock = {
|
||||
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 data_read = ctx.data.read().await;
|
||||
data_read.get::<Members>().expect("Expected Members in TypeMap.").clone()
|
||||
let db = db_lock.read().await;
|
||||
let config = match get_server_config(&db, &new_member.guild_id).await {
|
||||
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![];
|
||||
|
||||
if !new_member.roles.contains(&config.member_role_past) {
|
||||
roles.push(config.member_role_past);
|
||||
if let Some(role) = &config.role_past {
|
||||
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 {
|
||||
if !new_member.roles.contains(role) {
|
||||
roles.push(role.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = new_member.add_roles(&ctx, &roles).await {
|
||||
|
@ -56,171 +49,32 @@ impl EventHandler for Handler {
|
|||
}
|
||||
}
|
||||
|
||||
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 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;
|
||||
async fn ready(&self, _ctx: Context, ready: Ready) {
|
||||
println!("[Main] {} is connected!", ready.user.name);
|
||||
}
|
||||
}
|
||||
|
||||
#[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 {
|
||||
is_loop_running: AtomicBool::new(false),
|
||||
})
|
||||
.event_handler(Handler {})
|
||||
.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::<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::<DataBase>(Arc::new(RwLock::new(db)));
|
||||
}
|
||||
|
||||
// Finally, start a single shard, and start listening to events.
|
||||
|
@ -231,62 +85,3 @@ async fn main() {
|
|||
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