From 9fd85e38fb915a58d3f3b00f0d26a3ce847da376 Mon Sep 17 00:00:00 2001
From: minneelyyyy <minneelyyyy@proton.me>
Date: Wed, 20 Nov 2024 19:32:33 -0500
Subject: [PATCH] new anti-ping-spam functionality

---
 Cargo.toml        |  1 +
 src/common.rs     |  3 ++-
 src/main.rs       | 49 +++++++++++++++++++++++++++++++++++++++++++++--
 src/ping_limit.rs | 13 +++++++++++++
 4 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 src/ping_limit.rs

diff --git a/Cargo.toml b/Cargo.toml
index 92f48f9..4896e94 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,3 +10,4 @@ dotenv = "0.15.0"
 rand = "0.8.5"
 clap = { version = "4.5.20", features = ["derive"] }
 lamm = { git = "https://github.com/minneelyyyy/lamm", branch = "dev" }
+regex = "1.11.1"
\ No newline at end of file
diff --git a/src/common.rs b/src/common.rs
index 70bf39c..1950350 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -4,7 +4,8 @@ use std::collections::HashMap;
 use poise::serenity_prelude::UserId;
 
 pub struct Data {
-    pub users: Arc<Mutex<HashMap<UserId, usize>>>
+    pub users: Arc<Mutex<HashMap<UserId, usize>>>,
+    pub mentions: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
 }
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;
diff --git a/src/main.rs b/src/main.rs
index dfc2d6b..2db2cfa 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 mod commands;
+mod ping_limit;
 
 pub mod common;
 use crate::common::{Context, Error, Data};
@@ -6,6 +7,7 @@ use crate::common::{Context, Error, Data};
 use std::collections::HashMap;
 use std::env;
 use std::sync::Arc;
+use std::time::{Instant, Duration};
 
 use poise::serenity_prelude as serenity;
 use tokio::sync::Mutex;
@@ -19,6 +21,43 @@ struct BotArgs {
     prefix: Option<String>,
 }
 
+async fn event_handler(
+    ctx: &serenity::Context,
+    event: &serenity::FullEvent,
+    _framework: poise::FrameworkContext<'_, Data, Error>,
+    data: &Data,
+) -> Result<(), Error> {
+    match event {
+        serenity::FullEvent::Message { new_message: message } => {
+            let mentions = ping_limit::extract_mentions(&message.content);
+            let mut cooldowns = data.mentions.lock().await;
+
+            if mentions.iter()
+                .filter(|&&id| id != message.author.id)
+                .any(|mention| cooldowns.get(mention).map(|t| Instant::now().duration_since(*t) < Duration::from_secs(20)).unwrap_or(false))
+            {
+                message.reply(ctx, "stop spamming!").await?;
+
+                let guild = match message.guild_id {
+                    Some(g) => g,
+                    None => return Ok(()),
+                };
+
+                let mut member = guild.member(ctx, message.author.id).await.unwrap();
+                member.disable_communication_until_datetime(ctx,
+                    serenity::Timestamp::from_unix_timestamp(serenity::Timestamp::now().unix_timestamp() + 60 as i64).unwrap()).await?;
+            }
+
+            for mention in mentions {
+                cooldowns.insert(mention, Instant::now());
+            }
+        }
+        _ => (),
+    }
+
+    Ok(())
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Error> {
     dotenv::dotenv().ok();
@@ -33,16 +72,22 @@ async fn main() -> Result<(), Error> {
             prefix_options: poise::PrefixFrameworkOptions {
                 prefix: args.prefix,
                 edit_tracker: Some(Arc::new(
-                    poise::EditTracker::for_timespan(std::time::Duration::from_secs(3600)))),
+                    poise::EditTracker::for_timespan(std::time::Duration::from_secs(10)))),
                 case_insensitive_commands: true,
                 ..Default::default()
             },
+            event_handler: |ctx, event, framework, data| {
+                Box::pin(event_handler(ctx, event, framework, data))
+            },
             ..Default::default()
         })
         .setup(|ctx, _ready, framework| {
             Box::pin(async move {
                 poise::builtins::register_globally(ctx, &framework.options().commands).await?;
-                Ok(Data { users: Arc::new(Mutex::new(HashMap::new())) })
+                Ok(Data {
+                    users: Arc::new(Mutex::new(HashMap::new())),
+                    mentions: Arc::new(Mutex::new(HashMap::new())),
+                })
             })
         })
         .build();
diff --git a/src/ping_limit.rs b/src/ping_limit.rs
new file mode 100644
index 0000000..07c33b6
--- /dev/null
+++ b/src/ping_limit.rs
@@ -0,0 +1,13 @@
+
+use poise::serenity_prelude::*;
+use regex::Regex;
+
+pub fn extract_mentions(content: &str) -> Vec<UserId> {
+    // Define the regex pattern for user mentions
+    let re = Regex::new(r"<@(\d+)>").unwrap();
+
+    // Find all matches and capture the IDs
+    re.captures_iter(content)
+        .filter_map(|cap| cap.get(1).map(|id| id.as_str().parse().unwrap()))
+        .collect()
+}