Added commands + persistence

+ `say`: used by the host to speak in the game chat + `end`: ends the game and cleans up + `broadcast`: Send an important game message that that is pinned in each channel + Persistent allows games to continue across bot restarts + clippy + fmt
2022-03-06 10:03:49 -07:00 · 2022-03-06 10:03:49 -07:00 · e72d603149
parent fd878c3721
commit e72d603149
7 changed files with 278 additions and 83 deletions
--- a/.gitignore
+++ b/.gitignore
@ -1,3 +1,4 @@
 /target
 config.toml
+wOxlf_data.toml
 .idea
--- a/Cargo.lock
+++ b/Cargo.lock
@ -1729,6 +1729,7 @@ dependencies = [
 "serenity",
 "structopt",
 "tokio",
+ "toml",
 ]

 [[package]]
--- a/Cargo.toml
+++ b/Cargo.toml
@ -11,6 +11,7 @@ structopt = "0.3.26"
 chrono = "0.4.19"
 serde = "1.0.136"
 rand = "0.8.5"
+toml = "0.5.8"

 [dependencies.serenity]
 version = "0.10.10"
--- a/src/commands.rs
+++ b/src/commands.rs
@ -1,18 +1,19 @@
-use serenity::prelude::Context;
-use serenity::model::prelude::{Message, UserId};
-use serenity::framework::standard::{CommandResult, Args};
-use serenity::framework::StandardFramework;
-use serenity::framework::standard::macros::{command, group};
-use crate::data::{GlobalData, BotConfig, PlayerData};
+use crate::data::{BotConfig, GlobalData, MessageSource, PlayerData};
+use crate::helper::{clear_game_state, save_game_state, send_msg_to_player_channels};
 use rand::Rng;
-use serenity::model::guild::{Member, Guild};
-use serenity::model::id::ChannelId;
-use serenity::model::Permissions;
+use serenity::framework::standard::macros::{command, group};
+use serenity::framework::standard::{Args, CommandResult};
+use serenity::framework::StandardFramework;
 use serenity::model::channel::{PermissionOverwrite, PermissionOverwriteType};
+use serenity::model::guild::{Guild, Member};
+use serenity::model::id::ChannelId;
+use serenity::model::prelude::{Message, UserId};
+use serenity::model::Permissions;
+use serenity::prelude::Context;
 use serenity::utils::MessageBuilder;

 #[group]
-#[commands(start)]
+#[commands(start, say, end, broadcast)]
 struct Host;

 fn generate_codename(config: &BotConfig) -> String {
@ -24,21 +25,28 @@ fn generate_codename(config: &BotConfig) -> String {
    format!("{} {}", adj, occupation)
 }

-async fn add_user_to_game(ctx: &Context, guild: &Guild, global_data: &mut GlobalData, discord_user: &Member) {
+async fn add_user_to_game(
+    ctx: &Context,
+    guild: &Guild,
+    global_data: &mut GlobalData,
+    discord_user: &Member,
+) -> CommandResult {
    let mut codename = generate_codename(&global_data.cfg);

-    while global_data.game_state.codename_exists(&mut codename) {
+    while global_data.game_state.codename_exists(&codename) {
        codename = generate_codename(&global_data.cfg);
    }

+    let channel = guild
+        .create_channel(&ctx.http, |c| {
+            c.category(&ChannelId::from(global_data.cfg.category))
+                .name(format!("{}'s Channel", discord_user.display_name()))
+        })
+        .await?;

-    let channel = guild.create_channel(&ctx.http, |c| {
-        c
-            .category(&ChannelId::from(global_data.cfg.category))
-            .name(format!("{}'s Channel", discord_user.display_name()))
-    }).await.unwrap();
-
-    let allow = Permissions::SEND_MESSAGES | Permissions::READ_MESSAGE_HISTORY | Permissions::READ_MESSAGE_HISTORY;
+    let allow = Permissions::SEND_MESSAGES
+        | Permissions::READ_MESSAGE_HISTORY
+        | Permissions::READ_MESSAGE_HISTORY;

    let overwrite = PermissionOverwrite {
        allow,
@ -46,7 +54,7 @@ async fn add_user_to_game(ctx: &Context, guild: &Guild, global_data: &mut Global
        kind: PermissionOverwriteType::Member(discord_user.user.id),
    };

-    channel.create_permission(&ctx.http, &overwrite).await.unwrap();
+    channel.create_permission(&ctx.http, &overwrite).await?;

    let msg = channel.send_message(&ctx.http, |m| {
        m.content(MessageBuilder::new()
@ -59,18 +67,20 @@ async fn add_user_to_game(ctx: &Context, guild: &Guild, global_data: &mut Global
            .push("SUBJECT CODENAME: ")
            .push_line(&codename)
        )
-    }).await.unwrap();
+    }).await?;

-    channel.pin(&ctx.http, msg.id).await.unwrap();
+    channel.pin(&ctx.http, msg.id).await?;

    let player_data = PlayerData {
+        channel: channel.id.0,
        discord_id: discord_user.user.id.0,
-        codename: codename
+        codename,
    };

-    global_data.game_state.player_channels.insert(channel.id.0, player_data);
-}
+    global_data.game_state.player_data.push(player_data);

+    Ok(())
+}

 #[command]
 #[only_in(guilds)]
@ -78,33 +88,99 @@ async fn add_user_to_game(ctx: &Context, guild: &Guild, global_data: &mut Global
 async fn start(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
    msg.channel_id.say(&ctx.http, "Starting game").await?;

-    let mut data =  ctx.data.write().await;
+    let mut data = ctx.data.write().await;
    let global_data = data.get_mut::<GlobalData>().unwrap();
    let guild = msg.guild(&ctx.cache).await.unwrap();

    let mut global_data = global_data.lock().await;

-    global_data.game_state.clear();
+    clear_game_state(&mut global_data).unwrap();

-    for player in args.iter::<u64>() {
-        if let Ok(player) = player {
-            if let Some(discord_user) = guild.members.get(&UserId::from(player)) {
-                add_user_to_game(ctx, &guild, &mut global_data, &discord_user).await;
-            }
-            else {
-                msg.reply(&ctx.http,format!("User {} is invalid or not in this server!", player)).await.unwrap();
-                break;
-            }
+    for player in args.iter::<u64>().flatten() {
+        if let Some(discord_user) = guild.members.get(&UserId::from(player)) {
+            add_user_to_game(ctx, &guild, &mut global_data, discord_user).await?;
+        } else {
+            msg.reply(
+                &ctx.http,
+                format!("User {} is invalid or not in this server!", player),
+            )
+            .await?;
+
+            break;
        }
    }

+    save_game_state(&global_data).unwrap();
+
+    Ok(())
+}
+
+#[command]
+#[only_in(guilds)]
+#[allowed_roles("wolfx host")]
+async fn end(ctx: &Context, msg: &Message, mut _args: Args) -> CommandResult {
+    let mut data = ctx.data.write().await;
+    let global_data = data.get_mut::<GlobalData>().unwrap();
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+
+    let mut global_data = global_data.lock().await;
+
+    for player_data in &global_data.game_state.player_data {
+        let channel = guild
+            .channels
+            .get(&ChannelId::from(player_data.channel))
+            .unwrap();
+
+        channel.delete(&ctx.http).await?;
+    }
+
+    clear_game_state(&mut global_data).unwrap();
+
+    Ok(())
+}
+
+#[command]
+#[only_in(guilds)]
+#[allowed_roles("wolfx host")]
+async fn say(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
+    let data = ctx.data.read().await;
+    let global_data = data.get::<GlobalData>().unwrap();
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+
+    let global_data = global_data.lock().await;
+
+    let msg = format!("**wOxlf **> {}", args.rest());
+
+    send_msg_to_player_channels(ctx, &guild, &global_data, MessageSource::Host, &msg, false).await;
+
+    Ok(())
+}
+
+#[command]
+#[only_in(guilds)]
+#[allowed_roles("wolfx host")]
+async fn broadcast(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
+    let data = ctx.data.read().await;
+    let global_data = data.get::<GlobalData>().unwrap();
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+
+    let global_data = global_data.lock().await;
+
+    let msg = MessageBuilder::new()
+        .push_bold_line("\\*\\*IMPORTANT wOxlf SYSTEM MESSAGE\\*\\*")
+        .push_line("")
+        .push_line(args.rest())
+        .push_line("")
+        .push_bold_line("\\*\\*END OF SYSTEM MESSAGE\\*\\*")
+        .build();
+
+    send_msg_to_player_channels(ctx, &guild, &global_data, MessageSource::Host, &msg, true).await;
+
    Ok(())
 }

 pub fn command_framework() -> StandardFramework {
    StandardFramework::new()
-        .configure(|c| {
-            c.prefix("!")
-        })
+        .configure(|c| c.prefix("!"))
        .group(&HOST_GROUP)
 }
--- a/src/data.rs
+++ b/src/data.rs
@ -1,17 +1,16 @@
 use config::{Config, File};
 use serde::{Deserialize, Serialize};
-use std::path::Path;
-use std::collections::HashMap;
 use serenity::prelude::TypeMapKey;
-use std::sync::Arc;
+use std::path::Path;
 use std::path::PathBuf;
+use std::sync::Arc;
 use structopt::StructOpt;
 use tokio::sync::Mutex;

 #[derive(Debug, StructOpt)]
 #[structopt(name = "WOXlf", about = "WOXlf discord bot")]
 pub struct Args {
-    pub cfg_path: PathBuf
+    pub cfg_path: PathBuf,
 }

 #[derive(Debug, Deserialize, Serialize, Clone)]
@ -20,8 +19,9 @@ pub struct BotConfig {
    pub app_id: u64,
    pub host_channel: u64,
    pub category: u64,
+    pub game_state_dir: PathBuf,
    pub occupation: Vec<String>,
-    pub adjective: Vec<String>
+    pub adjective: Vec<String>,
 }

 impl BotConfig {
@ -32,12 +32,16 @@ impl BotConfig {

        cfg.try_deserialize()
    }
+
+    pub fn get_game_state_path(&self) -> PathBuf {
+        self.game_state_dir.join("wOxlf_data.toml")
+    }
 }

 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub enum Phase {
    Day,
-    Night
+    Night,
 }

 impl Default for Phase {
@ -48,41 +52,45 @@ impl Default for Phase {

 #[derive(Debug, Deserialize, Serialize, Clone, Default, Hash)]
 pub struct PlayerData {
+    pub channel: u64,
    pub discord_id: u64,
    pub codename: String,
 }

 #[derive(Debug, Deserialize, Serialize, Clone, Default)]
 pub struct GameState {
-    pub player_channels: HashMap<u64, PlayerData>,
    pub current_phase: Phase,
+    pub player_data: Vec<PlayerData>,
 }

 impl GameState {
    pub fn codename_exists(&self, codename: &str) -> bool {
-        self.player_channels.iter().any(|(_, data)| {
-            data.codename.to_lowercase() == codename
-        })
+        self.player_data
+            .iter()
+            .any(|data| data.codename.to_lowercase() == codename)
    }

    pub fn clear(&mut self) {
-        self.player_channels.clear();
+        self.player_data.clear();
        self.current_phase = Phase::Night;
    }
+
+    pub fn get_player_from_channel(&self, channel_id: u64) -> Option<&PlayerData> {
+        self.player_data.iter().find(|p| p.channel == channel_id)
+    }
 }

 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct GlobalData {
    pub cfg: BotConfig,
-    pub game_state: GameState
+    pub game_state: GameState,
 }

-
 impl GlobalData {
    pub fn new(cfg: BotConfig) -> Self {
        Self {
            cfg,
-            game_state: GameState::default()
+            game_state: GameState::default(),
        }
    }
 }
@ -91,3 +99,9 @@ impl TypeMapKey for GlobalData {
    type Value = Arc<Mutex<GlobalData>>;
 }

+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub enum MessageSource {
+    Player(u64),
+    Host,
+    Automated,
+}
--- a/src/helper.rs
+++ b/src/helper.rs
@ -0,0 +1,99 @@
+use crate::data::{GlobalData, MessageSource};
+use serenity::model::id::UserId;
+use serenity::model::prelude::ChannelId;
+use serenity::model::prelude::Guild;
+use serenity::prelude::Context;
+use std::fs::File;
+use std::io::prelude::*;
+
+pub async fn send_msg_to_player_channels(
+    ctx: &Context,
+    guild: &Guild,
+    global_data: &GlobalData,
+    msg_source: MessageSource,
+    msg: &str,
+    pin: bool,
+) {
+    for player_data in &global_data.game_state.player_data {
+        if let MessageSource::Player(channel_id) = msg_source {
+            if channel_id == player_data.channel {
+                continue;
+            }
+        }
+
+        let channel = guild
+            .channels
+            .get(&ChannelId::from(player_data.channel))
+            .unwrap();
+
+        let msg = channel
+            .send_message(&ctx.http, |m| m.content(&msg))
+            .await
+            .unwrap();
+
+        if pin {
+            // pin system messages
+            msg.pin(&ctx.http).await.unwrap();
+        }
+    }
+
+    let host_channel = guild
+        .channels
+        .get(&ChannelId::from(global_data.cfg.host_channel))
+        .unwrap();
+
+    let source = match msg_source {
+        MessageSource::Player(channel_id) => {
+            let discord_id = global_data
+                .game_state
+                .get_player_from_channel(channel_id)
+                .unwrap()
+                .discord_id;
+            let name = guild
+                .members
+                .get(&UserId::from(discord_id))
+                .unwrap()
+                .display_name();
+
+            name.to_string()
+        }
+        MessageSource::Host => "Host".to_string(),
+        MessageSource::Automated => "Automated".to_string(),
+    };
+
+    host_channel
+        .send_message(&ctx.http, |m| m.content(format!("({}): {}", source, msg)))
+        .await
+        .unwrap();
+}
+
+pub fn save_game_state(global_data: &GlobalData) -> std::io::Result<()> {
+    let s = toml::to_string_pretty(&global_data.game_state).unwrap();
+
+    let mut file = File::create(global_data.cfg.get_game_state_path())?;
+
+    file.write_all(s.as_bytes())
+}
+
+pub fn get_game_state(global_data: &mut GlobalData) -> std::io::Result<()> {
+    let mut file = File::open(global_data.cfg.get_game_state_path())?;
+
+    let mut data = String::new();
+
+    file.read_to_string(&mut data)?;
+
+    global_data.game_state = toml::from_str(&data).unwrap();
+
+    Ok(())
+}
+
+pub fn clear_game_state(global_data: &mut GlobalData) -> std::io::Result<()> {
+    global_data.game_state.clear();
+
+    let state_path = global_data.cfg.get_game_state_path();
+    if state_path.exists() {
+        std::fs::remove_file(global_data.cfg.get_game_state_path())?;
+    }
+
+    Ok(())
+}
--- a/src/main.rs
+++ b/src/main.rs
@ -1,15 +1,16 @@
-mod data;
 mod commands;
+mod data;
+mod helper;

-use serenity::prelude::*;
-use serenity::model::prelude::{Message, Ready};
-use std::sync::Arc;
-use serenity::client::bridge::gateway::GatewayIntents;
-use crate::data::{GlobalData, Args, BotConfig};
-use serenity::async_trait;
 use crate::commands::command_framework;
+use crate::data::{Args, BotConfig, GlobalData, MessageSource};
+use crate::helper::{get_game_state, send_msg_to_player_channels};
+use serenity::async_trait;
+use serenity::client::bridge::gateway::GatewayIntents;
+use serenity::model::prelude::{Message, Ready};
+use serenity::prelude::*;
+use std::sync::Arc;
 use structopt::StructOpt;
-use serenity::model::id::ChannelId;

 struct Handler {}

@ -20,36 +21,32 @@ impl EventHandler for Handler {
            return;
        }

+        if msg.content.starts_with('!') {
+            return;
+        }
+
        let data = ctx.data.read().await;

        let game_data = data.get::<GlobalData>().unwrap();

        let game_data = game_data.lock().await;

-
-        if let Some(player_data) = game_data.game_state.player_channels.get(&msg.channel_id.0) {
+        if let Some(player_data) = game_data
+            .game_state
+            .get_player_from_channel(msg.channel_id.0)
+        {
            let guild = msg.guild(&ctx.cache).await.unwrap();
            let user_msg = format!("{} > {}", player_data.codename, msg.content);

-            for (channel, _) in &game_data.game_state.player_channels {
-                if *channel == msg.channel_id.0 {
-                    continue;
-                }
-
-                let channel = guild.channels.get(&ChannelId::from(*channel)).unwrap();
-
-                channel.send_message(&ctx.http, |m| {
-                    m
-                        .content(&user_msg)
-                }).await.unwrap();
-            }
-
-            let host_channel = guild.channels.get(&ChannelId::from(game_data.cfg.host_channel)).unwrap();
-
-            host_channel.send_message(&ctx.http, |m| {
-                m
-                    .content(&user_msg)
-            }).await.unwrap();
+            send_msg_to_player_channels(
+                &ctx,
+                &guild,
+                &game_data,
+                MessageSource::Player(msg.channel_id.0),
+                &user_msg,
+                false,
+            )
+            .await;
        }
    }

@ -64,11 +61,17 @@ async fn main() {

    let bot_cfg: BotConfig = BotConfig::new(&args.cfg_path).expect("Unable to parse cfg");

+    let mut global_data = GlobalData::new(bot_cfg.clone());
+
+    if get_game_state(&mut global_data).is_ok() {
+        println!("Resuming game...")
+    }
+
    let mut client = Client::builder(&bot_cfg.token)
        .event_handler(Handler {})
        .framework(command_framework())
        .intents(GatewayIntents::all())
-        .type_map_insert::<GlobalData>(Arc::new(Mutex::new(GlobalData::new(bot_cfg))))
+        .type_map_insert::<GlobalData>(Arc::new(Mutex::new(global_data)))
        .await
        .expect("Err creating client");