From 9507d80303d7088def37488e393528715d198ea6 Mon Sep 17 00:00:00 2001
From: jamesbt365 <jamesbt365@gmail.com>
Date: Sun, 26 May 2024 20:14:16 +0100
Subject: [PATCH] Allow only specific team roles to be initailized to owners

---
 src/framework/mod.rs             | 18 ++++++++++++++----
 src/structs/framework_options.rs |  6 ++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/framework/mod.rs b/src/framework/mod.rs
index 2fbb4ee87d23..c391277be677 100644
--- a/src/framework/mod.rs
+++ b/src/framework/mod.rs
@@ -138,7 +138,13 @@ impl<U: Send + Sync, E: Send + Sync> serenity::Framework for Framework<U, E> {
         self.shard_manager = Some(client.shard_manager.clone());
 
         if self.options.initialize_owners {
-            if let Err(e) = insert_owners_from_http(&client.http, &mut self.options.owners).await {
+            if let Err(e) = insert_owners_from_http(
+                &client.http,
+                &mut self.options.owners,
+                &self.options.initialized_team_roles,
+            )
+            .await
+            {
                 tracing::warn!("Failed to insert owners from HTTP: {e}");
             }
         }
@@ -236,6 +242,7 @@ fn message_content_intent_sanity_check<U, E>(
 pub async fn insert_owners_from_http(
     http: &serenity::Http,
     owners: &mut std::collections::HashSet<serenity::UserId>,
+    initialized_teams: &Option<Vec<serenity::TeamMemberRole>>,
 ) -> Result<(), serenity::Error> {
     let application_info = http.get_current_application_info().await?;
 
@@ -245,9 +252,12 @@ pub async fn insert_owners_from_http(
 
     if let Some(team) = application_info.team {
         for member in team.members {
-            // This `if` currently always evaluates to true but it becomes important once
-            // Discord implements more team roles than Admin
-            if member.permissions.iter().any(|p| p == "*") {
+            let Some(initialized_teams) = initialized_teams else {
+                owners.insert(member.user.id);
+                continue;
+            };
+
+            if initialized_teams.iter().any(|r| *r == member.role) {
                 owners.insert(member.user.id);
             }
         }
diff --git a/src/structs/framework_options.rs b/src/structs/framework_options.rs
index 3ec310e9c3c5..0e859067a745 100644
--- a/src/structs/framework_options.rs
+++ b/src/structs/framework_options.rs
@@ -67,6 +67,11 @@ pub struct FrameworkOptions<U, E> {
     ///
     /// True by default.
     pub initialize_owners: bool,
+    /// If set and [`Self::initialize_owners`] is true, only select teams will be initialized by
+    /// the results of [`serenity::Http::get_current_application_info()`].
+    ///
+    /// None by default.
+    pub initialized_team_roles: Option<Vec<serenity::TeamMemberRole>>,
     // #[non_exhaustive] forbids struct update syntax for ?? reason
     #[doc(hidden)]
     pub __non_exhaustive: (),
@@ -120,6 +125,7 @@ where
             prefix_options: Default::default(),
             owners: Default::default(),
             initialize_owners: true,
+            initialized_team_roles: None,
             __non_exhaustive: (),
         }
     }