tron/src/main.rs

use anyhow::Result;
use async_trait::async_trait;
use discord_webhook::Body;
use furbooru::{Client, Comment, FirehoseAdaptor, Image};
use regex::Regex;
use serde::Deserialize;

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct Config {
    discord_webhook_url: String,
    furbooru_api_key: String,
    bot_owner_furbooru_account: String,
    regexes: std::path::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 Rules(Vec<CompiledRule>, Config);

#[async_trait]
impl FirehoseAdaptor for Rules {
    async fn image_created(&self, img: Image) -> Result<()> {
        let mut buf: String = String::new();
        let mut found = false;

        for rule in &self.0 {
            if rule.regex.is_match(&img.description.to_lowercase()) {
                log::debug!("{:?} matches {}", img.description, rule.raw);
                found = true;
                buf.push_str(&format!("\n- match on rule `{}` ({})", rule.raw, rule.why));
            }
        }

        if found {
            discord_webhook::execute(
                self.1.discord_webhook_url.clone(),
                Body::new(format!("matches found on <{}>:{}", img.view_url, buf)),
            )
            .await?;
            log::info!("the description of {} has naughty words", img.view_url);
        }

        Ok(())
    }

    async fn comment_created(&self, cmt: Comment) -> Result<()> {
        let mut buf: String = String::new();
        let mut found = false;

        for rule in &self.0 {
            if rule.regex.is_match(&cmt.body.to_lowercase()) {
                log::debug!("{:?} matches {}", cmt.body, rule.raw);
                found = true;
                buf.push_str(&format!("\n- match on rule `{}` ({})", rule.raw, rule.why));
            }
        }

        if found {
            discord_webhook::execute(
                self.1.discord_webhook_url.clone(),
                Body::new(format!(
                    "matches found on <https://furbooru.org/{}#comment_{}>:{}",
                    cmt.image_id, cmt.id, buf
                )),
            )
            .await?;
            log::info!(
                "comment https://furbooru.org/{}#comment_{} has naughty words",
                cmt.image_id,
                cmt.id
            );
        }

        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    kankyo::init()?;
    pretty_env_logger::init();
    let cfg: Config = envy::from_env()?;
    let rexes: Vec<Rule> = serde_dhall::from_file(cfg.regexes.clone()).parse()?;
    let mut compiled_rules: Vec<CompiledRule> = Vec::new();

    for rule in rexes {
        log::debug!("{} -> {}", rule.regex, rule.why);
        compiled_rules.push(CompiledRule {
            raw: rule.regex.clone(),
            regex: Regex::new(rule.regex.as_str())?,
            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");
    cli.firehose(Rules(compiled_rules, cfg)).await?;

    Ok(())
}