From c529e6b2ac7a3f4a03c3d4276f07e819bdb3d932 Mon Sep 17 00:00:00 2001 From: Joey Hines Date: Sat, 9 Mar 2024 16:31:34 -0700 Subject: [PATCH] Added group functionality --- Cargo.lock | 2 ++ Cargo.toml | 2 ++ src/discord/mod.rs | 82 +++++++++++++++++++++++++++++++++++++++---- src/error.rs | 7 ++++ src/main.rs | 1 + src/models/service.rs | 79 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 166 insertions(+), 7 deletions(-) create mode 100644 src/error.rs diff --git a/Cargo.lock b/Cargo.lock index 594b8d2..6855c0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -411,10 +411,12 @@ dependencies = [ "config", "env_logger", "j_db", + "log", "poise", "serde", "structopt", "systemctl", + "thiserror", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index 53aaa3c..737733f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,8 @@ j_db = { version = "0.1.2", registry = "jojo-dev" } serde = "1.0.197" structopt = "0.3.26" env_logger = "0.11.3" +thiserror = "1.0.57" +log = "0.4.21" [dependencies.tokio] version = "1.36.0" diff --git a/src/discord/mod.rs b/src/discord/mod.rs index 2fb83b3..3c0acc9 100644 --- a/src/discord/mod.rs +++ b/src/discord/mod.rs @@ -1,6 +1,8 @@ use crate::config::DaemonConfig; -use crate::models::service::Service; +use crate::models::service::{Service, ServiceGroup}; use j_db::database::Database; +use j_db::model::JdbModel; +use log::info; use poise::serenity_prelude as serenity; use poise::serenity_prelude::MessageBuilder; @@ -24,12 +26,11 @@ async fn add_service( .iter() .find(|s| s.unit_file.contains(&service_name)) { - let service = Service::new(&service.unit_file); + let service = Service::add_service(&ctx.data().db, &service.unit_file)?; - ctx.data().db.insert(service)?; ctx.reply(format!( "`{}` was added to the list of known services.", - service_name + service.name )) .await?; } else { @@ -41,16 +42,80 @@ async fn add_service( } #[poise::command(slash_command)] -async fn list_services(ctx: Context<'_>) -> Result<(), Error> { - let mut message_builder = serenity::MessageBuilder::new(); +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)] +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)] +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 services: Vec = ctx .data() .db .filter(|_, _s: &Service| true)? + .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 = serenity::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 { @@ -58,7 +123,7 @@ async fn list_services(ctx: Context<'_>) -> Result<(), Error> { } else { ":red_circle:" }; - message_builder.push_line(format!("* **{}** {}", service, status_emoji)); + message_builder.push_line(format!("* {} **{}**", status_emoji, service)); } ctx.reply(message_builder.build()).await.unwrap(); @@ -128,8 +193,10 @@ pub async fn run_bot(db: Database, config: DaemonConfig) { commands: vec![ list_services(), add_service(), + add_service_group(), restart_service(), service_status(), + add_service_to_group(), ], ..Default::default() }) @@ -144,5 +211,6 @@ pub async fn run_bot(db: Database, config: DaemonConfig) { let client = serenity::ClientBuilder::new(config.discord_api_token, intents) .framework(framework) .await; + info!("Starting bot..."); client.unwrap().start().await.unwrap(); } diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..dd9dd46 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,7 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum DaemonError { + #[error("DB error")] + DBError(#[from] j_db::error::JDbError), +} diff --git a/src/main.rs b/src/main.rs index 9c3478b..27fb0f5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ use structopt::StructOpt; mod config; mod discord; +mod error; mod models; #[tokio::main] diff --git a/src/models/service.rs b/src/models/service.rs index cfd1cc8..12a17ba 100644 --- a/src/models/service.rs +++ b/src/models/service.rs @@ -1,4 +1,6 @@ +use j_db::database::Database; use serde::{Deserialize, Serialize}; +use std::collections::HashSet; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Service { @@ -14,6 +16,21 @@ impl Service { id: None, } } + + pub fn add_service(db: &Database, name: &str) -> Result { + let service = Self::new(name); + db.insert(service) + } + + pub fn find_service_by_name( + db: &Database, + name: &str, + ) -> Result, j_db::error::JDbError> { + Ok(db + .filter(|_, service: &Service| service.name.contains(name))? + .next() + .clone()) + } } impl j_db::model::JdbModel for Service { @@ -28,4 +45,66 @@ impl j_db::model::JdbModel for Service { fn tree() -> String { "Service".to_string() } + + fn check_unique(&self, other: &Self) -> bool { + !self.name.eq_ignore_ascii_case(&other.name) + } +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct ServiceGroup { + pub name: String, + pub description: String, + pub services: HashSet, + + id: Option, +} + +impl ServiceGroup { + pub fn new(name: &str, description: &str) -> Self { + Self { + name: name.to_string(), + description: description.to_string(), + services: HashSet::new(), + id: None, + } + } + + pub fn add_group( + db: &Database, + name: &str, + description: &str, + ) -> Result { + let service_group = Self::new(name, description); + + db.insert(service_group) + } + + pub fn find_group_by_name( + db: &Database, + name: &str, + ) -> Result, j_db::error::JDbError> { + Ok(db + .filter(|_, group: &ServiceGroup| group.name.eq_ignore_ascii_case(name))? + .next() + .clone()) + } +} + +impl j_db::model::JdbModel for ServiceGroup { + fn id(&self) -> Option { + self.id + } + + fn set_id(&mut self, id: u64) { + self.id = Some(id) + } + + fn tree() -> String { + "ServiceGroup".to_string() + } + + fn check_unique(&self, other: &Self) -> bool { + !self.name.eq_ignore_ascii_case(&other.name) + } }