From 8903b0afca8a7e1b18b33c92f41051bd5ceb9b3d Mon Sep 17 00:00:00 2001 From: Ashelyn Rose Date: Tue, 9 Jul 2024 12:35:58 -0600 Subject: handle attachments and replies --- src/listener.rs | 39 ++++++++++++----------- src/system.rs | 97 +++++++++++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 101 insertions(+), 35 deletions(-) (limited to 'src') diff --git a/src/listener.rs b/src/listener.rs index 9d0e062..bc9868d 100644 --- a/src/listener.rs +++ b/src/listener.rs @@ -1,5 +1,6 @@ use tokio::sync::mpsc::Sender; -use twilight_model::{channel::ChannelMention, id::{Id, marker::{ChannelMarker, MessageMarker, UserMarker}}, user::User}; +use twilight_http::Client; +use twilight_model::{channel::Message, id::{Id, marker::{ChannelMarker, MessageMarker, UserMarker}}, user::User}; use twilight_gateway::{error::ReceiveMessageError, Intents, Shard, ShardId}; use twilight_model::util::Timestamp; @@ -20,6 +21,7 @@ impl Listener { let intents = Intents::GUILD_MEMBERS | Intents::GUILD_PRESENCES | Intents::GUILD_MESSAGES | Intents::MESSAGE_CONTENT; let mut shard = Shard::new(ShardId::ONE, self.config.discord_token.clone(), intents); + let mut client = Client::new(self.config.discord_token.clone()); loop { match shard.next_event().await { @@ -29,40 +31,44 @@ impl Listener { println!("Bot started for {}#{}", client.user.name, client.user.discriminator); }, - twilight_gateway::Event::MessageCreate(message) => { + twilight_gateway::Event::MessageCreate(message_create) => { + let message = message_create.0; + if message.author.id != self.reference_user_id { continue } if let Err(_) = channel.send(ClientEvent::Message { event_time: message.timestamp, - message_id: message.id, - channel_id: message.channel_id, - author: message.author.clone(), - content: message.content.clone() + message }).await { println!("Client listener error: System context has already closed"); return } }, - twilight_gateway::Event::MessageUpdate(message) => { - if message.author.is_none() || message.author.as_ref().unwrap().id != self.reference_user_id { + twilight_gateway::Event::MessageUpdate(message_update) => { + if message_update.author.is_none() || message_update.author.as_ref().unwrap().id != self.reference_user_id { continue } - if message.edited_timestamp.is_none() { + if message_update.edited_timestamp.is_none() { println!("Message update but no edit timestamp"); continue; } + if message_update.content.is_none() { + println!("Message update but no content"); + continue; + } + + let message = client.message(message_update.channel_id, message_update.id) + .await.expect("Could not load message") + .model().await.expect("Could not deserialize message"); if let Err(_) = channel.send(ClientEvent::Message { - event_time: message.edited_timestamp.unwrap(), - message_id: message.id, - channel_id: message.channel_id, - author: message.author.unwrap(), - content: message.content.unwrap() + event_time: message_update.edited_timestamp.unwrap(), + message, }).await { println!("Client listener error: System context has already closed"); return @@ -91,10 +97,7 @@ impl Listener { pub enum ClientEvent { Message { event_time: Timestamp, - message_id: Id, - channel_id: Id, - author: User, - content: String, + message: Message, }, Error(ReceiveMessageError) } diff --git a/src/system.rs b/src/system.rs index 3fe4195..cf59f0d 100644 --- a/src/system.rs +++ b/src/system.rs @@ -1,9 +1,11 @@ use std::{collections::HashMap, num::NonZeroUsize, str::FromStr}; use tokio::sync::mpsc::channel; +use futures::future::join_all; use twilight_http::Client; -use twilight_model::id::{Id, marker::{ChannelMarker, MessageMarker, UserMarker}}; +use twilight_model::{channel::{message::MessageType, Message}, id::{Id, marker::{ChannelMarker, MessageMarker, UserMarker}}}; use twilight_model::util::Timestamp; +use twilight_model::http::attachment::Attachment; use crate::{config::{AutoproxyConfig, Member, MemberName}, listener::{Listener, ClientEvent}}; @@ -49,9 +51,9 @@ impl System { loop { match rx.recv().await { Some(event) => match event { - ClientEvent::Message { event_time, message_id, channel_id, content, author: _ } => { - if self.is_new_message(message_id, event_time) { - self.handle_message(message_id, channel_id, content, event_time).await; + ClientEvent::Message { event_time, message } => { + if self.is_new_message(message.id, event_time) { + self.handle_message(message, event_time).await; } }, ClientEvent::Error(_err) => { @@ -78,7 +80,7 @@ impl System { } } - async fn handle_message(&mut self, message_id: Id, channel_id: Id, content: String, timestamp: Timestamp) { + async fn handle_message(&mut self, message: Message, timestamp: Timestamp) { // Check for command // TODO: Commands // TODO: Escaping @@ -86,9 +88,9 @@ impl System { // TODO: Non-latching prefixes maybe? // Check for prefix - let match_prefix = self.config.members.iter().find_map(|member| Some((member, member.matches_proxy_prefix(&content)?))); + let match_prefix = self.config.members.iter().find_map(|member| Some((member, member.matches_proxy_prefix(&message)?))); if let Some((member, matched_content)) = match_prefix { - self.proxy_message(message_id, channel_id, member, matched_content).await; + self.proxy_message(&message, member, matched_content).await; self.update_autoproxy_state_after_message(member.clone(), timestamp); return } @@ -99,7 +101,7 @@ impl System { match autoproxy_config { AutoproxyConfig::Member {name} => { let member = self.config.members.iter().find(|member| member.name == *name).expect("Invalid autoproxy member name"); - self.proxy_message(message_id, channel_id, member, content.as_str()).await; + self.proxy_message(&message, member, message.content.as_str()).await; }, // TODO: Do something with the latch scope // TODO: Do something with presence setting @@ -107,7 +109,7 @@ impl System { if let Some((member, last_timestamp)) = &self.latch_state { let time_since_last = timestamp.as_secs() - last_timestamp.as_secs(); if time_since_last <= (*timeout_seconds).into() { - self.proxy_message(message_id, channel_id, &member, content.as_str()).await; + self.proxy_message(&message, &member, message.content.as_str()).await; self.latch_state = Some((member.clone(), timestamp)); } } @@ -116,13 +118,43 @@ impl System { } } - async fn proxy_message(&self, message_id: Id, channel_id: Id, member: &Member, content: &str) { + async fn proxy_message(&self, message: &Message, member: &Member, content: &str) { let client = self.clients.get(&member.name).expect("No client for member"); - if let Ok(_) = client.create_message(channel_id) - .content(content).expect("Cannot set content").await { - client.delete_message(channel_id, message_id).await.expect("Could not delete message"); - } + if let Ok(_) = self.duplicate_message(message, client, content).await { + client.delete_message(message.channel_id, message.id).await.expect("Could not delete message"); + } + } + + async fn duplicate_message(&self, message: &Message, client: &Client, content: &str) -> Result { + let mut create_message = client.create_message(message.channel_id) + .content(content)?; + + if message.kind == MessageType::Reply { + create_message = create_message.reply( + message.referenced_message.as_ref().expect("Message was reply but no referenced message").id + ); + } + + let attachments = join_all(message.attachments.iter().map(|attachment| async { + let filename = attachment.filename.clone(); + let description = attachment.description.clone(); + let bytes = reqwest::get(attachment.proxy_url.clone()).await?.bytes().await?; + + // TODO: keep description + Ok(Attachment::from_bytes(filename, bytes.try_into().unwrap(), attachment.id.into())) + })).await.iter().filter_map(|result: &Result| match result { + Ok(attachment) => Some(attachment.clone()), + Err(_) => None, + }).collect::>(); + + if attachments.len() > 0 { + create_message = create_message.attachments(attachments.as_slice())?; + } + + let new_message = create_message.await?.model().await?; + + Ok(new_message) } fn update_autoproxy_state_after_message(&mut self, member: Member, timestamp: Timestamp) { @@ -139,13 +171,13 @@ impl System { impl crate::config::Member { - pub fn matches_proxy_prefix<'a>(&self, content: &'a String) -> Option<&'a str> { - match self.message_pattern.captures(content.as_str()) { + pub fn matches_proxy_prefix<'a>(&self, message: &'a Message) -> Option<&'a str> { + match self.message_pattern.captures(message.content.as_str()) { None => None, Some(captures) => { let full_match = captures.get(0).unwrap(); - if full_match.len() != content.len() { + if full_match.len() != message.content.len() { return None } @@ -157,3 +189,34 @@ impl crate::config::Member { } } } + +enum MessageDuplicateError { + MessageValidation(twilight_validate::message::MessageValidationError), + AttachmentRequest(reqwest::Error), + MessageCreate(twilight_http::error::Error), + ResponseDeserialization(twilight_http::response::DeserializeBodyError) +} + +impl From for MessageDuplicateError { + fn from(value: twilight_validate::message::MessageValidationError) -> Self { + MessageDuplicateError::MessageValidation(value) + } +} + +impl From for MessageDuplicateError { + fn from(value: reqwest::Error) -> Self { + MessageDuplicateError::AttachmentRequest(value) + } +} + +impl From for MessageDuplicateError { + fn from(value: twilight_http::error::Error) -> Self { + MessageDuplicateError::MessageCreate(value) + } +} + +impl From for MessageDuplicateError { + fn from(value: twilight_http::response::DeserializeBodyError) -> Self { + MessageDuplicateError::ResponseDeserialization(value) + } +} -- cgit 1.4.1