From 84b742ee6ed9b792725955153af39b687146cb44 Mon Sep 17 00:00:00 2001 From: Joscha Date: Tue, 24 Jan 2023 12:15:22 +0100 Subject: [PATCH] Add bot::command and bot::commands --- Cargo.toml | 1 + src/bot.rs | 2 + src/bot/command.rs | 76 ++++++++++++++ src/bot/commands.rs | 242 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 321 insertions(+) create mode 100644 src/bot/command.rs create mode 100644 src/bot/commands.rs diff --git a/Cargo.toml b/Cargo.toml index 423712a..8aaca88 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ edition = "2021" bot = ["dep:cookie"] [dependencies] +async-trait = "0.1.63" cookie = { version = "0.16.2", optional = true } futures-util = { version = "0.3.25", default-features = false, features = ["sink"] } log = "0.4.17" diff --git a/src/bot.rs b/src/bot.rs index ea7f5f0..03ed2ef 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,4 +1,6 @@ //! Building blocks for bots. +pub mod command; +pub mod commands; pub mod instance; pub mod instances; diff --git a/src/bot/command.rs b/src/bot/command.rs new file mode 100644 index 0000000..5b9f749 --- /dev/null +++ b/src/bot/command.rs @@ -0,0 +1,76 @@ +use async_trait::async_trait; + +use crate::api::{self, Message, MessageId}; +use crate::conn::{self, ConnTx, Joined}; + +use super::instance::InstanceConfig; + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum Kind { + /// Global commands always respond. They override any specific or general + /// commands of the same name. + Global, + /// General commands only respond if no nick is specified. + General, + /// Specific commands only respond if the bot's current nick is specified. + Specific, +} + +impl Kind { + pub fn global_and_general_usage(cmd_name: &str) -> String { + format!("!{cmd_name}") + } + + pub fn specific_nick(nick: &str) -> String { + nick.replace(char::is_whitespace, "") + } + + pub fn specific_usage(cmd_name: &str, nick: &str) -> String { + format!("!{cmd_name} @{}", Self::specific_nick(nick)) + } + + pub fn usage(self, cmd_name: &str, nick: &str) -> String { + match self { + Self::Global | Self::General => Self::global_and_general_usage(cmd_name), + Self::Specific => Self::specific_usage(cmd_name, nick), + } + } +} + +pub struct Context { + pub kind: Kind, + pub config: InstanceConfig, + pub conn_tx: ConnTx, + pub joined: Joined, +} + +impl Context { + pub async fn send(&self, content: S) -> Result { + let cmd = api::Send { + content: content.to_string(), + parent: None, + }; + Ok(self.conn_tx.send(cmd).await?.0) + } + + pub async fn reply( + &self, + parent: MessageId, + content: S, + ) -> Result { + let cmd = api::Send { + content: content.to_string(), + parent: Some(parent), + }; + Ok(self.conn_tx.send(cmd).await?.0) + } +} + +#[async_trait] +pub trait Command { + fn description(&self) -> Option { + None + } + + async fn execute(&self, arg: &str, msg: &Message, ctx: &Context, bot: &mut B); +} diff --git a/src/bot/commands.rs b/src/bot/commands.rs new file mode 100644 index 0000000..cef0a3a --- /dev/null +++ b/src/bot/commands.rs @@ -0,0 +1,242 @@ +use std::collections::HashMap; + +use crate::api::packet::ParsedPacket; +use crate::api::{Data, SendEvent}; +use crate::conn; + +use super::command::{Command, Context, Kind}; +use super::instance::{InstanceConfig, Snapshot}; + +fn normalize_specific_nick(nick: &str) -> String { + Kind::specific_nick(nick).to_lowercase() +} + +/// Parse leading whitespace followed by an `!`-initiated command. +/// +/// Returns the command name and the remaining text with one leading whitespace +/// removed. The remaining text may be the empty string. +fn parse_command(text: &str) -> Option<(&str, &str)> { + let text = text.trim_start(); + let text = text.strip_prefix('!')?; + let (name, rest) = text.split_once(char::is_whitespace).unwrap_or((text, "")); + if name.is_empty() { + return None; + } + Some((name, rest)) +} + +/// Parse leading whitespace followed by an `@`-initiated nick. +/// +/// Returns the nick and the remaining text with one leading whitespace removed. +/// The remaining text may be the empty string. +fn parse_specific(text: &str) -> Option<(&str, &str)> { + let text = text.trim_start(); + let text = text.strip_prefix('@')?; + let (name, rest) = text.split_once(char::is_whitespace).unwrap_or((text, "")); + if name.is_empty() { + return None; + } + Some((name, rest)) +} + +pub struct CommandInfo { + pub kind: Kind, + pub name: String, + pub description: Option, + pub visible: bool, +} + +struct CommandWrapper { + command: Box>, + visible: bool, +} + +pub struct Commands { + global: HashMap>, + general: HashMap>, + specific: HashMap>, +} + +impl Commands { + /// Global commands always respond. They override any specific or general + /// commands of the same name. + /// + /// Use this if your bot "owns" the command and no other bot uses it. + pub fn global(mut self, name: S, command: C, visible: bool) -> Self + where + S: ToString, + C: Command + 'static, + { + let command = Box::new(command); + let info = CommandWrapper { command, visible }; + self.global.insert(name.to_string(), info); + self + } + + /// General commands only respond if no nick is specified. + /// + /// Use this if your or any other bot has a specific command of the same + /// name. + pub fn general(mut self, name: S, command: C, visible: bool) -> Self + where + S: ToString, + C: Command + 'static, + { + let command = Box::new(command); + let info = CommandWrapper { command, visible }; + self.general.insert(name.to_string(), info); + self + } + + /// Specific commands only respond if the bot's current nick is specified. + pub fn specific(mut self, name: S, command: C, visible: bool) -> Self + where + S: ToString, + C: Command + 'static, + { + let command = Box::new(command); + let info = CommandWrapper { command, visible }; + self.specific.insert(name.to_string(), info); + self + } + + pub fn descriptions(&self) -> Vec { + let mut keys = (self.global.keys()) + .chain(self.general.keys()) + .chain(self.specific.keys()) + .collect::>(); + keys.sort_unstable(); + keys.dedup(); + + let mut result = vec![]; + for name in keys { + if let Some(wrapper) = self.global.get(name) { + result.push(CommandInfo { + name: name.clone(), + kind: Kind::Global, + visible: wrapper.visible, + description: wrapper.command.description(), + }); + continue; // Shadows general and specific commands + } + + if let Some(wrapper) = self.general.get(name) { + result.push(CommandInfo { + name: name.clone(), + kind: Kind::General, + visible: wrapper.visible, + description: wrapper.command.description(), + }); + } + + if let Some(wrapper) = self.specific.get(name) { + result.push(CommandInfo { + name: name.clone(), + kind: Kind::Specific, + visible: wrapper.visible, + description: wrapper.command.description(), + }); + } + } + + result + } + + /// Returns `true` if a command was found and executed, `false` otherwise. + pub async fn handle_packet( + &self, + config: &InstanceConfig, + packet: &ParsedPacket, + snapshot: &Snapshot, + bot: &mut B, + ) -> bool { + let msg = match &packet.content { + Ok(Data::SendEvent(SendEvent(msg))) => msg, + _ => return false, + }; + + let joined = match &snapshot.state { + conn::State::Joining(_) => return false, + conn::State::Joined(joined) => joined.clone(), + }; + + let (cmd_name, rest) = match parse_command(&msg.content) { + Some(parsed) => parsed, + None => return false, + }; + + let mut ctx = Context { + kind: Kind::Global, + config: config.clone(), + conn_tx: snapshot.conn_tx.clone(), + joined, + }; + + if let Some(wrapper) = self.global.get(cmd_name) { + ctx.kind = Kind::Global; + wrapper.command.execute(rest, msg, &ctx, bot).await; + return true; + } + + if let Some((cmd_nick, rest)) = parse_specific(rest) { + if let Some(wrapper) = self.specific.get(cmd_name) { + let nick_norm = normalize_specific_nick(&ctx.joined.session.name); + let cmd_nick_norm = normalize_specific_nick(cmd_nick); + if nick_norm == cmd_nick_norm { + ctx.kind = Kind::Specific; + wrapper.command.execute(rest, msg, &ctx, bot).await; + return true; + } + } + + // The command looks like a specific command. If we treated it like + // a general command just because the nick doesn't match, we would + // interpret other bots' specific commands as general commands. + // + // To call a specific command with a mention as its first positional + // argument, -- can be used. + return false; + } + + if let Some(wrapper) = self.general.get(cmd_name) { + ctx.kind = Kind::General; + wrapper.command.execute(rest, msg, &ctx, bot).await; + return true; + } + + false + } +} + +#[cfg(test)] +mod test { + use super::{parse_command, parse_specific}; + + #[test] + fn test_parse_command() { + assert_eq!(parse_command("!foo"), Some(("foo", ""))); + assert_eq!(parse_command(" !foo"), Some(("foo", ""))); + assert_eq!(parse_command("!foo "), Some(("foo", " "))); + assert_eq!(parse_command(" !foo "), Some(("foo", " "))); + assert_eq!(parse_command("!foo @bar"), Some(("foo", "@bar"))); + assert_eq!(parse_command("!foo @bar"), Some(("foo", " @bar"))); + assert_eq!(parse_command("!foo @bar "), Some(("foo", "@bar "))); + assert_eq!(parse_command("! foo @bar"), None); + assert_eq!(parse_command("!"), None); + assert_eq!(parse_command("?foo"), None); + } + + #[test] + fn test_parse_specific() { + assert_eq!(parse_specific("@foo"), Some(("foo", ""))); + assert_eq!(parse_specific(" @foo"), Some(("foo", ""))); + assert_eq!(parse_specific("@foo "), Some(("foo", " "))); + assert_eq!(parse_specific(" @foo "), Some(("foo", " "))); + assert_eq!(parse_specific("@foo !bar"), Some(("foo", "!bar"))); + assert_eq!(parse_specific("@foo !bar"), Some(("foo", " !bar"))); + assert_eq!(parse_specific("@foo !bar "), Some(("foo", "!bar "))); + assert_eq!(parse_specific("@ foo !bar"), None); + assert_eq!(parse_specific("@"), None); + assert_eq!(parse_specific("?foo"), None); + } +}