use futures::StreamExt; use irc::client::prelude::Config as IrcConfig; use irc::client::prelude::*; use serde::Deserialize; use std::collections::HashMap; #[derive(Deserialize)] struct Config { quit_message: String, timeout_limit: Option, irc: IrcConfig, } macro_rules! unwrap_or_continue { ($o:expr) => { match $o { Some(v) => v, None => continue, } }; } #[derive(Clone, Copy, Debug)] enum Status { TimeoutCount(u8), Banned, } impl std::default::Default for Status { fn default() -> Self { Self::TimeoutCount(0) } } #[tokio::main] async fn main() -> anyhow::Result<()> { let filename = std::env::var("SLEEPERAGENT_CONFIG").unwrap_or("sleeperagent.toml".into()); let contents = tokio::fs::read_to_string(filename).await?; let conf: Config = toml::from_str(&contents)?; let mut client = Client::from_config(conf.irc).await?; client.identify()?; let mut stream = client.stream()?; let mut channel_users: HashMap> = HashMap::new(); while let Some(message) = stream.next().await.transpose()? { if let Command::Response(response, args) = message.command { if response == Response::RPL_NAMREPLY { let channel = args[2].to_string(); let users = args[3] .split(' ') .map(|s| (s.to_owned(), Status::TimeoutCount(0))) .collect(); channel_users.insert(channel, users); } continue; } let nick = unwrap_or_continue!(message.source_nickname()); let nick = nick.strip_prefix(['~', '&', '@', '%', '+']).unwrap_or(nick); match message.command { Command::JOIN(ref channel, _, _) => { let users = unwrap_or_continue!(channel_users.get_mut(channel)); for (user, status) in users.iter_mut() { if nick.starts_with(user) { // They're joining back for real, unban them *status = match status { Status::Banned => Status::TimeoutCount(0), _ => *status, }; let mode = Mode::Minus(ChannelMode::Ban, Some(format!("{}!*@*", user))); client.send_mode(channel, &[mode])?; } } if !users.contains_key(nick) { users.insert(nick.to_string(), Status::TimeoutCount(0)); } } Command::PRIVMSG(ref channel, ref message) => { if let Some(mut arg) = message.strip_prefix("!dbg") { arg = arg.trim(); if !arg.is_empty() { if let Some(h) = channel_users.get(channel) { client.send_privmsg( nick, format!( "{}: {}", arg, h.get(arg) .map(|v| format!("{:?}", v)) .unwrap_or_else(|| "No such nick".into()) ) .replace("\n", "\r\n"), )?; } else { client.send_privmsg( nick, "!dbg with a nickname can only be used in a channel!", )?; } } else { client.send_privmsg( nick, format!("{:#?}", channel_users).replace("\n", "\r\n"), )?; } } else if let Some(mut user) = message.strip_prefix("!unban ") { user = user.trim(); if !user.is_empty() { if let Some(users) = channel_users.get_mut(channel) { if let Some(status) = users.get_mut(user) { *status = Status::TimeoutCount(0); let mode = Mode::Minus(ChannelMode::Ban, Some(format!("{}!*@*", user))); client.send_mode(channel, &[mode])?; client.send_privmsg(channel, format!("Unbanned {}", user))?; } } } } let users = unwrap_or_continue!(channel_users.get_mut(channel)); let user = unwrap_or_continue!(users.get_mut(nick)); *user = match user { Status::TimeoutCount(count) if *count > conf.timeout_limit.unwrap_or(1) => { Status::TimeoutCount(0) } _ => *user, }; } Command::NICK(ref new_nick) => { for (_, users) in &mut channel_users { let status = users.remove(nick).unwrap_or_default(); users.insert(new_nick.clone(), status); } } Command::PART(ref channel, Some(ref message)) => { if message == &conf.quit_message { let users = unwrap_or_continue!(channel_users.get_mut(channel)); let user = unwrap_or_continue!(users.get_mut(nick)); *user = match user { Status::TimeoutCount(count) => Status::TimeoutCount(*count + 1), _ => *user, }; if let Status::TimeoutCount(count) = user { if *count > conf.timeout_limit.unwrap_or(1) { let mode = Mode::Plus(ChannelMode::Ban, Some(format!("{}!*@*", nick))); client.send_mode(channel, &[mode])?; *user = Status::Banned; } } } } Command::QUIT(Some(ref message)) => { for (channel, users) in &mut channel_users { let user = unwrap_or_continue!(users.get_mut(nick)); if message == &conf.quit_message { *user = match user { Status::TimeoutCount(count) => Status::TimeoutCount(*count + 1), _ => *user, }; if let Status::TimeoutCount(count) = user { if *count > conf.timeout_limit.unwrap_or(1) { let mode = Mode::Plus(ChannelMode::Ban, Some(format!("{}!*@*", nick))); client.send_mode(channel, &[mode])?; *user = Status::Banned; } } } } } _ => {} } } Ok(()) }