Merge pull request #59 from tinaun/command-editing

Track edits and deletions of commands

Author: technetos

Cargo.lock

@@ -18,4 +18,4 @@ lazy_static = "1.4.0"
 log = "0.4.0"
 env_logger = "0.7.1"
 envy = "0.4"
-indexmap = "1.3"
+indexmap = "1.6"

Cargo.toml

@@ -2,16 +2,36 @@ use crate::{
     commands::{Args, Result},
     db::DB,
     schema::roles,
+    CommandHistory,
 };
 use diesel::prelude::*;
 use serenity::{model::prelude::*, utils::parse_username};
 
 /// Send a reply to the channel the message was received on.  
 pub(crate) fn send_reply(args: &Args, message: &str) -> Result<()> {
-    args.msg.channel_id.say(&args.cx, message)?;
+    if let Some(response_id) = response_exists(args) {
+        info!("editing message: {:?}", response_id);
+        args.msg
+            .channel_id
+            .edit_message(&args.cx, response_id, |msg| msg.content(message))?;
+    } else {
+        let command_id = args.msg.id;
+        let response = args.msg.channel_id.say(&args.cx, message)?;
+
+        let mut data = args.cx.data.write();
+        let history = data.get_mut::<CommandHistory>().unwrap();
+        history.insert(command_id, response.id);
+    }
+
     Ok(())
 }
 
+fn response_exists(args: &Args) -> Option<MessageId> {
+    let data = args.cx.data.read();
+    let history = data.get::<CommandHistory>().unwrap();
+    history.get(&args.msg.id).cloned()
+}
+
 /// Determine if a member sending a message has the `Role`.  
 pub(crate) fn has_role(args: &Args, role: &RoleId) -> Result<bool> {
     Ok(args

src/api.rs

@@ -1,3 +1,4 @@
+use super::CommandHistory;
 use crate::{
     api,
     commands::{Args, Result},
@@ -48,7 +49,7 @@ pub(crate) fn save_unban(user_id: String, guild_id: String) -> Result<()> {
     Ok(())
 }
 
-pub(crate) fn start_unban_thread(cx: Context) {
+pub(crate) fn start_cleanup_thread(cx: Context) {
     use std::str::FromStr;
     if !UNBAN_THREAD_INITIALIZED.load(Ordering::SeqCst) {
         UNBAN_THREAD_INITIALIZED.store(true, Ordering::SeqCst);
@@ -69,6 +70,18 @@ pub(crate) fn start_unban_thread(cx: Context) {
                     info!("Unbanning user {}", &row.1);
                     guild_id.unban(&cx, u64::from_str(&row.1)?)?;
                 }
+
+                let mut data = cx.data.write();
+                let history = data.get_mut::<CommandHistory>().unwrap();
+
+                // always keep the last command in history
+                if history.len() > 0 {
+                    info!("Clearing command history");
+                    history.drain(..history.len() - 1);
+                }
+
+                drop(data);
+
                 sleep(Duration::new(HOUR, 0));
             }
         });

src/ban.rs

@@ -25,8 +25,10 @@ use diesel::prelude::*;
 use envy;
 use indexmap::IndexMap;
 use serde::Deserialize;
-use serenity::{model::prelude::*, prelude::*};
+use serenity::{model::prelude::*, prelude::*, utils::CustomMessage};
+use std::time::Duration;
 
+const MESSAGE_AGE_MAX: Duration = Duration::from_secs(3600);
 #[derive(Deserialize)]
 struct Config {
     tags: bool,
@@ -186,7 +188,7 @@ fn app() -> Result<()> {
     });
 
     let mut client = Client::new_with_extras(&config.discord_token, |e| {
-        e.raw_event_handler(Events { cmds });
+        e.event_handler(Events { cmds });
         e
     })?;
 
@@ -221,33 +223,75 @@ fn main() {
     }
 }
 
+struct CommandHistory {}
+impl TypeMapKey for CommandHistory {
+    type Value = IndexMap<MessageId, MessageId>;
+}
+
 struct Events {
     cmds: Commands,
 }
 
-impl RawEventHandler for Events {
-    fn raw_event(&self, cx: Context, event: Event) {
-        match event {
-            Event::Ready(ev) => {
-                info!("{} connected to discord", ev.ready.user.name);
-                ban::start_unban_thread(cx);
-            }
-            Event::MessageCreate(ev) => {
-                self.cmds.execute(cx, &ev.message);
-            }
-            Event::ReactionAdd(ev) => {
-                if let Err(e) = welcome::assign_talk_role(&cx, &ev) {
-                    error!("{}", e);
-                }
-            }
-            Event::GuildBanRemove(ev) => {
-                if let Err(e) =
-                    ban::save_unban(format!("{}", ev.user.id), format!("{}", ev.guild_id))
-                {
-                    error!("{}", e);
-                }
-            }
-            _ => (),
+impl EventHandler for Events {
+    fn ready(&self, cx: Context, ready: Ready) {
+        info!("{} connected to discord", ready.user.name);
+
+        let mut data = cx.data.write();
+        data.insert::<CommandHistory>(IndexMap::new());
+        drop(data);
+
+        ban::start_cleanup_thread(cx);
+    }
+
+    fn message(&self, cx: Context, message: Message) {
+        self.cmds.execute(cx, &message);
+    }
+
+    fn message_update(
+        &self,
+        cx: Context,
+        _: Option<Message>,
+        _: Option<Message>,
+        ev: MessageUpdateEvent,
+    ) {
+        let age = ev.timestamp.and_then(|create| {
+            ev.edited_timestamp
+                .and_then(|edit| edit.signed_duration_since(create).to_std().ok())
+        });
+
+        if age.is_some() && age.unwrap() < MESSAGE_AGE_MAX {
+            let mut msg = CustomMessage::new();
+            msg.id(ev.id)
+                .channel_id(ev.channel_id)
+                .content(ev.content.unwrap_or_else(|| String::new()));
+
+            let msg = msg.build();
+            info!(
+                "sending edited message - {:?} {:?}",
+                msg.content, msg.author
+            );
+            self.cmds.execute(cx, &msg);
+        }
+    }
+
+    fn message_delete(&self, cx: Context, channel_id: ChannelId, message_id: MessageId) {
+        let mut data = cx.data.write();
+        let history = data.get_mut::<CommandHistory>().unwrap();
+        if let Some(response_id) = history.remove(&message_id) {
+            info!("deleting message: {:?}", response_id);
+            let _ = channel_id.delete_message(&cx, response_id);
+        }
+    }
+
+    fn reaction_add(&self, cx: Context, reaction: Reaction) {
+        if let Err(e) = welcome::assign_talk_role(&cx, &reaction) {
+            error!("{}", e);
+        }
+    }
+
+    fn guild_ban_removal(&self, _cx: Context, guild_id: GuildId, user: User) {
+        if let Err(e) = ban::save_unban(format!("{}", user.id), format!("{}", guild_id)) {
+            error!("{}", e);
         }
     }
 }

src/main.rs

@@ -63,9 +63,7 @@ pub(crate) fn post_message(args: Args) -> Result<()> {
     Ok(())
 }
 
-pub(crate) fn assign_talk_role(cx: &Context, ev: &ReactionAddEvent) -> Result<()> {
-    let reaction = &ev.reaction;
-
+pub(crate) fn assign_talk_role(cx: &Context, reaction: &Reaction) -> Result<()> {
     let channel = reaction.channel(cx)?;
     let channel_id = ChannelId::from(&channel);
     let message = reaction.message(cx)?;
@@ -114,13 +112,13 @@ pub(crate) fn assign_talk_role(cx: &Context, ev: &ReactionAddEvent) -> Result<()
 
                     // Requires ManageMessage permission
                     if let Some((_, _, user_id)) = me {
-                        if ev.reaction.user_id.0.to_string() != user_id {
-                            ev.reaction.delete(cx)?;
+                        if reaction.user_id.0.to_string() != user_id {
+                            reaction.delete(cx)?;
                         }
                     }
                 }
             } else {
-                ev.reaction.delete(cx)?;
+                reaction.delete(cx)?;
             }
         }
     }