daemon/src/discord/mod.rs

use crate::config::DaemonConfig;
use crate::models::service::{Service, ServiceGroup};
use j_db::database::Database;
use j_db::model::JdbModel;
use log::info;
use poise::serenity_prelude::MessageBuilder;
use poise::{serenity_prelude as serenity, ChoiceParameter};

#[allow(dead_code)]
struct Data {
    db: Database,
    config: DaemonConfig,
}
type Error = Box<dyn std::error::Error + Send + Sync>;
type Context<'a> = poise::Context<'a, Data, Error>;

#[poise::command(slash_command, ephemeral)]
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)]
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)]
async fn list_services(
    ctx: Context<'_>,
    #[description = "Service group name"] group: Option<String>,
) -> Result<(), Error> {
    let group = if let Some(group) = group {
        ServiceGroup::find_group_by_name(&ctx.data().db, &group)?
    } else {
        None
    };

    let services: Vec<String> = 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 {
            ":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,
}

#[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> = Service::find_service_by_name(&ctx.data().db, &service_name)?;

    if let Some(service) = service {
        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(),
            ],
            ..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();
}