sleeperagent/src/main.rs

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<u8>,
    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<String, HashMap<String, Status>> = 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(())
}