From 8dbc7ee86352ba0259de584ec49b46415764d06d Mon Sep 17 00:00:00 2001
From: Oliver Atkinson <oatkinson@ehps.k12.mt.us>
Date: Thu, 25 Jul 2024 14:42:15 -0600
Subject: [PATCH] Put scrape behind owner privliages I don't want normies
 tyring to scrape the whole server

---
 src/command.rs | 12 +++++++++---
 src/main.rs    | 14 ++++++++------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/command.rs b/src/command.rs
index 43414e2..ce2dd7c 100644
--- a/src/command.rs
+++ b/src/command.rs
@@ -205,8 +205,8 @@ impl Server {
 
 }
 
-#[poise::command(slash_command, rename = "scrape_all", guild_only)]
-pub async fn scrape_all(ctx: Context<'_>) -> Result<(), Error> {
+#[poise::command(slash_command, rename = "scrape_all", guild_only, owners_only)]
+pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Error> {
     let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap();
 
     let invoker = ctx.author().name.clone();
@@ -225,7 +225,13 @@ pub async fn scrape_all(ctx: Context<'_>) -> Result<(), Error> {
                 let end = start.elapsed().as_millis();
                 let msg_count = server.message_count();
 
-                match serde_json::to_string(&server) {
+                let print = if pretty_print {
+                    serde_json::to_string_pretty(&server)
+                } else {
+                    serde_json::to_string(&server)
+                };
+
+                match print {
                     Ok(ok) => {
                         if let Err(e) = fs::write("server.json", ok) {
                             error!("Problem writing server to disk: {e}");
diff --git a/src/main.rs b/src/main.rs
index df6d93b..5810e3f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,7 @@
+use std::collections::HashSet;
+
 use once_cell::sync::Lazy;
-use poise::serenity_prelude::{self as serenity, GatewayIntents};
+use poise::serenity_prelude::{self as serenity, GatewayIntents, UserId};
 use tracing::{debug, error, info, warn, Level};
 use tracing_subscriber::EnvFilter;
 mod command;
@@ -30,14 +32,13 @@ async fn main() {
 
     // Generate sick text like this: 
     // http://www.patorjk.com/software/taag/#p=testall&f=Graffiti&t=hello%20world
-    info!(r#"
-    Invite this bot with:
-    "#);
-    info!("https://discord.com/api/oauth2/authorize?client_id={}&permissions={}&scope=bot",
+    info!("Invite with: https://discord.com/api/oauth2/authorize?client_id={}&permissions={}&scope=bot",
         ENV.id,
         ENV.intents.bits(),
     );
-    info!("\n");
+
+    let mut owners = HashSet::new();
+    owners.insert(UserId::new(423970006334832650));
 
     // Setup framework
     let framework = poise::Framework::builder()
@@ -46,6 +47,7 @@ async fn main() {
                 command::index_cmd(),
                 command::scrape_all(),
             ],
+            owners,
             ..Default::default()
         })
         .setup(|ctx, _ready, framework| {