use std::{future::Future, io, path::Path}; use ::serenity::all::{ChannelId, GuildId, Mentionable, UserId}; use eyre::{Context as _, Error, OptionExt, Result}; use futures::{FutureExt, StreamExt, TryStreamExt}; use poise::serenity_prelude as serenity; use rand::seq::SliceRandom; use rusqlite::OptionalExtension; use tracing_subscriber::{ fmt::{self, format::FmtSpan}, layer::SubscriberExt, EnvFilter, }; #[derive(Copy, Clone)] struct GuildData { guild: serenity::GuildId, town_square: serenity::ChannelId, cottages: serenity::ChannelId, st: serenity::RoleId, currently_playing: serenity::RoleId, } struct Data(tokio::sync::Mutex); impl std::fmt::Debug for Data { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "Data {{}}") } } impl Data { #[tracing::instrument(skip(path))] fn open(path: impl AsRef) -> Result { let conn = rusqlite::Connection::open(path)?; conn.execute_batch("CREATE TABLE IF NOT EXISTS guilds(guild TEXT PRIMARY KEY NOT NULL, town_square TEXT NOT NULL, cottages TEXT NOT NULL, st TEXT, currently_playing TEXT);")?; Ok(Self(tokio::sync::Mutex::new(conn))) } #[tracing::instrument] async fn get(&self, guild: serenity::GuildId) -> Result { self.0 .lock() .await .prepare_cached("SELECT guild, town_square, cottages, st, currently_playing FROM guilds WHERE guild = ?;")? .query_row(rusqlite::params![guild.to_string()], |row| { Ok(GuildData { guild: row.get_unwrap::<_, String>("guild").parse().unwrap(), town_square: row.get_unwrap::<_, String>("town_square").parse().unwrap(), cottages: row.get_unwrap::<_, String>("cottages").parse().unwrap(), st: row.get_unwrap::<_, String>("st").parse().unwrap(), currently_playing: row.get_unwrap::<_, String>("currently_playing").parse().unwrap(), }) }) .optional()? .ok_or_eyre("This bot has not been configured in this server!") } #[tracing::instrument] async fn insert( &self, GuildData { guild, town_square, cottages, st, currently_playing, }: GuildData, ) -> Result<()> { self.0 .lock() .await .prepare_cached("INSERT INTO guilds (guild, town_square, cottages, st, currently_playing) VALUES (?, ?, ?, ?, ?) ON CONFLICT(guild) DO UPDATE SET town_square = excluded.town_square, cottages = excluded.cottages, st = excluded.st, currently_playing = excluded.currently_playing; ")? .execute(rusqlite::params![ guild.to_string(), town_square.to_string(), cottages.to_string(), st.to_string(), currently_playing.to_string(), ])?; Ok(()) } } type Context<'a> = poise::Context<'a, Data, Error>; #[tracing::instrument] async fn channel_children( ctx: &Context<'_>, channel: serenity::ChannelId, ) -> Result> { let channels = ctx .guild_id() .ok_or_eyre("This bot only works in servers")? .channels(&ctx) .await?; Ok(channels .into_values() .filter(|chan| chan.parent_id == Some(channel)) .collect()) } #[tracing::instrument] fn move_users<'a>( ctx: Context<'a>, guild: GuildId, users: Vec<(ChannelId, UserId)>, ) -> impl Future> + Send + 'a { futures::stream::iter(users.into_iter().map(move |(channel, user)| async move { tracing::info!(?channel, ?user, "Moving user"); guild.move_member(ctx, user, channel).await?; tracing::info!(?channel, ?user, "Moved user"); Ok(()) })) .buffer_unordered(10) .try_collect::<()>() } #[tracing::instrument] async fn spread( ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.defer_ephemeral().await?; let from = from .to_channel(ctx) .await? .guild() .ok_or_eyre("This bot only works in servers")?; let mut to = channel_children(&ctx, to).await?; let mut users = from.members(ctx)?; users.shuffle(&mut rand::thread_rng()); to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty())); to.shuffle(&mut rand::thread_rng()); move_users( ctx, guild, to.into_iter() .zip(&users) .map(|(channel, user)| (channel.id, user.user.id)) .collect(), ) .await?; ctx.reply(format!("Moved {} users", users.len())).await?; Ok(()) } #[tracing::instrument] async fn collect( ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.defer_ephemeral().await?; let users = channel_children(&ctx, from) .await? .into_iter() .flat_map(|x| x.members(ctx)) .flatten() .collect::>(); move_users( ctx, guild, users.iter().map(|user| (to, user.user.id)).collect(), ) .await?; futures::future::try_join_all(users.iter().map(|user| guild.move_member(&ctx, user, to))) .await?; ctx.reply(format!("Moved {} users", users.len())).await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] #[tracing::instrument] 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().get(guild).await?; spread(ctx, town_square, cottages).await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] #[tracing::instrument] 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().get(guild).await?; collect(ctx, cottages, town_square).await?; Ok(()) } #[poise::command(slash_command, ephemeral)] #[tracing::instrument] async fn join( ctx: Context<'_>, to: serenity::UserId, user: Option, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { cottages, .. } = ctx.data().get(guild).await?; let channel = channel_children(&ctx, cottages) .await? .into_iter() .find(|chan| { chan.members(ctx) .is_ok_and(|members| members.iter().any(|member| member.user.id == to)) }); match channel { Some(channel) => { guild .move_member(ctx, user.unwrap_or(ctx.author().id), channel) .await?; ctx.reply("Joined player").await?; } None => { ctx.reply("User is not in cottages").await?; } } Ok(()) } #[poise::command(slash_command, ephemeral)] #[tracing::instrument] async fn st(ctx: Context<'_>) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { st, .. } = ctx.data().get(guild).await?; let sender = ctx.author().id; futures::future::try_join_all(guild.members(&ctx, None, None).await?.into_iter().map( |member| async move { match (sender == member.user.id, member.roles.contains(&st)) { (true, true) | (false, false) => Ok(()), (true, false) => member.add_role(&ctx, st).await, (false, true) => member.remove_role(&ctx, st).await, } }, )) .await?; ctx.reply("You are now the ST").await?; Ok(()) } #[poise::command(slash_command, ephemeral)] #[tracing::instrument] async fn spectate(ctx: Context<'_>, player: Option) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { st, .. } = ctx.data().get(guild).await?; let member = guild .member(&ctx, player.unwrap_or(ctx.author().id)) .await?; let msg = if member.roles.contains(&st) { member.remove_role(&ctx, st).await?; "no longer a spectator" } else { member.add_role(&ctx, st).await?; "now a spectator" }; ctx.reply(if player.is_some() { format!("{} is {}", member.mention(), msg) } else { format!("You are {}", msg) }) .await?; Ok(()) } #[poise::command(slash_command, ephemeral)] #[tracing::instrument] async fn currently_playing(ctx: Context<'_>, player: Option) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { currently_playing, .. } = ctx.data().get(guild).await?; let member = guild .member(&ctx, player.unwrap_or(ctx.author().id)) .await?; let msg = if member.roles.contains(¤tly_playing) { member.remove_role(&ctx, currently_playing).await?; "no longer" } else { member.add_role(&ctx, currently_playing).await?; "now" }; ctx.reply(if player.is_some() { format!("{} is {} {}", member.mention(), msg, currently_playing.mention()) } else { format!("You are {} {}", msg, currently_playing.mention()) }) .await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] #[tracing::instrument] async fn configure( ctx: Context<'_>, #[channel_types("Voice")] town_square: serenity::ChannelId, #[channel_types("Category")] cottages: serenity::ChannelId, st: serenity::RoleId, currently_playing: serenity::RoleId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.data() .insert(GuildData { guild, town_square, cottages, st, currently_playing, }) .await?; ctx.reply("Configured successfully!").await?; Ok(()) } #[tracing::instrument(skip(ctx))] 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")] #[tracing::instrument] async fn main() -> Result<()> { color_eyre::install()?; tracing::subscriber::set_global_default( tracing_subscriber::registry() .with( fmt::layer() .event_format(fmt::format().with_ansi(true).pretty()) .with_span_events(FmtSpan::ACTIVE) .with_writer(io::stderr), ) .with(EnvFilter::from_default_env()) .with(tracing_error::ErrorLayer::default()), )?; 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 dbfile = std::env::var_os("DB_FILE").ok_or_eyre("A database file is required")?; let intents = serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::GUILD_MEMBERS; let framework = poise::Framework::builder() .options(poise::FrameworkOptions { commands: vec![ dusk(), dawn(), st(), spectate(), currently_playing(), join(), 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?; } Data::open(dbfile) }) }) .build(); serenity::ClientBuilder::new(token, intents) .framework(framework) .await? .start() .await?; Ok(()) }