use crate::config::DaemonConfig; use crate::error::DaemonError; use crate::models::service::{Service, ServiceGroup}; use crate::models::user::{Permissions, User, UserGroup}; use j_db::database::Database; use j_db::model::JdbModel; use log::info; use poise::serenity_prelude::{CreateAttachment, Mentionable, MessageBuilder, UserId}; use poise::{serenity_prelude as serenity, ChoiceParameter, CreateReply}; #[allow(dead_code)] struct Data { db: Database, config: DaemonConfig, } type Error = Box; type Context<'a> = poise::Context<'a, Data, Error>; async fn check_owner(ctx: Context<'_>) -> Result { Ok(ctx.data().config.admins.contains(&ctx.author().id)) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_service( ctx: Context<'_>, #[description = "Service unit name to add"] service_name: String, ) -> Result<(), Error> { let filter = format!("*{}*", service_name); let services = systemctl::list_units_full(Some("service"), None, Some(&filter))?; if let Some(service) = services .iter() .find(|s| s.unit_file.contains(&service_name)) { let service = Service::add_service(&ctx.data().db, &service.unit_file)?; ctx.reply(format!( "`{}` was added to the list of known services.", service.name )) .await?; } else { ctx.reply(format!("`{}` was not found.", service_name)) .await?; } Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_service_group( ctx: Context<'_>, #[description = "New group name"] group_name: String, #[description = "Description of the group"] group_description: String, ) -> Result<(), Error> { ServiceGroup::add_group(&ctx.data().db, &group_name, &group_description)?; ctx.reply(format!("Added group `{}`", group_name)).await?; Ok(()) } #[poise::command(slash_command, ephemeral)] async fn add_service_to_group( ctx: Context<'_>, #[description = "Service name"] service_name: String, #[description = "Service group name"] group_name: String, ) -> Result<(), Error> { let service = Service::find_service_by_name(&ctx.data().db, &service_name)?; if let Some(service) = service { let group = ServiceGroup::find_group_by_name(&ctx.data().db, &group_name)?; if let Some(mut group) = group { group.services.insert(service.id().unwrap()); let group = ctx.data().db.insert(group)?; ctx.reply(format!("Added `{}` to `{}`", service.name, group.name)) .await?; } else { ctx.reply(format!("Unknown group `{}`", group_name)).await?; } } else { ctx.reply(format!("`{}` was not found", service_name)) .await?; } Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_user( ctx: Context<'_>, #[description = "User to add to db"] user: UserId, ) -> Result<(), Error> { let user = User::new(user); let user = ctx.data().db.insert(user)?; let discord_user = user.discord_uuid.to_user(&ctx.http()).await?; ctx.reply(format!("{} has been added.", discord_user.mention())) .await?; Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_user_group( ctx: Context<'_>, #[description = "New group name"] group_name: String, #[description = "Description of the group"] group_description: String, ) -> Result<(), Error> { let group = UserGroup::new(&group_name, &group_description); ctx.data().db.insert(group)?; ctx.reply(format!("Added group `{}`", group_name)).await?; Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_user_to_group( ctx: Context<'_>, #[description = "User"] user: UserId, #[description = "User Group"] group_name: String, ) -> Result<(), Error> { let daemon_user = User::find_user_by_id(&ctx.data().db, user)?; if let Some(user) = daemon_user { UserGroup::add_user_to_group(&ctx.data().db, &group_name, user.id().unwrap())?; ctx.reply("Added user to group.").await?; } else { let discord_user = user.to_user(&ctx.http()).await?; ctx.reply(format!( "{} is not a user in the db!", discord_user.mention() )) .await?; } Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn set_user_group_permission( ctx: Context<'_>, #[description = "User Group"] user_group: String, #[description = "Service Group"] service_group: String, #[description = "Permission"] permissions: u32, ) -> Result<(), Error> { let service_group = ServiceGroup::find_group_by_name(&ctx.data().db, &service_group)? .ok_or(DaemonError::ServiceGroupNotFound)?; UserGroup::set_group_permission( &ctx.data().db, &user_group, service_group.id().unwrap(), Permissions::from_bits(permissions).unwrap(), )?; ctx.reply("User group permission updated.").await?; Ok(()) } #[poise::command(slash_command, ephemeral, check = "check_owner")] async fn dump_db(ctx: Context<'_>) -> Result<(), Error> { let db_state = ctx.data().db.dump_db()?; let db_state = db_state.pretty(4); ctx.send( CreateReply::default() .attachment(CreateAttachment::bytes(db_state.into_bytes(), "db.json")), ) .await?; Ok(()) } #[poise::command(slash_command, ephemeral)] async fn list_services( ctx: Context<'_>, #[description = "Service group name"] group: Option, ) -> Result<(), Error> { let group = if let Some(group) = group { ServiceGroup::find_group_by_name(&ctx.data().db, &group)? } else { None }; let user = User::find_user_by_id(&ctx.data().db, ctx.author().id)? .ok_or(DaemonError::UserNotAuthorized)?; let services: Vec = ctx .data() .db .filter(|_, s: &Service| { UserGroup::check_if_user_has_permission( &ctx.data().db, user.id().unwrap(), s.id().unwrap(), Permissions::Status, ) .unwrap_or(false) })? .filter(|s: &Service| { if let Some(group) = &group { group.services.contains(&s.id().unwrap()) } else { true } }) .map(|s: Service| s.name) .collect(); let mut message_builder = MessageBuilder::new(); if let Some(group) = group { message_builder.push_line(format!("## `{}` Services", group.name)); message_builder.push_line(group.description); } else { message_builder.push_line("## Services"); } for service in &services { let status = systemctl::is_active(service)?; let status_emoji = if status { ":green_circle:" } else { ":red_circle:" }; message_builder.push_line(format!("* {} **{}**", status_emoji, service)); } ctx.reply(message_builder.build()).await.unwrap(); Ok(()) } #[derive(Debug, Clone, ChoiceParameter)] pub enum ServiceAction { Start, Stop, Restart, Status, } impl ServiceAction { pub fn permission(&self) -> Permissions { match self { ServiceAction::Start => Permissions::Start, ServiceAction::Stop => Permissions::Stop, ServiceAction::Restart => Permissions::Restart, ServiceAction::Status => Permissions::Status, } } } #[poise::command(slash_command, ephemeral)] async fn service( ctx: Context<'_>, #[description = "Action to preform on service"] action: ServiceAction, #[description = "Service unit name"] service_name: String, ) -> Result<(), Error> { let service: Option = Service::find_service_by_name(&ctx.data().db, &service_name)?; let user = User::find_user_by_id(&ctx.data().db, ctx.author().id)? .ok_or(DaemonError::UserNotAuthorized)?; if let Some(service) = service { let check_perm = UserGroup::check_if_user_has_permission( &ctx.data().db, user.id().unwrap(), service.id().unwrap(), action.permission(), )?; if !check_perm { return Err(Error::from(DaemonError::UserNotAuthorized)); } match action { ServiceAction::Start => { systemctl::start(&service_name)?; ctx.reply(format!("`{}` has been started", service_name)) .await?; } ServiceAction::Stop => { systemctl::stop(&service_name)?; ctx.reply(format!("`{}` has been stopped", service_name)) .await?; } ServiceAction::Restart => { systemctl::stop(&service_name)?; ctx.reply(format!("`{}` has been restarted", service_name)) .await?; } ServiceAction::Status => { let status = systemctl::status(&service.name)?; let mut msg = MessageBuilder::new(); msg.push_codeblock_safe(status, None); ctx.reply(msg.build()).await?; } } } else { ctx.reply(format!("Unknown service `{}`", service_name)) .await?; } Ok(()) } pub async fn run_bot(db: Database, config: DaemonConfig) { let intents = serenity::GatewayIntents::non_privileged(); let data = Data { db, config: config.clone(), }; let framework = poise::Framework::builder() .options(poise::FrameworkOptions { commands: vec![ list_services(), add_service(), add_service_group(), add_service_to_group(), service(), add_user(), add_user_group(), add_user_to_group(), dump_db(), set_user_group_permission(), ], ..Default::default() }) .setup(|ctx, _ready, framework| { Box::pin(async move { poise::builtins::register_globally(ctx, &framework.options().commands).await?; Ok(data) }) }) .build(); let client = serenity::ClientBuilder::new(config.discord_api_token, intents) .framework(framework) .await; info!("Starting bot..."); client.unwrap().start().await.unwrap(); }