diff --git a/Cargo.lock b/Cargo.lock index 7bf6731..a591e0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -352,6 +352,7 @@ version = "0.1.0" dependencies = [ "anyhow", "dotenv", + "once_cell", "poise", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index a514037..06f8c10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,3 +14,4 @@ tokio = { version = "1.21.2", features = ["macros", "rt-multi-thread"] } poise = { version = "0.6", features = ["cache"] } dotenv = "0.15.0" anyhow = "1.0.75" +once_cell = "1.19.0" diff --git a/src/command.rs b/src/command.rs index 1d5d701..c3753b4 100644 --- a/src/command.rs +++ b/src/command.rs @@ -1,8 +1,8 @@ -use std::fmt::Display; +use std::{collections::HashMap, fmt::Display, sync::Arc}; use crate::Context; use anyhow::Error; -use poise::serenity_prelude::{ChannelId, ChannelType, GuildChannel}; +use poise::serenity_prelude::{Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message}; struct Server { channels: Vec, @@ -12,7 +12,19 @@ struct Server { struct Channel { this: GuildChannel, - children: Vec + children: Vec, + messages: Vec, +} + +impl Channel { + fn new(this: GuildChannel) -> Self { + Self { + this, + // Empty vecs don't allocate until a push + children: Vec::new(), + messages: Vec::new(), + } + } } impl Display for Server { @@ -69,12 +81,12 @@ impl Server { } // TODO this might be broken - fn search<'a>(target: &'a mut Vec, find: &ChannelId) -> Option<&'a mut Channel> { + fn search_by_id<'a>(target: &'a mut Vec, find: &ChannelId) -> Option<&'a mut Channel> { for child in target { if child.this.id == *find { return Some(child); } - match Self::search(&mut child.children, find) { + match Self::search_by_id(&mut child.children, find) { Some(x) => return Some(x), None => {}, } @@ -87,9 +99,9 @@ impl Server { if let Some(parent_id) = &insert.parent_id { // find the parent (needs to go thru all nodes) - match Self::search(&mut self.channels, &parent_id) { + match Self::search_by_id(&mut self.channels, &parent_id) { Some(parent_node) => { - parent_node.children.push(Channel { this: insert, children: Vec::new() }); + parent_node.children.push(Channel::new(insert)); }, None => { // couldn't find parent, store somewhere else until it's parent is added... @@ -98,18 +110,19 @@ impl Server { }, } } else { - self.channels.push(Channel { this: insert, children: Vec::new() }) + self.channels.push(Channel::new(insert)); } } + /// Cleans out the orphan channels, finding them parents. You'll want to use this before displaying anything. fn clean(&mut self) { if !self.needs_clean {return;} // Look thru the orphanage and try to find parents for orphan in &self.orphanage { if let Some(parent_id) = orphan.parent_id { - if let Some(found) = Self::search(&mut self.channels, &parent_id) { - found.children.push(Channel { this: orphan.clone(), children: Vec::new() }); + if let Some(found) = Self::search_by_id(&mut self.channels, &parent_id) { + found.children.push(Channel::new(orphan.clone())); } else { panic!("⚠️ Couldn't find parent for orphan!"); } @@ -120,50 +133,83 @@ impl Server { self.orphanage.clear(); self.needs_clean = false; } + + /// Scrapes messages for all the channels in `self`. + async fn scrape_all(&mut self) { + let settings = GetMessages::default().limit(5); + let cache: (&Arc, &Http) = (&Arc::new(Cache::new()), &Http::new(&crate::ENV.token)); + walk_channels(&mut self.channels, cache, settings).await; + + // recursive walk thru the channels + async fn walk_channels(all: &mut Vec, cache: impl CacheHttp + Clone, settings: GetMessages) { + for channel in all { + // get the messages + match channel.this.messages(cache.clone(), settings).await { + Ok(mesgs) => { + // store messages in our server object + channel.messages = mesgs; + if channel.messages.is_empty() { + eprintln!("{} was empty - (Or incorrect permissions)", channel.this.name); + } + }, + Err(_) => todo!(), + } + // Clone *should* be cheap - it's Arc under the hood + walk_channels(&mut channel.children, cache.clone(), settings).await; + } + } + } + + /// Walk thru all the channels and count the saved messages. Will only give relevant data if + /// done after `scrape_all()`. + fn message_count(&self) -> usize { + fn walk(this: &Vec) -> usize { + let mut total = 0; + for channel in this { + total += walk(&channel.children); + }; + total + } + walk(&self.channels) + } } +#[poise::command(slash_command, rename = "scrape_all", guild_only)] +pub async fn scrape_all(ctx: Context<'_>) -> Result<(), Error> { + let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap(); + if let Ok(map) = guild.channels(ctx.http()).await { + let mut server = index(map).await; + server.scrape_all().await; + + let _ = ctx.reply(&format!("Scraped {} messages", server.message_count())).await; + } + Ok(()) +} + +/// Get server's topology (and runs clean) +async fn index(map: HashMap) -> Server { + let mut server = Server::new(); + // iterate thru all channels + map.into_iter().for_each(|(_id, current)| { + // println!("{} {} {:?}", current.name, current.id, current.parent_id); + server.add(current); + // TODO take note of position + // Take node of vc user limit + }); + server.clean(); + server +} + + // NOTE!!! Make sure these names in quotes are lowercase! #[poise::command(slash_command, rename = "index", guild_only)] -pub async fn index(ctx: Context<'_>) -> Result<(), Error> { +pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> { let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap(); match guild.channels(ctx.http()).await { Ok(ok) => { - - let mut server = Server::new(); - - // iterate thru all channels - ok.into_iter().for_each(|(_id, current)| { - match current.kind { - poise::serenity_prelude::ChannelType::Text => { - server.add(current); - // current.position, - }, - poise::serenity_prelude::ChannelType::Private => todo!(), - poise::serenity_prelude::ChannelType::Voice => { - server.add(current); - // current.user_limit, - // current.parent_id, - }, - poise::serenity_prelude::ChannelType::GroupDm => todo!(), - poise::serenity_prelude::ChannelType::Category => { - server.add(current); - }, - poise::serenity_prelude::ChannelType::News => todo!(), - poise::serenity_prelude::ChannelType::NewsThread => todo!(), - poise::serenity_prelude::ChannelType::PublicThread => todo!(), - poise::serenity_prelude::ChannelType::PrivateThread => todo!(), - poise::serenity_prelude::ChannelType::Stage => todo!(), - poise::serenity_prelude::ChannelType::Directory => todo!(), - poise::serenity_prelude::ChannelType::Forum => todo!(), - poise::serenity_prelude::ChannelType::Unknown(_) => todo!(), - _ => todo!(), - } - }); - - server.clean(); - println!("{}", server); - + let server = index(ok).await; + let _ = ctx.reply(server.to_string()).await; }, Err(_) => todo!(), } diff --git a/src/main.rs b/src/main.rs index ce0a620..d10c238 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,20 +1,25 @@ +use once_cell::sync::Lazy; use poise::serenity_prelude::{self as serenity, GatewayIntents}; mod command; pub struct Data {} // User data, which is stored and accessible in all command invocations type Context<'a> = poise::Context<'a, Data, anyhow::Error>; +static ENV: Lazy = Lazy::new(|| { + read_env() +}); + #[tokio::main] async fn main() { - let env = read_env(); + // Generate sick text like this: // http://www.patorjk.com/software/taag/#p=testall&f=Graffiti&t=hello%20world println!(r#" Invite this bot with: "#); println!("https://discord.com/api/oauth2/authorize?client_id={}&permissions={}&scope=bot", - env.id, - env.intents.bits(), + ENV.id, + ENV.intents.bits(), ); print!("\n"); @@ -22,7 +27,8 @@ async fn main() { let framework = poise::Framework::builder() .options(poise::FrameworkOptions { commands: vec![ - command::index(), + command::index_cmd(), + command::scrape_all(), ], ..Default::default() }) @@ -38,7 +44,7 @@ async fn main() { .build(); // Start the Bot. - let client = serenity::ClientBuilder::new(env.token, env.intents) + let client = serenity::ClientBuilder::new(&ENV.token, ENV.intents) .framework(framework) .await;