Initial commit

src/bot/commands.rs

@@ -0,0 +1,74 @@
+use serenity::all::{
+    CommandInteraction, CommandOptionType, Context, CreateCommand, CreateCommandOption,
+    CreateInputText, CreateLabel, CreateModal, CreateModalComponent, CreateTextDisplay,
+    InputTextStyle, Permissions,
+};
+
+use crate::handler::{Handler, get_guild};
+
+pub fn config_command() -> CreateCommand<'static> {
+    CreateCommand::new("starboard")
+        .description("스타보드 설정")
+        .add_option(CreateCommandOption::new(
+            CommandOptionType::SubCommand,
+            "script",
+            "스크립트 수정",
+        ))
+        .default_member_permissions(Permissions::ADMINISTRATOR)
+}
+
+impl Handler {
+    pub async fn process_starboard_command(
+        &self,
+        ctx: &Context,
+        interaction: &CommandInteraction,
+    ) -> anyhow::Result<()> {
+        let Some(guild_id) = interaction.guild_id else {
+            return Ok(());
+        };
+
+        // if let CommandType interaction.data.kind {}
+
+        for option in interaction.data.options() {
+            if option.name == "script" {
+                let script = get_guild(&self.db, guild_id)
+                    .await?
+                    .and_then(|x| x.script)
+                    .unwrap_or("".to_string());
+
+                interaction
+                    .create_response(
+                        &ctx.http,
+                        serenity::all::CreateInteractionResponse::Modal(
+                            CreateModal::new("edit_script", "스크립트 편집").components(vec![
+                                CreateModalComponent::TextDisplay(CreateTextDisplay::new(
+                                    r#"rhai 스크립트로 스타보드 동작을 지정할 수 있습니다(반복문은 제한됩니다)
+사용 가능 변수:
+- `reactions`: map<string, i64> (이모지 - 개수)
+- `channel`: string (채널 ID)
+
+`result("채널ID", 리액션 개수, "메시지에 사용될 아이콘(이모지 권장)")` 를 리턴해주세요(unit이나 다른 값 반환 시 무시됩니다)
+
+예시 스크립트:
+```rs
+if reactions["⭐"] >= 3 {
+    return result("0000000000000000000", reactions["⭐"], "⭐")
+}
+```"#,
+                                )),
+                                CreateModalComponent::Label(CreateLabel::input_text(
+                                    "스크립트",
+                                    CreateInputText::new(InputTextStyle::Paragraph, "script")
+                                        .value(script)
+                                        .placeholder("사용하지 않으려면 () 입력"),
+                                )),
+                            ]),
+                        ),
+                    )
+                    .await?;
+            }
+        }
+
+        Ok(())
+    }
+}

src/bot/config.rs

@@ -0,0 +1,18 @@
+use secrecy::SecretString;
+use serde::Deserialize;
+
+#[derive(Deserialize, Debug)]
+pub struct Config {
+    pub bot: BotConfig,
+    pub db: DatabaseConfig,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct DatabaseConfig {
+    pub url: SecretString,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct BotConfig {
+    pub token: SecretString,
+}

src/bot/db.rs

@@ -0,0 +1,23 @@
+use sqlx::prelude::FromRow;
+
+#[derive(Debug, FromRow)]
+pub struct GuildRow {
+    pub id: String,
+    pub script: Option<String>,
+}
+
+#[derive(Debug, FromRow)]
+pub struct MessageRow {
+    // pub guild_id: String,
+    pub message_id: String,
+    pub counter_id: String,
+    pub counter_channel_id: String,
+}
+
+#[derive(Debug, FromRow)]
+pub struct WebhookRow {
+    // pub guild_id: String,
+    // pub channel_id: String,
+    pub webhook_id: String,
+    pub webhook_token: String,
+}

src/bot/handler.rs

@@ -0,0 +1,393 @@
+use std::{collections::HashMap, sync::Arc};
+
+use anyhow::{Context as _, bail};
+use dashmap::DashMap;
+use serenity::{
+    all::{
+        Attachment, CacheHttp, ChannelId, Component, Context, CreateComponent, CreateMediaGallery,
+        CreateMediaGalleryItem, CreateSeparator, CreateTextDisplay, CreateUnfurledMediaItem,
+        CreateWebhook, EditWebhookMessage, EventHandler, ExecuteWebhook, FullEvent,
+        GenericChannelId, GuildId, Interaction, InteractionType, Message, MessageFlags, MessageId,
+        MessageSnapshot, Reaction, StickerFormatType, StickerItem, Webhook, WebhookId,
+    },
+    async_trait,
+    futures::lock::Mutex,
+    small_fixed_array::{FixedArray, FixedString},
+};
+use sqlx::{PgExecutor, PgPool};
+
+use crate::{
+    commands,
+    db::{GuildRow, MessageRow, WebhookRow},
+};
+
+struct CloneBuilderMessage {
+    pub content: FixedString<u16>,
+    pub attachments: FixedArray<Attachment>,
+    pub sticker_items: FixedArray<StickerItem>,
+}
+
+impl From<Message> for CloneBuilderMessage {
+    fn from(value: Message) -> Self {
+        Self {
+            content: value.content,
+            attachments: value.attachments,
+            sticker_items: value.sticker_items,
+        }
+    }
+}
+
+impl From<MessageSnapshot> for CloneBuilderMessage {
+    fn from(value: MessageSnapshot) -> Self {
+        CloneBuilderMessage {
+            content: value.content,
+            attachments: value.attachments,
+            sticker_items: value.sticker_items,
+        }
+    }
+}
+
+pub struct Handler {
+    pub db: PgPool,
+    pub message_lock: Arc<DashMap<MessageId, Arc<Mutex<()>>>>,
+}
+
+#[async_trait]
+impl EventHandler for Handler {
+    async fn dispatch(&self, context: &Context, event: &FullEvent) {
+        match event {
+            FullEvent::Ready { data_about_bot, .. } => {
+                info!("Bot is ready as {}", data_about_bot.user.tag());
+
+                if let Err(e) = context
+                    .http
+                    .create_global_commands(&vec![commands::config_command()])
+                    .await
+                {
+                    error!("Failed to register config command: {e:?}");
+                }
+            }
+            FullEvent::ReactionAdd { add_reaction, .. } => {
+                let lock = self
+                    .message_lock
+                    .entry(add_reaction.message_id)
+                    .or_insert_with(|| Arc::new(Mutex::new(())));
+
+                {
+                    let _guard = lock.lock().await;
+
+                    if let Err(e) = self.reaction_add(context, add_reaction).await {
+                        error!("Error while processing reaction add event: {e:?}");
+                    }
+                }
+
+                if Arc::strong_count(&lock) == 2 {
+                    self.message_lock
+                        .remove_if(&add_reaction.message_id, |_, val| {
+                            Arc::strong_count(val) <= 2
+                        });
+                }
+            }
+            FullEvent::InteractionCreate { interaction } => {
+                if let Err(e) = self.interaction_create(context, interaction).await {
+                    error!("Error while processing interaction: {e:?}");
+                }
+            }
+            _ => return,
+        }
+    }
+}
+
+impl Handler {
+    async fn interaction_create(
+        &self,
+        ctx: &Context,
+        interaction: &Interaction,
+    ) -> anyhow::Result<()> {
+        if let Some(command) = interaction.as_command() {
+            if command.data.name == "starboard" {
+                self.process_starboard_command(ctx, command).await?;
+
+                return Ok(());
+            }
+        }
+
+        if let Some(modal) = interaction.as_modal_submit() {
+            self.process_modal(ctx, modal).await?;
+
+            return Ok(());
+        }
+
+        Ok(())
+    }
+
+    async fn reaction_add(&self, ctx: &Context, reaction: &Reaction) -> anyhow::Result<()> {
+        if reaction.user(&ctx.http).await?.bot() {
+            return Ok(());
+        }
+        let Some(guild_id) = reaction.guild_id else {
+            return Ok(());
+        };
+        let Some(guild) = get_guild(&self.db, guild_id).await? else {
+            return Ok(());
+        };
+        let Some(script) = &guild.script else {
+            return Ok(());
+        };
+
+        let existing_message = get_message(&self.db, reaction.message_id).await?;
+
+        let msg = reaction.message(&ctx.http).await?;
+
+        let reactions: HashMap<String, i64> = msg
+            .reactions
+            .iter()
+            .map(|x| (x.reaction_type.to_string(), x.count as i64))
+            .collect();
+
+        let reference_msg = if !msg.message_snapshots.is_empty() {
+            CloneBuilderMessage::from(msg.message_snapshots[0].clone())
+        } else {
+            CloneBuilderMessage::from(msg.clone())
+        };
+
+        let Some(result) =
+            paringboard::script::check(script, reactions, reaction.channel_id.to_string())?
+        else {
+            return Ok(());
+        };
+
+        let channel_id = result
+            .channel_id
+            .parse()
+            .context("failed to parse send channel id")?;
+        let webhook = get_or_create_webhook(&self.db, &ctx.http, guild_id, channel_id).await?;
+
+        let mut components: Vec<CreateComponent<'_>> = vec![];
+
+        if let Some(existing_message) = &existing_message {
+            let msg_id: MessageId = existing_message
+                .counter_id
+                .parse()
+                .context("unable to parse existing counter message id")?;
+            let chn_id: GenericChannelId = existing_message
+                .counter_channel_id
+                .parse()
+                .context("unable to parse counter channel id")?;
+
+            let existing_message = ctx.http.get_message(chn_id, msg_id).await?;
+
+            for comp in existing_message
+                .components
+                .iter()
+                .take(existing_message.components.len().saturating_sub(2) as _)
+            {
+                match comp {
+                    Component::TextDisplay(text) => {
+                        components.push(CreateComponent::TextDisplay(CreateTextDisplay::new(
+                            text.content.clone().unwrap(),
+                        )));
+                    }
+                    Component::Separator(separator) => {
+                        components.push(CreateComponent::Separator(CreateSeparator::new(
+                            separator.divider.unwrap(),
+                        )));
+                    }
+                    Component::MediaGallery(gallery) => {
+                        let items = gallery
+                            .items
+                            .iter()
+                            .map(|x| {
+                                let media = CreateUnfurledMediaItem::new(x.media.url.to_string());
+                                CreateMediaGalleryItem::new(media)
+                            })
+                            .collect::<Vec<_>>();
+
+                        components.push(CreateComponent::MediaGallery(CreateMediaGallery::new(
+                            items,
+                        )));
+                    }
+                    unknown => {
+                        bail!("unhandle component value received: {unknown:?}");
+                    }
+                }
+            }
+        } else {
+            if !reference_msg.content.is_empty() {
+                components.push(CreateComponent::TextDisplay(CreateTextDisplay::new(
+                    &reference_msg.content,
+                )));
+            }
+
+            let images = reference_msg
+                .attachments
+                .iter()
+                .filter(|x| x.width.is_some())
+                .collect::<Vec<_>>();
+
+            if !images.is_empty() {
+                components.push(CreateComponent::MediaGallery(CreateMediaGallery::new(
+                    images
+                        .into_iter()
+                        .map(|x| {
+                            CreateMediaGalleryItem::new(CreateUnfurledMediaItem::new(x.url.clone()))
+                        })
+                        .collect::<Vec<_>>(),
+                )));
+            }
+
+            if let Some(sticker_url) = reference_msg
+                .sticker_items
+                .iter()
+                .filter_map(|x| {
+                    let ext = match x.format_type {
+                        StickerFormatType::Png | StickerFormatType::Apng => "png",
+                        StickerFormatType::Gif => "gif",
+                        _ => return None,
+                    };
+
+                    let url = format!(
+                        "https://media.discordapp.net/stickers/{}.{}?size=160&quality=lossless",
+                        &x.id, ext
+                    );
+
+                    Some(url)
+                })
+                .next()
+            {
+                let item = CreateMediaGalleryItem::new(CreateUnfurledMediaItem::new(sticker_url));
+
+                components.push(CreateComponent::MediaGallery(CreateMediaGallery::new(
+                    vec![item],
+                )));
+            }
+        }
+
+        components.push(CreateComponent::Separator(CreateSeparator::new(true)));
+
+        components.push(CreateComponent::TextDisplay(CreateTextDisplay::new(
+            format!(
+                "-# {} {} · [원본 메시지]({})",
+                result.icon,
+                result.count,
+                msg.link().to_string()
+            ),
+        )));
+
+        if let Some(existing) = existing_message {
+            let msg_id = existing
+                .counter_id
+                .parse()
+                .context("unable to parse counter message id")?;
+            webhook
+                .edit_message(
+                    &ctx.http,
+                    msg_id,
+                    EditWebhookMessage::new().components(components),
+                )
+                .await?;
+        } else {
+            let counter_msg = webhook
+                .execute(
+                    &ctx.http,
+                    true,
+                    ExecuteWebhook::new()
+                        .username(msg.author.display_name())
+                        .avatar_url(
+                            msg.author
+                                .avatar_url()
+                                .unwrap_or_else(|| msg.author.default_avatar_url()),
+                        )
+                        .components(components)
+                        .flags(MessageFlags::IS_COMPONENTS_V2),
+                )
+                .await?
+                .unwrap();
+
+            sqlx::query(
+                    "insert into messages (guild_id, message_id, counter_id, counter_channel_id) values ($1, $2, $3, $4)",
+                )
+                .bind(guild_id.to_string())
+                .bind(msg.id.to_string())
+                .bind(counter_msg.id.to_string())
+                .bind(counter_msg.channel_id.to_string())
+                .execute(&self.db).await?;
+        }
+
+        Ok(())
+    }
+}
+
+pub async fn get_guild(
+    executor: impl PgExecutor<'_>,
+    id: GuildId,
+) -> anyhow::Result<Option<GuildRow>> {
+    Ok(
+        sqlx::query_as::<_, GuildRow>("select * from guilds where id = $1")
+            .bind(id.to_string())
+            .fetch_optional(executor)
+            .await
+            .context("failed to get guild data from db")?,
+    )
+}
+
+async fn get_message(
+    executor: impl PgExecutor<'_>,
+    message_id: MessageId,
+) -> anyhow::Result<Option<MessageRow>> {
+    Ok(
+        sqlx::query_as::<_, MessageRow>("select * from messages where message_id = $1")
+            .bind(message_id.to_string())
+            .fetch_optional(executor)
+            .await
+            .context("failed to get message data from db")?,
+    )
+}
+
+async fn get_or_create_webhook(
+    executor: impl PgExecutor<'_> + Copy,
+    http: impl CacheHttp,
+    guild_id: GuildId,
+    channel_id: ChannelId,
+) -> anyhow::Result<Webhook> {
+    let existing = sqlx::query_as::<_, WebhookRow>(
+        "select * from webhooks where guild_id = $1 and channel_id = $2",
+    )
+    .bind(guild_id.to_string())
+    .bind(channel_id.to_string())
+    .fetch_optional(executor)
+    .await
+    .context("failed to get webhook data from db")?;
+
+    if let Some(existing) = existing {
+        let webhook_id: WebhookId = existing
+            .webhook_id
+            .parse()
+            .context("failed to parse webhook id")?;
+        let webhook =
+            Webhook::from_id_with_token(http.http(), webhook_id, &existing.webhook_token).await?;
+
+        return Ok(webhook);
+    }
+
+    let webhook = channel_id
+        .create_webhook(
+            http.http(),
+            CreateWebhook::new("스타보드").audit_log_reason("스타보드 메시지 전송용 웹훅"),
+        )
+        .await
+        .context("failed to create a webhook")?;
+
+    let token = webhook
+        .token
+        .clone()
+        .context("why token does not exist???")?;
+
+    sqlx::query(
+        "insert into webhooks (guild_id, channel_id, webhook_id, webhook_token) values ($1, $2, $3, $4)",
+    ).bind(guild_id.to_string()).bind(channel_id.to_string()).bind(webhook.id.to_string()).bind(token.expose_secret())
+    .execute(executor)
+    .await?;
+
+    Ok(webhook)
+}

src/bot/main.rs

@@ -0,0 +1,63 @@
+use std::{str::FromStr, sync::Arc};
+
+use anyhow::Context;
+use dashmap::DashMap;
+use figment::{
+    Figment,
+    providers::{Env, Format, Toml},
+};
+use secrecy::ExposeSecret;
+use serenity::{
+    Client,
+    all::{GatewayIntents, Token},
+};
+use sqlx::{migrate, postgres::PgPoolOptions};
+
+use crate::{config::Config, handler::Handler};
+
+mod commands;
+mod config;
+mod db;
+mod handler;
+mod modal;
+
+#[macro_use]
+extern crate tracing;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt::init();
+
+    let config: Config = Figment::new()
+        .merge(Toml::file("config.toml"))
+        .merge(Env::prefixed("PB_"))
+        .extract()?;
+
+    info!("config: {config:?}");
+
+    let db = PgPoolOptions::new()
+        .max_connections(5)
+        .connect(config.db.url.expose_secret())
+        .await
+        .context("unable to connect to database")?;
+
+    let migrator = migrate!("./migrations");
+    migrator.run(&db).await.context("failed to migrate")?;
+
+    info!("migrated database");
+
+    let mut client = Client::builder(
+        Token::from_str(config.bot.token.expose_secret()).unwrap(),
+        GatewayIntents::GUILD_MESSAGES | GatewayIntents::GUILD_MESSAGE_REACTIONS,
+    )
+    .event_handler(Arc::new(Handler {
+        db,
+        message_lock: Arc::new(DashMap::new()),
+    }))
+    .await
+    .expect("Err creating client");
+
+    client.start().await?;
+
+    Ok(())
+}

src/bot/modal.rs

@@ -0,0 +1,64 @@
+use anyhow::Context as _;
+use serenity::all::{
+    Component, Context, CreateInteractionResponseMessage, Label, LabelComponent, ModalInteraction,
+};
+
+use crate::{db, handler::Handler};
+
+impl Handler {
+    pub async fn process_modal(
+        &self,
+        ctx: &Context,
+        interaction: &ModalInteraction,
+    ) -> anyhow::Result<()> {
+        if interaction.data.custom_id == "edit_script" {
+            return self.process_script_modal(ctx, interaction).await;
+        }
+
+        Ok(())
+    }
+
+    async fn process_script_modal(
+        &self,
+        ctx: &Context,
+        interaction: &ModalInteraction,
+    ) -> anyhow::Result<()> {
+        let Some(guild_id) = interaction.guild_id else {
+            return Ok(());
+        };
+
+        let components = &interaction.data.components;
+
+        let value = components
+            .get(1)
+            .and_then(|x| match x {
+                Component::Label(label) => Some(label),
+                _ => None,
+            })
+            .and_then(|x| match &x.component {
+                LabelComponent::InputText(input) => Some(input),
+                _ => None,
+            })
+            .and_then(|x| x.value.clone())
+            .context("unable to get value text from input")?;
+
+        sqlx::query("update guilds set script = $1 where id = $2")
+            .bind(value.as_str())
+            .bind(guild_id.to_string())
+            .execute(&self.db)
+            .await?;
+
+        interaction
+            .create_response(
+                &ctx.http,
+                serenity::all::CreateInteractionResponse::Message(
+                    CreateInteractionResponseMessage::new()
+                        .content("스타보드 스크립트가 설정되었습니다.")
+                        .ephemeral(true),
+                ),
+            )
+            .await?;
+
+        Ok(())
+    }
+}

src/lib.rs

@@ -0,0 +1 @@
+pub mod script;

src/script.rs

@@ -0,0 +1,63 @@
+use std::collections::HashMap;
+
+use anyhow::anyhow;
+use rhai::{Dynamic, Engine, Map, Scope};
+use serde::Deserialize;
+
+#[derive(Debug, Deserialize, Clone)]
+pub struct ReactionResult {
+    pub channel_id: String,
+    pub count: i64,
+    pub icon: String,
+}
+
+pub fn check(
+    script: &str,
+    reactions: HashMap<String, i64>,
+    channel: String,
+) -> anyhow::Result<Option<ReactionResult>> {
+    let mut engine = Engine::new();
+    engine.set_max_operations(1000);
+    engine.register_type::<ReactionResult>();
+    engine.register_fn("result", |webhook_url: String, count: i64, icon: String| {
+        ReactionResult {
+            channel_id: webhook_url,
+            count,
+            icon,
+        }
+    });
+
+    engine
+        .disable_symbol("for")
+        .disable_symbol("while")
+        .disable_symbol("loop")
+        .set_max_expr_depths(50, 5)
+        .set_max_string_size(60)
+        .set_max_map_size(512)
+        .set_max_array_size(512);
+
+    let mut emotes_input = Map::new();
+    for (emoji, count) in reactions {
+        emotes_input.insert(emoji.into(), Dynamic::from(count));
+    }
+
+    let mut scope = Scope::new();
+    scope.set_value("reactions", Dynamic::from(emotes_input));
+    scope.set_value("channel", Dynamic::from(channel));
+
+    let result: Dynamic = engine
+        .eval_with_scope(&mut scope, script)
+        .map_err(|e| anyhow!("{e:?}"))?;
+
+    let result: Option<ReactionResult> = if result.is_unit() {
+        None
+    } else {
+        Some(
+            result
+                .try_cast_result()
+                .map_err(|e| anyhow!("failed to cast: {e:?}"))?,
+        )
+    };
+
+    return Ok(result);
+}