Refactor bot into separate client + gateway

3 files changed, 390 insertions, 0 deletions

diff --git a/src/system/bot/client.rs b/src/system/bot/client.rs
new file mode 100644
index 0000000..c55759a
--- /dev/null
+++ b/src/system/bot/client.rs
@@ -0,0 +1,187 @@
+use std::sync::Arc;
+use futures::future::join_all;
+use tokio::sync::RwLock;
+use tokio::sync::Mutex;
+use twilight_http::client::Client as TwiClient;
+use twilight_http::error::Error as TwiError;
+use twilight_http::request::channel::reaction::RequestReactionType;
+use twilight_model::channel::message::{AllowedMentions, MentionType, MessageType};
+use twilight_model::http::attachment::Attachment;
+
+use super::*;
+
+pub struct Client {
+    client: Arc<Mutex<TwiClient>>,
+    bot_conf: Arc<RwLock<BotConfig>>,
+}
+
+impl Client {
+    pub fn new(discord_token: &String, bot_conf: &Arc<RwLock<BotConfig>>) -> Self {
+        Self {
+            client: Arc::new(Mutex::new(TwiClient::new(discord_token.clone()))),
+            bot_conf: bot_conf.clone(),
+        }
+    }
+
+    pub async fn refetch_message(&self, message_id: MessageId, channel_id: ChannelId) {
+        let client = self.client.lock().await;
+        let bot_conf = self.bot_conf.read().await;
+        let message_channel = bot_conf.message_handler.as_ref().expect("No message handler");
+
+        let message = client
+            .message(channel_id, message_id)
+            .await
+            .expect("Could not load message")
+            .model()
+            .await
+            .expect("Could not deserialize message");
+
+        let timestamp = if message.edited_timestamp.is_some() {
+            message.edited_timestamp.unwrap()
+        } else {
+            message.timestamp
+        };
+
+        message_channel
+            .send((timestamp, Message::Complete(message)))
+            .await;
+    }
+
+    pub async fn delete_message(&self, channel_id: ChannelId, message_id: MessageId) -> Result<(), TwiError> {
+        let client = self.client.lock().await;
+        let delete_result = client.delete_message(channel_id, message_id).await;
+        let member_id = self.bot_conf.read().await.member_id;
+
+        match delete_result {
+            Err(err) => {
+                match &err.kind() {
+                    twilight_http::error::ErrorType::Response { body: _, error, status: _ } => match error {
+                        twilight_http::api_error::ApiError::General(err) => {
+                            // Code for "Missing Permissions": https://discord.com/developers/docs/topics/opcodes-and-status-codes#json-json-error-codes
+                            if err.code == 50013 {
+                                println!("ERROR: Client {} doesn't have permissions to delete message", member_id);
+                                let _ = client.create_reaction(
+                                    channel_id,
+                                    message_id,
+                                    &RequestReactionType::Unicode { name: "🔐" }
+                                ).await;
+                            }
+                        },
+                        _ => (),
+                    },
+                    _ => (),
+                };
+
+                Err(err)
+            },
+            _ => Ok(()),
+        }
+    }
+
+    pub async fn duplicate_message(&self, message: &TwiMessage, content: &str) -> Result<TwiMessage, MessageDuplicateError> {
+        let client = self.client.lock().await;
+
+        let mut create_message = client.create_message(message.channel_id).content(content)?;
+
+        let mut allowed_mentions = AllowedMentions {
+            parse: Vec::new(),
+            replied_user: false,
+            roles: message.mention_roles.clone(),
+            users: message.mentions.iter().map(|user| user.id).collect(),
+        };
+
+        if message.mention_everyone {
+            allowed_mentions.parse.push(MentionType::Everyone);
+        }
+
+        if message.kind == MessageType::Reply {
+            if let Some(ref_message) = message.referenced_message.as_ref() {
+                create_message = create_message.reply(ref_message.id);
+
+                let pings_referenced_author = message
+                    .mentions
+                    .iter()
+                    .any(|user| user.id == ref_message.author.id);
+
+                if pings_referenced_author {
+                    allowed_mentions.replied_user = true;
+                } else {
+                    allowed_mentions.replied_user = false;
+                }
+            } else {
+                panic!("Cannot proxy message: Was reply but no referenced message");
+            }
+        }
+
+        let attachments = join_all(message.attachments.iter().map(|attachment| async {
+            let filename = attachment.filename.clone();
+            let description_opt = attachment.description.clone();
+            let bytes = reqwest::get(attachment.proxy_url.clone())
+                .await?
+                .bytes()
+                .await?;
+            let mut new_attachment =
+                Attachment::from_bytes(filename, bytes.try_into().unwrap(), attachment.id.into());
+
+            if let Some(description) = description_opt {
+                new_attachment.description(description);
+            }
+
+            Ok(new_attachment)
+        }))
+        .await
+        .iter()
+        .filter_map(
+            |result: &Result<Attachment, MessageDuplicateError>| match result {
+                Ok(attachment) => Some(attachment.clone()),
+                Err(_) => None,
+            },
+        )
+        .collect::<Vec<_>>();
+
+        if attachments.len() > 0 {
+            create_message = create_message.attachments(attachments.as_slice())?;
+        }
+
+        if let Some(flags) = message.flags {
+            create_message = create_message.flags(flags);
+        }
+
+        create_message = create_message.allowed_mentions(Some(&allowed_mentions));
+        let new_message = create_message.await?.model().await?;
+
+        Ok(new_message)
+    }
+}
+
+#[derive(Debug)]
+pub enum MessageDuplicateError {
+    MessageValidation(twilight_validate::message::MessageValidationError),
+    AttachmentRequest(reqwest::Error),
+    MessageCreate(twilight_http::error::Error),
+    ResponseDeserialization(twilight_http::response::DeserializeBodyError),
+}
+
+impl From<twilight_validate::message::MessageValidationError> for MessageDuplicateError {
+    fn from(value: twilight_validate::message::MessageValidationError) -> Self {
+        MessageDuplicateError::MessageValidation(value)
+    }
+}
+
+impl From<reqwest::Error> for MessageDuplicateError {
+    fn from(value: reqwest::Error) -> Self {
+        MessageDuplicateError::AttachmentRequest(value)
+    }
+}
+
+impl From<twilight_http::error::Error> for MessageDuplicateError {
+    fn from(value: twilight_http::error::Error) -> Self {
+        MessageDuplicateError::MessageCreate(value)
+    }
+}
+
+impl From<twilight_http::response::DeserializeBodyError> for MessageDuplicateError {
+    fn from(value: twilight_http::response::DeserializeBodyError) -> Self {
+        MessageDuplicateError::ResponseDeserialization(value)
+    }
+}
diff --git a/src/system/bot/gateway.rs b/src/system/bot/gateway.rs
new file mode 100644
index 0000000..4a83086
--- /dev/null
+++ b/src/system/bot/gateway.rs
@@ -0,0 +1,123 @@
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use tokio::sync::Mutex;
+use twilight_model::gateway::OpCode;
+use twilight_model::gateway::payload::outgoing::{update_presence::UpdatePresencePayload, UpdatePresence};
+use twilight_gateway::{
+    Intents, Shard, ShardId, 
+};
+
+use super::{Message, Status, SystemEvent, BotConfig};
+
+pub struct Gateway {
+    shard: Arc<Mutex<Shard>>,
+    bot_conf: Arc<RwLock<BotConfig>>,
+}
+
+impl Gateway {
+    pub fn new(discord_token: &String, bot_conf: &Arc<RwLock<BotConfig>>) -> Self {
+        let intents = Intents::GUILD_MEMBERS
+            | Intents::GUILD_PRESENCES
+            | Intents::GUILD_MESSAGES
+            | Intents::MESSAGE_CONTENT;
+
+        Self {
+            shard: Arc::new(Mutex::new(Shard::new(
+                ShardId::ONE,
+                discord_token.clone(),
+                intents,
+            ))),
+            bot_conf: bot_conf.clone(),
+        }
+    }
+
+    pub async fn set_status(&self, status: Status) {
+        {
+            let last_status = { (*self.bot_conf.read().await).last_status };
+
+            if status == last_status {
+                return
+            }
+        }
+
+
+        {
+            let mut shard = self.shard.lock().await;
+
+            shard.command(&UpdatePresence {
+                d: UpdatePresencePayload {
+                    activities: Vec::new(),
+                    afk: false,
+                    since: None,
+                    status,
+                },
+                op: OpCode::PresenceUpdate,
+            }).await.expect("Could not send command to gateway");
+        }
+
+        self.bot_conf.write().await.last_status = status;
+    }
+
+    pub fn start_listening(&self) {
+        let bot_conf = self.bot_conf.clone();
+        let shard = self.shard.clone();
+        tokio::spawn(async move {
+            loop {
+                let bot_conf = { (*bot_conf.read().await).clone() };
+                let next_event = { shard.lock().await.next_event().await };
+                let system_channel = bot_conf.system_handler.as_ref().expect("No system channel");
+                let message_channel = bot_conf.message_handler.as_ref().expect("No message channel");
+
+                match next_event {
+                    Err(source) => {
+                        system_channel
+                            .send(SystemEvent::GatewayError(bot_conf.member_id, source.to_string()))
+                            .await;
+
+                        if source.is_fatal() {
+                            system_channel.send(SystemEvent::GatewayClosed(bot_conf.member_id)).await;
+                            break;
+                        }
+                    }
+                    Ok(event) => match event {
+                        twilight_gateway::Event::Ready(_) => {
+                            system_channel
+                                .send(SystemEvent::GatewayConnected(bot_conf.member_id))
+                                .await;
+                        }
+
+                        twilight_gateway::Event::MessageCreate(message_create) => {
+                            let message = message_create.0;
+
+                            if message.author.id != bot_conf.reference_user_id {
+                                continue;
+                            }
+
+                            message_channel
+                                .send((message.timestamp, Message::Complete(message)))
+                                .await;
+                        }
+
+                        twilight_gateway::Event::MessageUpdate(message_update) => {
+                            if message_update.author.is_none()
+                                || message_update.author.as_ref().unwrap().id != bot_conf.reference_user_id
+                            {
+                                continue;
+                            }
+
+                            if message_update.edited_timestamp.is_none() || message_update.content.is_none() {
+                                continue;
+                            }
+
+                            message_channel
+                                .send((message_update.edited_timestamp.unwrap(), Message::Partial(*message_update, bot_conf.member_id)))
+                                .await;
+                        }
+
+                        _ => (),
+                    },
+                };
+            }
+        });
+    }
+}
diff --git a/src/system/bot/mod.rs b/src/system/bot/mod.rs
new file mode 100644
index 0000000..6bf8d78
--- /dev/null
+++ b/src/system/bot/mod.rs
@@ -0,0 +1,80 @@
+mod client;
+mod gateway;
+
+use std::sync::Arc;
+use tokio::sync::mpsc::Sender;
+use tokio::sync::RwLock;
+use twilight_http::error::Error as TwiError;
+
+pub use super::types::*;
+pub use client::MessageDuplicateError;
+use gateway::Gateway;
+use client::Client;
+
+#[derive(Clone)]
+pub struct BotConfig {
+    pub member_id: MemberId,
+    pub reference_user_id: UserId,
+    pub discord_token: String,
+    pub last_status: Status,
+    pub message_handler: Option<Sender<MessageEvent>>,
+    pub system_handler: Option<Sender<SystemEvent>>,
+}
+
+pub struct Bot {
+    bot_conf: Arc<RwLock<BotConfig>>,
+    gateway: Gateway,
+    client: Client,
+}
+
+impl Bot {
+    pub fn new(
+        member_id: MemberId,
+        config: &crate::config::Member,
+        reference_user_id: UserId,
+    ) -> Self {
+        let bot_conf = Arc::new(RwLock::new(BotConfig {
+            member_id,
+            reference_user_id,
+            discord_token: config.discord_token.clone(),
+            last_status: Status::Online,
+            message_handler: None,
+            system_handler: None,
+        }));
+
+        Self {
+            gateway: Gateway::new(&config.discord_token, &bot_conf),
+            client: Client::new(&config.discord_token, &bot_conf),
+            bot_conf,
+        }
+    }
+
+    pub async fn set_message_handler(&mut self, handler: Sender<MessageEvent>) {
+        self.bot_conf.write().await.message_handler = Some(handler);
+    }
+
+    pub async fn set_system_handler(&mut self, handler: Sender<SystemEvent>) {
+        self.bot_conf.write().await.system_handler = Some(handler);
+    }
+
+    pub async fn set_status(&self, status: Status) {
+        self.gateway.set_status(status).await;
+    }
+
+    pub fn start(&self) {
+        self.gateway.start_listening()
+    }
+
+    pub async fn refetch_message(&self, message_id: MessageId, channel_id: ChannelId) {
+        self.client.refetch_message(message_id, channel_id).await;
+    }
+
+    pub async fn delete_message(&self, channel_id: ChannelId, message_id: MessageId) -> Result<(), TwiError> {
+        self.client.delete_message(channel_id, message_id).await
+    }
+
+    pub async fn duplicate_message(&self, message_id: &TwiMessage, content: &str) ->  Result<TwiMessage, MessageDuplicateError> {
+        self.client.duplicate_message(message_id, content).await
+    }
+}
+