tron/src/main.rs

127 lines
3.4 KiB
Rust

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<()> {
let _ = kankyo::init();
pretty_env_logger::init();
let cfg: Config = envy::from_env()?;
log::debug!("cfg: {:?}", cfg);
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(())
}