From 79a3188095777917b2747d7dfb0080077b47ca71 Mon Sep 17 00:00:00 2001 From: Joey Hines Date: Fri, 8 Mar 2024 20:04:53 -0700 Subject: [PATCH] Initial commit + Basic functionality w/ no permission --- .cargo/config.toml | 2 + src/config.rs | 29 +++++++++ src/discord/mod.rs | 148 ++++++++++++++++++++++++++++++++++++++++++ src/models/mod.rs | 1 + src/models/service.rs | 31 +++++++++ 5 files changed, 211 insertions(+) create mode 100644 .cargo/config.toml create mode 100644 src/config.rs create mode 100644 src/discord/mod.rs create mode 100644 src/models/mod.rs create mode 100644 src/models/service.rs diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..6629c6e --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[registries.jojo-dev] +index = "https://git.jojodev.com/joeyahines/_cargo-index.git" \ No newline at end of file diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..3f33356 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,29 @@ +use config::Config; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use structopt::StructOpt; + +#[derive(StructOpt)] +pub struct Args { + pub config_path: PathBuf, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DaemonConfig { + pub discord_api_token: String, + pub db_path: PathBuf, +} + +impl DaemonConfig { + pub fn new(config: PathBuf) -> DaemonConfig { + let daemon_config = Config::builder() + .add_source(config::File::new( + config.to_str().unwrap(), + config::FileFormat::Toml, + )) + .build() + .unwrap(); + + daemon_config.try_deserialize().unwrap() + } +} diff --git a/src/discord/mod.rs b/src/discord/mod.rs new file mode 100644 index 0000000..2fb83b3 --- /dev/null +++ b/src/discord/mod.rs @@ -0,0 +1,148 @@ +use crate::config::DaemonConfig; +use crate::models::service::Service; +use j_db::database::Database; +use poise::serenity_prelude as serenity; +use poise::serenity_prelude::MessageBuilder; + +#[allow(dead_code)] +struct Data { + db: Database, + config: DaemonConfig, +} +type Error = Box; +type Context<'a> = poise::Context<'a, Data, Error>; + +#[poise::command(slash_command)] +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::new(&service.unit_file); + + ctx.data().db.insert(service)?; + 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)] +async fn list_services(ctx: Context<'_>) -> Result<(), Error> { + let mut message_builder = serenity::MessageBuilder::new(); + + let services: Vec = ctx + .data() + .db + .filter(|_, _s: &Service| true)? + .map(|s: Service| s.name) + .collect(); + + 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!("* **{}** {}", service, status_emoji)); + } + + ctx.reply(message_builder.build()).await.unwrap(); + + Ok(()) +} + +#[poise::command(slash_command)] +async fn restart_service( + ctx: Context<'_>, + #[description = "Service unit name"] service_name: String, +) -> Result<(), Error> { + let service: Option = ctx + .data() + .db + .filter(|_, s: &Service| s.name.contains(&service_name))? + .next(); + + if let Some(service) = service { + systemctl::restart(&service.name)?; + ctx.reply(format!("`{}` has been restarted", service_name)) + .await?; + } else { + ctx.reply(format!("Unknown service `{}`", service_name)) + .await?; + } + + Ok(()) +} + +#[poise::command(slash_command)] +async fn service_status( + ctx: Context<'_>, + #[description = "Service unit name"] service_name: String, +) -> Result<(), Error> { + let service: Option = ctx + .data() + .db + .filter(|_, s: &Service| s.name.contains(&service_name))? + .next(); + + if let Some(service) = service { + 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(), + restart_service(), + service_status(), + ], + ..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; + client.unwrap().start().await.unwrap(); +} diff --git a/src/models/mod.rs b/src/models/mod.rs new file mode 100644 index 0000000..1f278a4 --- /dev/null +++ b/src/models/mod.rs @@ -0,0 +1 @@ +pub mod service; diff --git a/src/models/service.rs b/src/models/service.rs new file mode 100644 index 0000000..cfd1cc8 --- /dev/null +++ b/src/models/service.rs @@ -0,0 +1,31 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct Service { + pub name: String, + + id: Option, +} + +impl Service { + pub fn new(name: &str) -> Self { + Self { + name: name.to_string(), + id: None, + } + } +} + +impl j_db::model::JdbModel for Service { + fn id(&self) -> Option { + self.id + } + + fn set_id(&mut self, id: u64) { + self.id = Some(id) + } + + fn tree() -> String { + "Service".to_string() + } +}