From 23169ef19c362c14d896eb6d6a31d2928c400898 Mon Sep 17 00:00:00 2001 From: bluepython508 <16466646+bluepython508@users.noreply.github.com> Date: Sun, 3 Nov 2024 21:00:19 +0000 Subject: [PATCH] Functional bot: responds to commands, moves users --- .envrc | 5 ++ Cargo.lock | 156 ++++++++++++++++++++++++++++++++++++++++--- Cargo.toml | 5 ++ src/main.rs | 189 +++++++++++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 343 insertions(+), 12 deletions(-) diff --git a/.envrc b/.envrc index 3550a30..9f972e3 100644 --- a/.envrc +++ b/.envrc @@ -1 +1,6 @@ use flake + +export DISCORD_TOKEN_FILE=$(mktemp) +cat >$DISCORD_TOKEN_FILE <>>); + +type Context<'a> = poise::Context<'a, Data, Error>; + +async fn run(ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId) -> Result<()> { + let guild = ctx + .guild_id() + .ok_or_eyre("This bot only works in servers")?; + + let channels = guild.channels(&ctx).await?; + let channels_for = |id| { + let channel = channels + .get(&id) + .ok_or_eyre(format!("The channel {id:?} doesn't exist"))?; + Ok::<_, Error>(match channel.kind { + serenity::ChannelType::Voice => vec![channel], + serenity::ChannelType::Category => channels + .values() + .filter(|chan| { + chan.kind == serenity::ChannelType::Voice && chan.parent_id == Some(channel.id) + }) + .collect(), + ty => bail!( + "This bot doesn't work with non-voice channels: {:?} has type {:?}", + &channel.name, + ty + ), + }) + }; + let mut users = channels_for(from)? + .into_iter() + .flat_map(|chan| chan.members(ctx).unwrap()) + .collect::>(); + users.shuffle(&mut rand::thread_rng()); + + let to = channels_for(to)?; + + futures::future::try_join_all( + to.into_iter() + .cycle() + .zip(&users) + .map(|(channel, user)| guild.move_member(ctx, user.user.id, channel)), + ) + .await?; + + ctx.reply(format!("Moved {} users", users.len())).await?; + Ok(()) +} + +#[poise::command(slash_command, prefix_command, ephemeral)] +async fn dusk(ctx: Context<'_>) -> Result<()> { + let guild = ctx + .guild_id() + .ok_or_eyre("This bot only works in servers")?; + let GuildData { + town_square, + cottages, + } = ctx + .data() + .0 + .read() + .expect("No poisoned locks") + .get(&guild) + .copied() + .ok_or_eyre("This bot hasn't been configured for this server")?; + + run(ctx, town_square, cottages).await?; + + Ok(()) +} + +#[poise::command(slash_command, prefix_command, ephemeral)] +async fn dawn(ctx: Context<'_>) -> Result<()> { + let guild = ctx + .guild_id() + .ok_or_eyre("This bot only works in servers")?; + let GuildData { + town_square, + cottages, + } = ctx + .data() + .0 + .read() + .expect("No poisoned locks") + .get(&guild) + .copied() + .ok_or_eyre("This bot hasn't been configured for this server")?; + + run(ctx, cottages, town_square).await?; + + Ok(()) +} + +#[poise::command(slash_command, prefix_command, ephemeral)] +async fn configure( + ctx: Context<'_>, + #[channel_types("Voice")] town_square: serenity::ChannelId, + #[channel_types("Category")] cottages: serenity::ChannelId, +) -> Result<()> { + let guild = ctx + .guild_id() + .ok_or_eyre("This bot only works in servers")?; + + ctx.data().0.write().expect("No poisoned locks").insert( + guild, + GuildData { + town_square, + cottages, + }, + ); + ctx.reply("Configured successfully!").await?; + Ok(()) +} + +async fn register_commands( + ctx: impl AsRef, + options: &poise::FrameworkOptions, + guild: serenity::GuildId, +) -> Result<()> { + poise::builtins::register_in_guild(ctx, &options.commands, guild).await?; + Ok(()) +} + +async fn on_event( + ctx: &serenity::Context, + event: &serenity::FullEvent, + framework: poise::FrameworkContext<'_, Data, Error>, + _: &Data, +) -> Result<()> { + if let serenity::FullEvent::GuildCreate { guild, is_new: _ } = event { + register_commands(ctx, framework.options, guild.id).await?; + }; + Ok(()) +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<()> { + color_eyre::install()?; + let token = std::fs::read_to_string( + std::env::var_os("DISCORD_TOKEN_FILE").ok_or_eyre("A Discord token is required")?, + ) + .wrap_err("A Discord token is required")?; + let intents = serenity::GatewayIntents::non_privileged(); + + let framework = poise::Framework::builder() + .options(poise::FrameworkOptions { + commands: vec![dusk(), dawn(), configure()], + event_handler: |ctx, ev, ctxf, data| on_event(ctx, ev, ctxf, data).boxed(), + ..Default::default() + }) + .setup(|ctx, ready, framework| { + Box::pin(async move { + for guild in &ready.guilds { + register_commands(ctx, framework.options(), guild.id).await?; + } + Ok(Data::default()) + }) + }) + .build(); + + serenity::ClientBuilder::new(token, intents) + .framework(framework) + .await? + .start() + .await?; + + Ok(()) +}