feat: reintroduce PK support

Signed-off-by: seth <getchoo@tuta.io>

src/api/mod.rs

@@ -1,6 +1,7 @@
 use once_cell::sync::Lazy;
 
 pub mod dadjoke;
+pub mod pluralkit;
 pub mod prism_meta;
 pub mod rory;
 

src/api/pluralkit.rs

@@ -0,0 +1,37 @@
+use crate::api::REQWEST_CLIENT;
+
+use color_eyre::eyre::{eyre, Context, Result};
+use log::*;
+use poise::serenity_prelude::{MessageId, UserId};
+use reqwest::StatusCode;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct PluralKitMessage {
+    pub sender: String,
+}
+
+const PLURAL_KIT: &str = "https://api.pluralkit.me/v2";
+const MESSAGES_ENDPOINT: &str = "/messages";
+
+pub async fn get_sender(message_id: MessageId) -> Result<UserId> {
+    let req = REQWEST_CLIENT
+        .get(format!("{PLURAL_KIT}{MESSAGES_ENDPOINT}/{message_id}"))
+        .build()?;
+
+    info!("Making request to {}", req.url());
+    let resp = REQWEST_CLIENT.execute(req).await?;
+    let status = resp.status();
+
+    if let StatusCode::OK = status {
+        let data = resp.json::<PluralKitMessage>().await?;
+        let id: u64 = data.sender.parse().wrap_err_with(|| format!("Couldn't parse response from PluralKit as a UserId! Here's the response:\n{data:#?}"))?;
+        let sender = UserId::from(id);
+
+        Ok(sender)
+    } else {
+        Err(eyre!(
+            "Failed to get PluralKit message information from {PLURAL_KIT}{MESSAGES_ENDPOINT}/{message_id} with {status}",
+        ))
+    }
+}

src/handlers/event/mod.rs

@@ -8,13 +8,14 @@ use poise::{Event, FrameworkContext};
 mod delete;
 mod eta;
 mod expand_link;
+pub mod pluralkit;
 mod support_onboard;
 
 pub async fn handle(
     ctx: &Context,
     event: &Event<'_>,
     _framework: FrameworkContext<'_, Data, Report>,
-    _data: &Data,
+    data: &Data,
 ) -> Result<()> {
     match event {
         Event::Ready { data_about_bot } => {
@@ -35,6 +36,16 @@ pub async fn handle(
                 return Ok(());
             }
 
+            // detect PK users first to make sure we don't respond to unproxied messages
+            pluralkit::handle(ctx, new_message, data).await?;
+
+            if data.storage.is_user_plural(new_message.author.id).await?
+                && pluralkit::is_message_proxied(new_message).await?
+            {
+                debug!("Not replying to unproxied PluralKit message");
+                return Ok(());
+            }
+
             eta::handle(ctx, new_message).await?;
             expand_link::handle(ctx, new_message).await?;
         }

src/handlers/event/pluralkit.rs

@@ -0,0 +1,42 @@
+use crate::{api, Data};
+use std::time::Duration;
+
+use color_eyre::eyre::Result;
+use log::*;
+use poise::serenity_prelude::{Context, Message};
+use tokio::time::sleep;
+
+const PK_DELAY_SEC: Duration = Duration::from_secs(1000);
+
+pub async fn is_message_proxied(message: &Message) -> Result<bool> {
+    debug!(
+        "Waiting on PluralKit API for {} seconds",
+        PK_DELAY_SEC.as_secs()
+    );
+    sleep(PK_DELAY_SEC).await;
+
+    let proxied = api::pluralkit::get_sender(message.id).await.is_ok();
+
+    Ok(proxied)
+}
+
+pub async fn handle(_ctx: &Context, msg: &Message, data: &Data) -> Result<()> {
+    if msg.webhook_id.is_some() {
+        debug!(
+            "Message {} has a webhook ID. Checking if it was sent through PluralKit",
+            msg.id
+        );
+
+        debug!(
+            "Waiting on PluralKit API for {} seconds",
+            PK_DELAY_SEC.as_secs()
+        );
+        sleep(PK_DELAY_SEC).await;
+
+        if let Ok(sender) = api::pluralkit::get_sender(msg.id).await {
+            data.storage.store_user_plurality(sender).await?;
+        }
+    }
+
+    Ok(())
+}

src/main.rs

@@ -6,12 +6,14 @@ use log::*;
 use poise::{
     serenity_prelude as serenity, EditTracker, Framework, FrameworkOptions, PrefixFrameworkOptions,
 };
+use storage::Storage;
 
 mod api;
 mod commands;
 mod config;
 mod consts;
 mod handlers;
+mod storage;
 mod tags;
 mod utils;
 
@@ -19,20 +21,20 @@ type Context<'a> = poise::Context<'a, Data, Report>;
 
 #[derive(Clone)]
 pub struct Data {
-    config: config::Config,
+    config: Config,
-    redis: redis::Client,
+    storage: Storage,
     octocrab: Arc<octocrab::Octocrab>,
 }
 
 impl Data {
     pub fn new() -> Result<Self> {
         let config = Config::new_from_env()?;
-        let redis = redis::Client::open(config.redis_url.clone())?;
+        let storage = Storage::new(&config.redis_url)?;
         let octocrab = octocrab::instance();
 
         Ok(Self {
             config,
-            redis,
+            storage,
             octocrab,
         })
     }

src/storage.rs

@@ -0,0 +1,40 @@
+use color_eyre::eyre::Result;
+use log::*;
+use poise::serenity_prelude::UserId;
+use redis::{AsyncCommands as _, Client};
+
+pub const USER_KEY: &str = "users-v1";
+
+#[derive(Clone, Debug)]
+pub struct Storage {
+    client: Client,
+}
+
+impl Storage {
+    pub fn new(redis_url: &str) -> Result<Self> {
+        let client = Client::open(redis_url)?;
+
+        Ok(Self { client })
+    }
+
+    pub async fn store_user_plurality(&self, sender: UserId) -> Result<()> {
+        let mut con = self.client.get_async_connection().await?;
+
+        info!("Marking {sender} as a PluralKit user");
+        let key = format!("{USER_KEY}:{sender}:pk");
+
+        // Just store some value. We only care about the presence of this key
+        con.set(&key, 0).await?;
+        con.expire(key, 7 * 24 * 60 * 60).await?;
+
+        Ok(())
+    }
+
+    pub async fn is_user_plural(&self, user_id: UserId) -> Result<bool> {
+        let key = format!("{USER_KEY}:{user_id}:pk");
+        let mut con = self.client.get_async_connection().await?;
+
+        let exists: bool = con.exists(key).await?;
+        Ok(exists)
+    }
+}