diff --git a/src/api/mod.rs b/src/api/mod.rs index 149b974..a76d5bc 100644 --- a/src/api/mod.rs +++ b/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; diff --git a/src/api/pluralkit.rs b/src/api/pluralkit.rs new file mode 100644 index 0000000..fd4dad5 --- /dev/null +++ b/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 { + 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::().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}", + )) + } +} diff --git a/src/handlers/event/mod.rs b/src/handlers/event/mod.rs index 8db6df5..07a82d0 100644 --- a/src/handlers/event/mod.rs +++ b/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?; } diff --git a/src/handlers/event/pluralkit.rs b/src/handlers/event/pluralkit.rs new file mode 100644 index 0000000..b44aff0 --- /dev/null +++ b/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 { + 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(()) +} diff --git a/src/main.rs b/src/main.rs index b26863d..a4e5daf 100644 --- a/src/main.rs +++ b/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, - redis: redis::Client, + config: Config, + storage: Storage, octocrab: Arc, } impl Data { pub fn new() -> Result { 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, }) } diff --git a/src/storage.rs b/src/storage.rs new file mode 100644 index 0000000..6cec18f --- /dev/null +++ b/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 { + 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 { + 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) + } +}