Added back in admin commands

This commit is contained in:
Joey Hines 2025-03-15 20:39:50 -06:00
parent 6ac5129114
commit 05591d49e3
Signed by: joeyahines
GPG Key ID: 38BA6F25C94C9382
6 changed files with 166 additions and 255 deletions

View File

@ -14,7 +14,7 @@ use std::path::{Path, PathBuf};
use std::sync::Arc; use std::sync::Arc;
use structopt::StructOpt; use structopt::StructOpt;
use tokio::sync::mpsc::{channel, Receiver, Sender}; use tokio::sync::mpsc::{channel, Receiver, Sender};
use tokio::sync::{Mutex, MutexGuard}; use tokio::sync::Mutex;
#[derive(Debug, StructOpt)] #[derive(Debug, StructOpt)]
#[structopt(name = "fren", about = "Friend Bot")] #[structopt(name = "fren", about = "Friend Bot")]
@ -104,14 +104,6 @@ impl GlobalData {
picox: AlbumManager::new(cfg.picox.api_base_url, &cfg.picox.token), picox: AlbumManager::new(cfg.picox.api_base_url, &cfg.picox.token),
}) })
} }
pub async fn reload(&mut self) -> Result<(), Error> {
let cfg = BotConfig::new(&self.args.cfg_path)?;
self.cfg = cfg;
Ok(())
}
} }
impl TypeMapKey for GlobalData { impl TypeMapKey for GlobalData {

View File

@ -1,4 +1,6 @@
use crate::config::BotConfig; use crate::config::BotConfig;
use crate::discord::Context;
use crate::error::Error;
use crate::inventory::ItemType; use crate::inventory::ItemType;
use crate::models::api_key::Apikey; use crate::models::api_key::Apikey;
use crate::models::lil_fren::{ use crate::models::lil_fren::{
@ -7,69 +9,30 @@ use crate::models::lil_fren::{
}; };
use crate::models::task::Task; use crate::models::task::Task;
use crate::user::User; use crate::user::User;
use crate::{command, group, GlobalData};
use json::JsonValue; use json::JsonValue;
use serenity::all::{CreateAttachment, CreateMessage, FormattedTimestamp, FormattedTimestampStyle}; use log::info;
use serenity::client::Context; use poise::serenity_prelude::{
use serenity::framework::standard::{Args, CommandResult}; Attachment, CreateAttachment, CreateMessage, EmojiIdentifier, FormattedTimestamp,
use serenity::model::channel::Message; FormattedTimestampStyle, MessageBuilder, Timestamp, UserId,
use serenity::model::misc::EmojiIdentifier; };
use serenity::model::prelude::UserId;
use serenity::model::Timestamp;
use serenity::utils::MessageBuilder;
use std::borrow::Cow; use std::borrow::Cow;
#[group] pub async fn is_admin(ctx: Context<'_>) -> Result<bool, Error> {
#[commands( Ok(ctx.data().cfg.admins.contains(&ctx.author().id))
reload,
dump_db,
load_db,
add_key,
debug_buddy,
draw_buddy_states,
op_give,
list_tasks
)]
pub struct ADMIN;
pub fn is_admin(user_id: &UserId, cfg: &BotConfig) -> bool {
cfg.admins.contains(user_id)
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
async fn reload(ctx: &Context, msg: &Message, _args: Args) -> CommandResult { pub async fn dump_db(ctx: Context<'_>) -> Result<(), Error> {
let mut data = ctx.data.write().await; let db_dump = ctx.data().db.dump_db()?;
let global_data = data.get_mut::<GlobalData>().unwrap();
if !is_admin(&msg.author.id, &global_data.cfg) {
return Ok(());
}
global_data.reload().await?;
msg.reply(&ctx.http, "Reload done ;)").await?;
Ok(())
}
#[command]
async fn dump_db(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
let mut data = ctx.data.write().await;
let global_data = data.get_mut::<GlobalData>().unwrap();
if !is_admin(&msg.author.id, &global_data.cfg) {
return Ok(());
}
let db_dump = global_data.db.dump_db()?;
let output = db_dump.pretty(4); let output = db_dump.pretty(4);
msg.author ctx.author()
.id .id
.create_dm_channel(&ctx.http) .create_dm_channel(ctx.http())
.await? .await?
.send_message( .send_message(
&ctx.http, ctx.http(),
CreateMessage::new() CreateMessage::new()
.content("The current DB state") .content("The current DB state")
.add_file(CreateAttachment::bytes( .add_file(CreateAttachment::bytes(
@ -82,67 +45,56 @@ async fn dump_db(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
async fn load_db(ctx: &Context, msg: &Message, _args: Args) -> CommandResult { pub async fn load_db(
let mut data = ctx.data.write().await; ctx: Context<'_>,
let global_data = data.get_mut::<GlobalData>().unwrap(); #[description = "DB import file"] attachment: Attachment,
) -> Result<(), Error> {
let db_bytes = attachment.download().await?;
let db_string = String::from_utf8(db_bytes).unwrap();
if !is_admin(&msg.author.id, &global_data.cfg) { let json_value: JsonValue = match json::parse(&db_string) {
return Ok(()); Ok(v) => v,
} Err(e) => {
ctx.reply(format!("Error parsing json: {}", e)).await?;
return Ok(());
}
};
if let Some(attachment) = msg.attachments.first() { info!("Importing DB from user supplied file");
let db_bytes = attachment.download().await?; match ctx.data().db.import_db(json_value) {
let db_string = String::from_utf8(db_bytes)?; Ok(_) => ctx.reply("Database imported successfully").await?,
Err(err) => {
let json_value: JsonValue = match json::parse(&db_string) { ctx.reply(format!("Error importing db: {}", err)).await?;
Ok(v) => v, return Ok(());
Err(e) => { }
msg.reply(&ctx.http, format!("Error parsing json: {}", e)) };
.await?;
return Ok(());
}
};
match global_data.db.import_db(json_value) {
Ok(_) => {
msg.reply(&ctx.http, "Database imported successfully")
.await?
}
Err(err) => {
msg.reply(&ctx.http, format!("Error importing db: {}", err))
.await?;
return Ok(());
}
};
}
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
async fn add_key(ctx: &Context, msg: &Message, args: Args) -> CommandResult { pub async fn add_key(
let mut data = ctx.data.write().await; ctx: Context<'_>,
let global_data = data.get_mut::<GlobalData>().unwrap(); #[description = "User to make a key for"] user_id: Option<UserId>,
) -> Result<(), Error> {
if !is_admin(&msg.author.id, &global_data.cfg) { let (api_key, key) = if let Some(user_id) = user_id {
return Ok(()); let user = user_id.to_user(&ctx.http()).await?;
}
let (api_key, key) = if args.len() == 1 {
let user_id = args.parse::<UserId>()?;
let user = user_id.to_user(&ctx.http).await?;
Apikey::new(&format!("{}'s Key", user.name), Some(user_id)) Apikey::new(&format!("{}'s Key", user.name), Some(user_id))
} else { } else {
Apikey::new(&format!("{}'s Key", msg.author.name), Some(msg.author.id)) Apikey::new(
&format!("{}'s Key", ctx.author().name),
Some(ctx.author().id),
)
}; };
global_data.db.insert::<Apikey>(api_key.clone())?; info!("Adding key '{}'", api_key.name);
ctx.data().db.insert::<Apikey>(api_key.clone())?;
let dm = msg.author.create_dm_channel(&ctx.http).await?; let dm = ctx.author().create_dm_channel(&ctx.http()).await?;
dm.say( dm.say(
&ctx.http, ctx.http(),
format!("Key '{}' added. Api Key: {}", api_key.name, key), format!("Key '{}' added. Api Key: {}", api_key.name, key),
) )
.await?; .await?;
@ -150,108 +102,76 @@ async fn add_key(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
#[description("Check little buddy stats")] pub async fn debug_buddy(ctx: Context<'_>) -> Result<(), Error> {
async fn debug_buddy(ctx: &Context, msg: &Message, _args: Args) -> CommandResult { let lil_fren = LilFren::get_lil_fren(&ctx.data().db)?;
let mut data = ctx.data.write().await;
let global_data = data.get_mut::<GlobalData>().unwrap();
if !is_admin(&msg.author.id, &global_data.cfg) {
return Ok(());
}
let lil_fren = LilFren::get_lil_fren(&global_data.db)?;
if let Some(lil_fren) = lil_fren { if let Some(lil_fren) = lil_fren {
if lil_fren.is_alive() == AliveState::Alive { if lil_fren.is_alive() == AliveState::Alive {
msg.reply( ctx.reply(format!(
&ctx.http, "Hunger: {} Thirst: {} Entertainment: {} Smarts: {} Metabolism: {} State: {:?}",
format!( lil_fren.hunger,
"Hunger: {} Thirst: {} Entertainment: {} Smarts: {} Metabolism: {} State: {:?}", lil_fren.thirst,
lil_fren.hunger, lil_fren.entertainment,
lil_fren.thirst, lil_fren.smarts,
lil_fren.entertainment, lil_fren.metabolism,
lil_fren.smarts, lil_fren.state
lil_fren.metabolism, ))
lil_fren.state
),
)
.await?; .await?;
} else { } else {
msg.reply(&ctx.http, "Lil buddy is dead!").await?; ctx.reply("Lil buddy is dead!").await?;
} }
} else { } else {
msg.reply(&ctx.http, "Sorry no little buddy found!").await?; ctx.reply("Sorry no little buddy found!").await?;
} }
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
#[description("Check little buddy stats")] pub async fn draw_buddy_states(
async fn draw_buddy_states(ctx: &Context, msg: &Message, args: Args) -> CommandResult { ctx: Context<'_>,
let emoji = args.parse::<EmojiIdentifier>()?; #[description = "Emoji to use for debugging"] emoji: EmojiIdentifier,
) -> Result<(), Error> {
let emoji = ctx.guild_id().unwrap().emoji(ctx.http(), emoji.id).await?;
let emoji = msg.guild_id.unwrap().emoji(&ctx.http, emoji.id).await?; ctx.reply(draw_standing(&emoji)).await?;
ctx.reply(draw_dancing(&emoji)).await?;
msg.reply(&ctx.http, draw_standing(&emoji)).await?; ctx.reply(draw_frankenstein(&emoji)).await?;
msg.reply(&ctx.http, draw_dancing(&emoji)).await?; ctx.reply(draw_sick(&emoji)).await?;
msg.reply(&ctx.http, draw_frankenstein(&emoji)).await?; ctx.reply(draw_magic(&emoji)).await?;
msg.reply(&ctx.http, draw_sick(&emoji)).await?; ctx.reply(draw_resonance_cascade(&emoji)).await?;
msg.reply(&ctx.http, draw_magic(&emoji)).await?; ctx.reply(draw_tax_fraud(&emoji)).await?;
msg.reply(&ctx.http, draw_resonance_cascade(&emoji)).await?; ctx.reply(draw_sleep(&emoji)).await?;
msg.reply(&ctx.http, draw_tax_fraud(&emoji)).await?; ctx.reply(draw_mining(&emoji)).await?;
msg.reply(&ctx.http, draw_sleep(&emoji)).await?; ctx.reply(draw_gone()).await?;
msg.reply(&ctx.http, draw_mining(&emoji)).await?;
msg.reply(&ctx.http, draw_gone()).await?;
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin", check = "is_admin")]
#[description("Hey can I get op? I'm from planet minecraft")] pub async fn op_give(
async fn op_give(ctx: &Context, msg: &Message, args: Args) -> CommandResult { ctx: Context<'_>,
let mut data = ctx.data.write().await; #[description = "Item to give"]
let global_data = data.get_mut::<GlobalData>().unwrap(); #[rest]
item: ItemType,
) -> Result<(), Error> {
info!("'{}' has been given a {:?}", ctx.author().name, item);
User::give_item(&ctx.data().db, ctx.author().id, item, 1, None)?;
if args.is_empty() { ctx.reply(format!(
msg.reply(&ctx.http, "You need to select an item to give") "[Console] Op has given {} 1 {}",
.await?; ctx.author().name,
return Ok(()); item
} ))
if !is_admin(&msg.author.id, &global_data.cfg) {
return Ok(());
}
let msg_content = args.rest().to_lowercase();
let item = match msg_content.parse::<ItemType>() {
Ok(i) => i,
Err(_) => {
msg.reply(&ctx.http, "I don't know what the heck that is tbh.")
.await?;
return Ok(());
}
};
User::give_item(&global_data.db, msg.author.id, item, 1, None)?;
msg.reply(
&ctx,
format!("[Console] Op has given {} 1 {}", msg.author.name, item),
)
.await?; .await?;
Ok(()) Ok(())
} }
#[command] #[poise::command(prefix_command, category = "Admin")]
#[description("List scheduled tasks")] pub async fn list_tasks(ctx: Context<'_>) -> Result<(), Error> {
async fn list_tasks(ctx: &Context, msg: &Message) -> CommandResult { let tasks: Vec<Task> = ctx.data().db.filter(|_, _task: &Task| true)?.collect();
let data = ctx.data.read().await;
let global_data = data.get::<GlobalData>().unwrap();
let tasks: Vec<Task> = global_data.db.filter(|_, _task: &Task| true)?.collect();
let mut resp = MessageBuilder::new(); let mut resp = MessageBuilder::new();
@ -265,7 +185,7 @@ async fn list_tasks(ctx: &Context, msg: &Message) -> CommandResult {
)); ));
} }
msg.reply(&ctx.http, resp.build()).await?; ctx.reply(resp.build()).await?;
Ok(()) Ok(())
} }

View File

@ -1,3 +1,4 @@
mod admin;
mod celeryman; mod celeryman;
use crate::config::GlobalData; use crate::config::GlobalData;
@ -5,7 +6,7 @@ use crate::error::Error;
use poise::{serenity_prelude as serenity, FrameworkOptions}; use poise::{serenity_prelude as serenity, FrameworkOptions};
type Context<'a> = poise::Context<'a, GlobalData, Error>; pub type Context<'a> = poise::Context<'a, GlobalData, Error>;
pub async fn run_bot(global_data: GlobalData) { pub async fn run_bot(global_data: GlobalData) {
let framework_options: FrameworkOptions<GlobalData, Error> = poise::FrameworkOptions { let framework_options: FrameworkOptions<GlobalData, Error> = poise::FrameworkOptions {
@ -20,6 +21,14 @@ pub async fn run_bot(global_data: GlobalData) {
celeryman::nudetayne(), celeryman::nudetayne(),
celeryman::celeryman(), celeryman::celeryman(),
celeryman::tayne(), celeryman::tayne(),
admin::dump_db(),
admin::load_db(),
admin::add_key(),
admin::debug_buddy(),
admin::draw_buddy_states(),
admin::debug_buddy(),
admin::op_give(),
admin::list_tasks(),
], ],
..Default::default() ..Default::default()
}; };

View File

@ -20,22 +20,35 @@ impl Display for InventoryError {
} }
} }
impl std::error::Error for InventoryError {}
#[derive(Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq)] #[derive(Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq)]
#[serde(tag = "type", content = "data")] #[serde(tag = "type", content = "data")]
pub enum ItemData { pub enum ItemData {
Nft(String), Nft(String),
} }
#[derive(Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq, Copy)] #[derive(
Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq, Copy, poise::ChoiceParameter,
)]
pub enum ItemType { pub enum ItemType {
#[name = "Cancel Insurance"]
CancelInsurance, CancelInsurance,
#[name = "The Concept of Love"]
TheConceptOfLove, TheConceptOfLove,
#[name = "Good Fortune"]
GoodFortune, GoodFortune,
#[name = "NFT"]
Nft, Nft,
#[name = "License to be Horny"]
LicenseToBeHorny, LicenseToBeHorny,
#[name = "Kill Gun"]
KillGun, KillGun,
#[name = "Cancel Ray"]
CancelRay, CancelRay,
#[name = "Helmet"]
Helmet, Helmet,
#[name = "Tactical Nuke"]
TacticalNuke, TacticalNuke,
} }

View File

@ -4,4 +4,4 @@ pub mod insult_compliment;
pub mod lil_fren; pub mod lil_fren;
pub mod motivation; pub mod motivation;
pub mod random; pub mod random;
//pub mod task; pub mod task;

View File

@ -1,14 +1,15 @@
use crate::config::GlobalData; use crate::config::GlobalData;
//use crate::discord::shop::restock_shop; //use crate::discord::shop::restock_shop;
use crate::discord::Context;
use crate::error::Error; use crate::error::Error;
use crate::models::birthday::BirthdayEntry; use crate::models::birthday::BirthdayEntry;
use crate::models::insult_compliment::{RandomResponseTemplate, ResponseType}; use crate::models::insult_compliment::{RandomResponseTemplate, ResponseType};
use chrono::{Days, Duration, TimeZone, Timelike, Utc}; use chrono::{Days, Duration, TimeZone, Timelike, Utc};
use j_db::database::Database; use j_db::database::Database;
use j_db::model::JdbModel; use j_db::model::JdbModel;
use serde::{Deserialize, Serialize}; use log::{error, info};
use poise::serenity_prelude::all::Mentionable; use poise::serenity_prelude::all::Mentionable;
use poise:: use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq)] #[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq)]
pub enum TaskType { pub enum TaskType {
@ -79,14 +80,11 @@ impl Task {
Ok(()) Ok(())
} }
pub async fn create_reoccurring_tasks(ctx: &Context) -> Result<(), Error> { pub async fn create_reoccurring_tasks(ctx: Context<'_>) -> Result<(), Error> {
let mut data = ctx.data.write().await; Task::add_task(&ctx.data().db, TaskType::CheckBirthdays, Utc::now())?;
let global_data = data.get_mut::<GlobalData>().unwrap(); Task::add_task(&ctx.data().db, TaskType::RestockShop, Utc::now())?;
Task::add_task(&global_data.db, TaskType::CheckBirthdays, Utc::now())?;
Task::add_task(&global_data.db, TaskType::RestockShop, Utc::now())?;
Task::add_task( Task::add_task(
&global_data.db, &ctx.data().db,
TaskType::HandleReload, TaskType::HandleReload,
Utc::now() + Duration::hours(1), Utc::now() + Duration::hours(1),
)?; )?;
@ -94,63 +92,54 @@ impl Task {
Ok(()) Ok(())
} }
pub async fn run_tasks(ctx: &Context) -> Result<(), Error> { pub async fn run_tasks(ctx: Context<'_>) -> Result<(), Error> {
let mut data = ctx.data.write().await; let active_tasks: Vec<Task> = ctx
let global_data = data.get_mut::<GlobalData>().unwrap(); .data()
let active_tasks: Vec<Task> = global_data
.db .db
.filter(|_, task: &Task| task.time < chrono::Utc::now())? .filter(|_, task: &Task| task.time < Utc::now())?
.collect(); .collect();
for task in active_tasks { for task in active_tasks {
match task.task_type { match task.task_type {
TaskType::RemoveRole { user_id, role_id } => { TaskType::RemoveRole { user_id, role_id } => {
let user = global_data let user = ctx.data().cfg.guild_id.member(ctx.http(), user_id).await?;
.cfg
.guild_id
.member(&ctx.http, user_id)
.await
.unwrap();
println!("Removing role {} from {}", role_id, user.display_name()); info!("Removing role {} from {}", role_id, user.display_name());
user.remove_role(&ctx.http, role_id).await.unwrap(); user.remove_role(ctx.http(), role_id).await?;
} }
TaskType::CheckBirthdays => { TaskType::CheckBirthdays => {
println!("Checking Birthdays"); info!("Checking Birthdays");
let todays_birthdays = BirthdayEntry::todays_birthdays( let todays_birthdays = BirthdayEntry::todays_birthdays(
&global_data.db, &ctx.data().db,
chrono::Utc::now().date_naive(), chrono::Utc::now().date_naive(),
) )?;
.unwrap();
for birth in todays_birthdays { for birth in todays_birthdays {
if let Ok(user) = global_data if let Ok(user) = ctx
.data()
.cfg .cfg
.guild_id .guild_id
.member(&ctx.http, birth.discord_id) .member(ctx.http(), birth.discord_id)
.await .await
{ {
global_data ctx.data()
.cfg .cfg
.announcement_channel .announcement_channel
.say(&ctx.http, format!("Happy birthday {}!", user.mention())) .say(ctx.http(), format!("Happy birthday {}!", user.mention()))
.await .await?;
.unwrap();
let compliment = RandomResponseTemplate::get_random_response( let compliment = RandomResponseTemplate::get_random_response(
&global_data.db, &ctx.data().db,
ResponseType::Compliment, ResponseType::Compliment,
user.display_name(), user.display_name(),
) )?
.unwrap() .unwrap_or("I couldn't think of anything funny tbh...".to_string());
.unwrap();
global_data ctx.data()
.cfg .cfg
.announcement_channel .announcement_channel
.say(&ctx.http, compliment) .say(ctx.http(), compliment)
.await .await?;
.unwrap();
} }
} }
@ -166,44 +155,32 @@ impl Task {
.unwrap(); .unwrap();
Task::add_task( Task::add_task(
&global_data.db, &ctx.data().db,
TaskType::CheckBirthdays, TaskType::CheckBirthdays,
next_check.with_timezone(&Utc), next_check.with_timezone(&Utc),
)?; )?;
} }
TaskType::HandleReload => { TaskType::HandleReload => {
println!("Reloading config..."); // Do nothing for now
let res = global_data.reload().await;
match res {
Ok(_) => println!("Finished reloading config!"),
Err(err) => println!("Error reloading config: {:?}", err),
}
Task::add_task(
&global_data.db,
TaskType::HandleReload,
Utc::now() + Duration::hours(1),
)?;
} }
TaskType::RestockShop => { TaskType::RestockShop => {
println!("Restocking Shop..."); info!("Restocking Shop...");
let res = restock_shop(ctx, global_data).await; //let res = restock_shop(ctx, global_data).await;
match res { //match res {
Ok(_) => println!("Finished restocking shop!"), // Ok(_) => info!("Finished restocking shop!"),
Err(err) => println!("Error restocking shop: {:?}", err), // Err(err) => error!("Error restocking shop: {:?}", err),
} // }
Task::add_task( Task::add_task(
&global_data.db, &ctx.data().db,
TaskType::RestockShop, TaskType::RestockShop,
Utc::now() + Duration::hours(1), Utc::now() + Duration::hours(1),
)?; )?;
} }
} }
let _ = global_data.db.remove::<Task>(task.id().unwrap()).is_ok(); let _ = ctx.data().db.remove::<Task>(task.id().unwrap()).is_ok();
} }
Ok(()) Ok(())