From 10a69052630de3656cdba568c664562280b5ffc7 Mon Sep 17 00:00:00 2001 From: Joey Hines Date: Sun, 10 Mar 2024 16:53:24 -0600 Subject: [PATCH] Added permission handling and MVP status!!!! --- Cargo.lock | 1 + Cargo.toml | 1 + src/config.rs | 2 + src/discord/mod.rs | 147 +++++++++++++++++++++++++++++-- src/error.rs | 7 ++ src/models/mod.rs | 1 + src/models/service.rs | 9 ++ src/models/user.rs | 195 ++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 357 insertions(+), 6 deletions(-) create mode 100644 src/models/user.rs diff --git a/Cargo.lock b/Cargo.lock index 6855c0b..4d4c4ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -408,6 +408,7 @@ dependencies = [ name = "daemon" version = "0.1.0" dependencies = [ + "bitflags 2.4.2", "config", "env_logger", "j_db", diff --git a/Cargo.toml b/Cargo.toml index 737733f..ab1a4b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ structopt = "0.3.26" env_logger = "0.11.3" thiserror = "1.0.57" log = "0.4.21" +bitflags = "2.4.2" [dependencies.tokio] version = "1.36.0" diff --git a/src/config.rs b/src/config.rs index 3f33356..6c9627d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,5 @@ use config::Config; +use poise::serenity_prelude::UserId; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use structopt::StructOpt; @@ -12,6 +13,7 @@ pub struct Args { pub struct DaemonConfig { pub discord_api_token: String, pub db_path: PathBuf, + pub admins: Vec, } impl DaemonConfig { diff --git a/src/discord/mod.rs b/src/discord/mod.rs index 087309c..19158a4 100644 --- a/src/discord/mod.rs +++ b/src/discord/mod.rs @@ -1,10 +1,12 @@ 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::MessageBuilder; -use poise::{serenity_prelude as serenity, ChoiceParameter}; +use poise::serenity_prelude::{CreateAttachment, Mentionable, MessageBuilder, UserId}; +use poise::{serenity_prelude as serenity, ChoiceParameter, CreateReply}; #[allow(dead_code)] struct Data { @@ -14,7 +16,11 @@ struct Data { type Error = Box; type Context<'a> = poise::Context<'a, Data, Error>; -#[poise::command(slash_command, ephemeral)] +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, @@ -41,7 +47,7 @@ async fn add_service( Ok(()) } -#[poise::command(slash_command, ephemeral)] +#[poise::command(slash_command, ephemeral, check = "check_owner")] async fn add_service_group( ctx: Context<'_>, #[description = "New group name"] group_name: String, @@ -83,6 +89,96 @@ async fn add_service_to_group( 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<'_>, @@ -93,11 +189,21 @@ async fn list_services( } 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| true)? + .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()) @@ -108,7 +214,7 @@ async fn list_services( .map(|s: Service| s.name) .collect(); - let mut message_builder = serenity::MessageBuilder::new(); + let mut message_builder = MessageBuilder::new(); if let Some(group) = group { message_builder.push_line(format!("## `{}` Services", group.name)); @@ -139,6 +245,17 @@ pub enum ServiceAction { 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<'_>, @@ -146,8 +263,21 @@ async fn service( #[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)?; @@ -197,6 +327,11 @@ pub async fn run_bot(db: Database, config: DaemonConfig) { 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() }) diff --git a/src/error.rs b/src/error.rs index dd9dd46..97a4f42 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,7 +1,14 @@ use thiserror::Error; +#[allow(dead_code)] #[derive(Error, Debug)] pub enum DaemonError { #[error("DB error")] DBError(#[from] j_db::error::JDbError), + #[error("User group not found")] + UserGroupNotFound, + #[error("Service group not found")] + ServiceGroupNotFound, + #[error("User not authorized")] + UserNotAuthorized, } diff --git a/src/models/mod.rs b/src/models/mod.rs index 1f278a4..b7f85a8 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1 +1,2 @@ pub mod service; +pub mod user; diff --git a/src/models/service.rs b/src/models/service.rs index 12a17ba..63f0fc0 100644 --- a/src/models/service.rs +++ b/src/models/service.rs @@ -89,6 +89,15 @@ impl ServiceGroup { .next() .clone()) } + + pub fn get_service_groups( + db: &Database, + service: u64, + ) -> Result, j_db::error::JDbError> { + Ok(db + .filter(|_, group: &ServiceGroup| group.services.contains(&service))? + .collect()) + } } impl j_db::model::JdbModel for ServiceGroup { diff --git a/src/models/user.rs b/src/models/user.rs new file mode 100644 index 0000000..76ccde5 --- /dev/null +++ b/src/models/user.rs @@ -0,0 +1,195 @@ +use crate::models::service::ServiceGroup; +use bitflags::bitflags; +use j_db::database::Database; +use j_db::model::JdbModel; +use poise::serenity_prelude::UserId; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct User { + pub discord_uuid: UserId, + id: Option, +} + +impl User { + pub fn new(discord_uuid: UserId) -> Self { + Self { + discord_uuid, + id: None, + } + } + + pub fn find_user_by_id( + db: &Database, + uuid: UserId, + ) -> Result, j_db::error::JDbError> { + let user = db + .filter(|_, user: &User| user.discord_uuid == uuid)? + .next(); + + Ok(user) + } +} + +impl j_db::model::JdbModel for User { + fn id(&self) -> Option { + self.id + } + + fn set_id(&mut self, id: u64) { + self.id = Some(id); + } + + fn tree() -> String { + "User".to_string() + } +} + +bitflags! { + #[derive(Serialize, Deserialize, Copy, Debug, Clone, PartialEq)] + pub struct Permissions: u32 { + const Status = 0b00001; + const Start = 0b00010; + const Stop = 0b00100; + const Restart = 0b01000; + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UserGroup { + pub name: String, + pub description: String, + pub members: HashSet, + pub permissions: HashMap, + + id: Option, +} + +impl UserGroup { + pub fn new(name: &str, description: &str) -> Self { + Self { + name: name.to_string(), + description: description.to_string(), + members: HashSet::new(), + permissions: Default::default(), + id: None, + } + } + + pub fn find_group_by_name( + db: &Database, + name: &str, + ) -> Result, j_db::error::JDbError> { + let group = db + .filter(|_, group: &UserGroup| group.name.eq_ignore_ascii_case(name))? + .next(); + + Ok(group) + } + + pub fn add_user_to_group( + db: &Database, + group_name: &str, + user_id: u64, + ) -> Result { + let mut group = + Self::find_group_by_name(db, group_name)?.ok_or(j_db::error::JDbError::NotFound)?; + + group.members.insert(user_id); + + let group = db.insert(group)?; + + Ok(group) + } + + pub fn set_group_permission( + db: &Database, + group_name: &str, + service_group: u64, + permissions: Permissions, + ) -> Result { + let mut group = + Self::find_group_by_name(db, group_name)?.ok_or(j_db::error::JDbError::NotFound)?; + + group.permissions.insert(service_group, permissions); + + let group = db.insert(group)?; + + Ok(group) + } + + pub fn get_user_groups( + db: &Database, + user: u64, + ) -> Result, j_db::error::JDbError> { + Ok(db + .filter(|_, group: &UserGroup| group.members.contains(&user))? + .collect()) + } + + pub fn check_if_user_has_permission( + db: &Database, + user: u64, + service: u64, + permission: Permissions, + ) -> Result { + let user_groups = Self::get_user_groups(db, user)?; + let service_group = ServiceGroup::get_service_groups(db, service)?; + + Ok(Self::check_perm(permission, user_groups, service_group)) + } + + fn check_perm( + permission: Permissions, + user_groups: Vec, + service_group: Vec, + ) -> bool { + let mut unified_permission_map: HashMap = HashMap::new(); + + for user_group in user_groups { + for (group, permission) in user_group.permissions { + #[allow(clippy::map_entry)] + if unified_permission_map.contains_key(&group) { + let perm1 = unified_permission_map.get_mut(&group).unwrap(); + *perm1 |= permission; + } else { + unified_permission_map.insert(group, permission); + } + } + } + + let service_group_ids: HashSet = service_group + .iter() + .map(|service_group| service_group.id().unwrap()) + .collect(); + + for service_group in service_group_ids { + let perm = unified_permission_map.get(&service_group); + + if let Some(perm) = perm { + let check_perm = (permission & *perm) == permission; + + if check_perm { + return true; + } + } + } + + false + } +} + +impl JdbModel for UserGroup { + fn id(&self) -> Option { + self.id + } + + fn set_id(&mut self, id: u64) { + self.id = Some(id); + } + + fn tree() -> String { + "UserGroup".to_string() + } +}