daemon/src/discord/mod.rs

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<dyn std::error::Error + Send + Sync>;
type Context<'a> = poise::Context<'a, Data, Error>;

async fn check_owner(ctx: Context<'_>) -> Result<bool, Error> {
    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<String>,
) -> 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<String> = 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> = 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();
}