tron/src/main.rs

use anyhow::Result;
use async_trait::async_trait;
use discord_webhook::Body;
use furbooru::{Client, Comment, FirehoseAdaptor, Forum, Image, Post, Topic};
use regex::Regex;
use serde::Deserialize;
use std::{fmt, path::PathBuf};

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct Config {
    discord_webhook_url: String,
    furbooru_api_key: String,
    bot_owner_furbooru_account: String,
    regexes: PathBuf,
}

impl Config {}

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct Rule {
    regex: String,
    why: String,
}

struct CompiledRule {
    regex: Regex,
    raw: String,
    why: String,
}

pub(crate) fn user_agent(username: String) -> String {
    format!(
        "{}/{} ({}, +{})",
        env!("CARGO_PKG_NAME"),
        env!("CARGO_PKG_VERSION"),
        username,
        env!("CARGO_PKG_REPOSITORY")
    )
}

struct Hit {
    matches: String,
    why: String,
}

impl fmt::Display for Hit {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "- match on rule `{}` ({})", self.matches, self.why)
    }
}

// TODO(Xe): make this an actual struct for the love of god.
struct Rules<'a>(&'a Vec<CompiledRule>, Config, furbooru::Client);

impl<'a> Rules<'a> {
    fn check(&self, text: &String) -> Option<Vec<Hit>> {
        let mut result: Vec<Hit> = vec![];
        let mut found = false;

        for rule in self.0 {
            if rule.regex.is_match(text) {
                log::debug!("{:?} matches {}", text, rule.raw);
                found = true;
                result.push(Hit {
                    matches: rule.raw.clone(),
                    why: rule.why.clone(),
                });
            }
        }

        if found {
            Some(result)
        } else {
            None
        }
    }

    async fn tell(&self, message: String) -> Result<()> {
        discord_webhook::execute(self.1.discord_webhook_url.clone(), Body::new(message)).await?;

        Ok(())
    }
}

#[async_trait]
impl<'a> FirehoseAdaptor for Rules<'a> {
    async fn image_created(&self, img: Image) -> Result<()> {
        let url = format!("https://furbooru.org/{}", img.id);
        log::debug!("got image {}", url);
        if let Some(hits) = self.check(&img.description.to_lowercase()) {
            let mut buf: String = String::new();

            for hit in hits {
                buf.push_str(&format!("\n{}", hit));
            }

            self.tell(format!(
                "matches from **{}** found on <{}>:{}",
                img.uploader.or(Some("An anonymous user".into())).unwrap(),
                url,
                buf
            ))
            .await?;

            log::info!("the description of {} has naughty words", url);
        }

        Ok(())
    }

    async fn image_description_updated(
        &self,
        id: u64,
        added: String,
        _removed: String,
    ) -> Result<()> {
        let url = format!("https://furbooru.org/{}", id);

        if let Some(hits) = self.check(&added.to_lowercase()) {
            let mut buf: String = String::new();

            for hit in hits {
                buf.push_str(&format!("\n{}", hit));
            }

            let img = self.2.image(id.clone()).await?;

            self.tell(format!(
                "description of <{url}> from **{who}** matches:{why}",
                who = img.uploader.or(Some("An anonymous user".into())).unwrap(),
                url = url,
                why = buf
            ))
            .await?;
        }

        Ok(())
    }

    async fn comment_created(&self, cmt: Comment) -> Result<()> {
        let url = format!("https://furbooru.org/{}#comment_{}", cmt.image_id, cmt.id);
        log::debug!("got comment {}", url);
        if let Some(hits) = self.check(&cmt.body.to_lowercase()) {
            let mut buf: String = String::new();

            for hit in hits {
                buf.push_str(&format!("\n{}", hit));
            }
            self.tell(format!(
                "matches from **{}** found on <{}>:{}",
                cmt.author, url, buf
            ))
            .await?;
            log::info!("comment {} has naughty words", url);
        }

        Ok(())
    }

    async fn comment_updated(&self, cmt: Comment) -> Result<()> {
        let url = format!("https://furbooru.org/{}#comment_{}", cmt.image_id, cmt.id);
        log::debug!("got comment update {}", url);
        if let Some(hits) = self.check(&cmt.body.to_lowercase()) {
            let mut buf: String = String::new();

            for hit in hits {
                buf.push_str(&format!("\n{}", hit));
            }
            self.tell(format!(
                "matches from **{}** found on updated comment <{}>:{}",
                cmt.author, url, buf
            ))
            .await?;
            log::info!("comment {} has naughty words", url);
        }

        Ok(())
    }

    async fn post_created(&self, frm: Forum, top: Topic, pst: Post) -> Result<()> {
        let url = format!(
            "https://furbooru.org/forums/{forum}/topics/{topic}?post_id={post}#post_{post}",
            forum = frm.short_name,
            topic = top.slug,
            post = pst.id
        );
        log::debug!("got post {}", url);

        if let Some(hits) = self.check(&pst.body.to_lowercase()) {
            let mut buf: String = String::new();

            for hit in hits {
                buf.push_str(&format!("\n{}", hit));
            }
            self.tell(format!(
                "matches from **{}** found on <{}>:{}",
                pst.author, url, buf
            ))
            .await?;
            log::info!("post {} has naughty words", url);
        }

        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let _ = kankyo::init()
        .or_else(|_| kankyo::load_from_path("/run/keys/tron", true))
        .map_err(|why| println!("startup err: {}", why));
    pretty_env_logger::init();
    let cfg: Config = envy::from_env()?;
    log::debug!("cfg: {:?}", cfg);

    log::debug!("loaded list of words to watch");
    let rexes: Vec<Rule> = serde_dhall::from_file(cfg.regexes.clone()).parse()?;
    let mut compiled_rules: Vec<CompiledRule> = Vec::new();

    #[cfg(not(debug_assertions))]
    {
        discord_webhook::execute(
            cfg.discord_webhook_url.clone(),
            Body::new(format!(
                "I fight for the user! {} version {} active",
                env!("CARGO_PKG_NAME"),
                env!("CARGO_PKG_VERSION")
            )),
        )
        .await?;
    }

    for rule in rexes {
        log::debug!("{} -> {}", rule.regex, rule.why);
        compiled_rules.push(CompiledRule {
            raw: rule.regex.clone(),
            regex: Regex::new(&format!("\\b{}\\b", rule.regex))?,
            why: rule.why,
        })
    }

    let cli = Client::new(
        user_agent(cfg.bot_owner_furbooru_account.clone()),
        cfg.furbooru_api_key.clone(),
    )?;
    log::info!("listening on the firehose");
    let mut tries: u64 = 0;

    loop {
        if let Err(why) = cli
            .firehose(Rules(
                &compiled_rules,
                cfg.clone(),
                Client::new(
                    user_agent(cfg.bot_owner_furbooru_account.clone()),
                    cfg.furbooru_api_key.clone(),
                )?,
            ))
            .await
        {
            log::error!("firehose error: {:#?}", why);

            std::thread::sleep(std::time::Duration::from_millis(tries * 1000));
            if tries == 5 {
                return Err(why);
            }

            tries = tries + 1;
        }
    }
}