From 8a5cb04d306e25d4f2454656b0b80742fd163285 Mon Sep 17 00:00:00 2001
From: Ryan Cao <70191398+ryanccn@users.noreply.github.com>
Date: Sun, 29 Sep 2024 20:23:40 +0800
Subject: [PATCH] fix(intelligence): add configurable allowlist

---
 src/handlers/intelligence.rs | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/handlers/intelligence.rs b/src/handlers/intelligence.rs
index 7a33d2f..2d712fe 100644
--- a/src/handlers/intelligence.rs
+++ b/src/handlers/intelligence.rs
@@ -1,11 +1,23 @@
-use std::time::Duration;
+use poise::serenity_prelude as serenity;
 
 use color_eyre::eyre::Result;
-use poise::serenity_prelude as serenity;
+use once_cell::sync::Lazy;
+use std::time::Duration;
 use tokio::{task, time};
 
 use crate::intelligence;
 
+static ALLOWED_ROLES: Lazy<Vec<serenity::RoleId>> = Lazy::new(|| {
+    std::env::var("INTELLIGENCE_ALLOWED_ROLES")
+        .ok()
+        .map(|s| {
+            s.split(',')
+                .filter_map(|f| f.trim().parse::<serenity::RoleId>().ok())
+                .collect::<Vec<_>>()
+        })
+        .unwrap_or_default()
+});
+
 #[tracing::instrument(skip_all, fields(message_id = message.id.get()))]
 pub async fn handle(message: &serenity::Message, ctx: &serenity::Context) -> Result<()> {
     if message
@@ -15,6 +27,14 @@ pub async fn handle(message: &serenity::Message, ctx: &serenity::Context) -> Res
         return Ok(());
     }
 
+    let member = message.member(&ctx).await?;
+
+    if !member.permissions(&ctx.cache)?.administrator()
+        && !member.roles.iter().any(|r| ALLOWED_ROLES.contains(r))
+    {
+        return Ok(());
+    }
+
     if let Some(query) = message
         .content
         .strip_prefix(&format!("<@{}>", ctx.cache.current_user().id))
@@ -41,10 +61,7 @@ pub async fn handle(message: &serenity::Message, ctx: &serenity::Context) -> Res
         let username = message.author.tag();
         let display_name = message.author.global_name.clone().map(|s| s.into_string());
 
-        let nick = message
-            .member
-            .as_ref()
-            .and_then(|m| m.nick.as_ref().map(|s| s.to_owned().into_string()));
+        let nick = member.nick.as_ref().map(|s| s.to_owned().into_string());
 
         let resp = intelligence::query(intelligence::Request {
             query: query.to_owned(),