diff --git a/.gitignore b/.gitignore index 0b745e2..2483f0b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ /target -.env \ No newline at end of file +.env +*.db +*.wal \ No newline at end of file diff --git a/src/commands.rs b/src/commands.rs index f7e924b..42732cc 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -1,9 +1,6 @@ -use std::sync::Mutex; - -use once_cell::sync::OnceCell; use poise::command; -use crate::database::Database; +use crate::database::*; const HELP_MESSAGE: &str = " Hello fellow human! I am a bot that can help you adding new technologies to a git repository. @@ -34,15 +31,12 @@ To get help, just type: /help ``` "; -const AUTHORIZED_USERS: [u64; 1] = [252497456447750144]; type Error = Box; type Context<'a> = poise::Context<'a, MsgData, Error>; pub struct MsgData {} -pub static DB: OnceCell> = OnceCell::new(); - /// Show help for all commands #[command(slash_command, prefix_command)] pub async fn help(ctx: Context<'_>) -> Result<(), Error> { @@ -58,7 +52,9 @@ pub async fn add( #[description = "Technology name"] technology: String, #[description = "Git repository link"] link: String, ) -> Result<(), Error> { - ctx.say(format!("Added {technology} with link {link}")) + add_tech(link.clone(), technology.clone())?; + + ctx.say(format!("Added {technology} with link {link}",)) .await?; Ok(()) @@ -67,7 +63,21 @@ pub async fn add( /// List all available technologies #[command(slash_command, prefix_command)] pub async fn list(ctx: Context<'_>) -> Result<(), Error> { - ctx.say("Listed all technologies").await?; + let techs = list_tech()?; + if techs.len() == 0 { + ctx.say("No technologies saved").await?; + return Ok(()); + } + + ctx.say(format!( + "Saved technologies: {}", + techs + .iter() + .map(|tech| format!("[{}]({})", tech.name, tech.link)) + .collect::>() + .join(", ") + )) + .await?; Ok(()) } @@ -78,7 +88,12 @@ pub async fn search( ctx: Context<'_>, #[description = "Technology name"] technology: String, ) -> Result<(), Error> { - ctx.say(format!("Found {technology}")).await?; + if let Some(tech) = search_tech(technology.clone())? { + ctx.say(format!("Name: {}\nLink: {}", tech.name, tech.link)) + .await?; + } else { + ctx.say("No technology found").await?; + } Ok(()) } @@ -89,13 +104,37 @@ pub async fn remove( ctx: Context<'_>, #[description = "Technology name"] technology: String, ) -> Result<(), Error> { - if !AUTHORIZED_USERS.contains(&ctx.author().id.0) { + if is_auth_user(ctx.author().id.to_string())? { ctx.say("You don't have permission to remove a technology") .await?; return Ok(()); } - ctx.say(format!("Removed {technology}")).await?; + remove_tech(technology)?; + + ctx.say("Technology removed").await?; + + Ok(()) +} + +/// Remove a technology from the technologies list +#[command(slash_command, prefix_command)] +pub async fn add_auth_user( + ctx: Context<'_>, + #[description = "Discord user ID"] id: String, +) -> Result<(), Error> { + if !std::env::var("ADMIN_USERS") + .expect("missing ADMIN_USERS") + .contains(&ctx.author().id.to_string()) + { + ctx.say("You don't have permission to add a new user") + .await?; + return Ok(()); + } + + set_auth_user(id)?; + + ctx.say("User added!").await?; Ok(()) } diff --git a/src/database.rs b/src/database.rs index 736deeb..516b65a 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,78 +1,88 @@ -use polodb_core::{ - bson::{doc, oid::ObjectId}, - Database as DB, -}; +use once_cell::sync::OnceCell; +use polodb_core::{bson::doc, Database}; use serde::{Deserialize, Serialize}; -use std::{ - io::{Error, ErrorKind}, - str::FromStr, -}; +use std::io::{Error, ErrorKind}; + +pub static DB: OnceCell = OnceCell::new(); #[derive(Default, Debug, Serialize, Deserialize)] -struct Technology { - link: String, - name: String, +pub struct Technology { + pub link: String, + pub name: String, } #[derive(Debug, Serialize, Deserialize)] struct AuthorizedUser { - name: String, + discord_id: String, } -pub struct Database { - pub db: DB, +/// Add a new technology to the database. +pub fn add_tech(link: String, name: String) -> Result<(), Error> { + DB.get() + .unwrap() + .collection::("technologies") + .insert_one(Technology { link, name }) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))?; + + Ok(()) } -impl Database { - /// Create a new database instance. - pub fn new() -> Self { - Self { - db: DB::open_file("test-polo.db").expect("failed to initialize database"), - } - } +pub fn remove_tech(name: String) -> Result<(), Error> { + DB.get() + .unwrap() + .collection::("technologies") + .delete_one(doc! { "name": name }) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))?; - /// Add a new technology to the database. - pub fn add_tech(&self, link: String, name: String) -> Result { - Ok(self - .db - .collection::("technologies") - .insert_one(Technology { link, name }) - .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? - .inserted_id - .to_string()) - } + Ok(()) +} - pub fn remove_tech(&self, id: &str) -> Result<(), Error> { - self.db - .collection::("technologies") - .delete_one(doc! { "_id": Some(ObjectId::from_str(id).map_err(|err| Error::new(ErrorKind::InvalidInput, err))?) }) - .map_err(|err| Error::new(ErrorKind::InvalidInput, err))?; +pub fn list_tech() -> Result, Error> { + Ok(DB + .get() + .unwrap() + .collection("technologies") + .find(None) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? + .filter(|doc| doc.is_ok()) + .map(|doc| doc.unwrap()) + .collect()) +} - Ok(()) - } - - pub fn list_tech(&self) -> Result, Error> { - Ok(self - .db - .collection("technologies") - .find(None) - .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? - .filter(|doc| doc.is_ok()) - .map(|doc| doc.unwrap()) - .collect()) - } - - pub fn search_tech(&self, name: String) -> Result { - if let Some(tech) = self - .db - .collection::("technologies") - .find(doc! { "$eq": [{"name": name}] }) - .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? - .next() - { - Ok(tech.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?) - } else { - Err(Error::new(ErrorKind::NotFound, "Technology not found")) - } +pub fn search_tech(name: String) -> Result, Error> { + if let Some(tech) = DB + .get() + .unwrap() + .collection::("technologies") + .find(doc! { "name": {"$eq": name} }) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? + .next() + { + Ok(Some( + tech.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?, + )) + } else { + Ok(None) } } + +pub fn set_auth_user(discord_id: String) -> Result<(), Error> { + DB.get() + .unwrap() + .collection("authorized_users") + .insert_one(AuthorizedUser { discord_id }) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))?; + + Ok(()) +} + +pub fn is_auth_user(discord_id: String) -> Result { + Ok(DB + .get() + .unwrap() + .collection::("authorized_users") + .find(doc! { "$eq": [{"discord_id": discord_id}] }) + .map_err(|err| Error::new(ErrorKind::InvalidInput, err))? + .next() + .is_some()) +} diff --git a/src/main.rs b/src/main.rs index a667bff..e72e4f6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,15 +3,15 @@ mod database; use poise::serenity_prelude as serenity; -use commands::{add, help, list, remove, search, MsgData, DB}; +use commands::{add, add_auth_user, help, list, remove, search, MsgData}; +use database::DB; +use polodb_core::Database; #[tokio::main] async fn main() { - DB.set(Mutex) - let framework = poise::Framework::builder() .options(poise::FrameworkOptions { - commands: vec![help(), add(), list(), search(), remove()], + commands: vec![help(), add(), list(), search(), remove(), add_auth_user()], ..Default::default() }) .token(std::env::var("DISCORD_TOKEN").expect("missing DISCORD_TOKEN")) @@ -25,5 +25,9 @@ async fn main() { }) }); + DB.get_or_init(|| { + Database::open_file(std::env::var("DB_PATH").expect("missing DB_PATH")) + .expect("failed to initialize database") + }); framework.run().await.unwrap(); }