From f7c7526081dd01fc8e80ed280e9ef9a3f4434cee Mon Sep 17 00:00:00 2001 From: 0xf8 <0xf8.dev@proton.me> Date: Sun, 16 Apr 2023 12:04:12 -0400 Subject: [PATCH] Add room-specific configs --- src/config.rs | 77 +++++++++++++++++++++++++++++++++++++++++++++++++-- src/judge.rs | 45 ++++++++++++++++++------------ src/main.rs | 45 ++++++++++++++++++++++++------ src/matrix.rs | 10 ++----- 4 files changed, 139 insertions(+), 38 deletions(-) diff --git a/src/config.rs b/src/config.rs index 02ca17b..dfa2e7e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,12 +1,83 @@ use serde_json::Value; -use matrix_sdk::ruma::OwnedRoomId; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; +use matrix_sdk::{ruma::OwnedRoomId, room::Joined}; +use tokio::fs; -#[derive(Debug)] +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq, PartialOrd, Ord)] pub struct RoomConfig { - pub id: OwnedRoomId, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(rename = "muted")] pub no_reply: bool, + #[serde(rename = "add_reply")] + pub reply_to_scam: bool, } +impl Default for RoomConfig { + fn default() -> Self { + Self { id: None, no_reply: false, reply_to_scam: false } + } +} + + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct RoomConfigController { + pub rooms: BTreeSet, +} + +impl RoomConfigController { + pub async fn find(&self, id: OwnedRoomId) -> Option<&RoomConfig> { + let id = Some(id); + for room in self.rooms.to_owned() { + if room.id == id { return Some(self.rooms.get(&room).unwrap()) } + } + + None + } + + pub async fn create_config(&mut self, room: &Joined) -> anyhow::Result { + let config = RoomConfig { + id: Some(room.room_id().to_owned()), + ..Default::default() + }; + + Ok(self.rooms.insert(config)) + } + + pub async fn find_or_create(&mut self, room: &Joined, id: OwnedRoomId) -> anyhow::Result<&RoomConfig> { + if self.find(id.to_owned()).await.is_none() { + self.create_config(room).await?; + } + + Ok(self.find(id.to_owned()).await.unwrap()) + } + + pub async fn save(&mut self, config: &RoomConfig) -> anyhow::Result<()> { + self.rooms.take(self.to_owned().find(config.id.to_owned().unwrap()).await.unwrap()); + self.rooms.insert(config.to_owned()); + + let serialized = serde_json::to_string_pretty(self)?; + fs::write(crate::ROOMS_CONFIG_FILE.to_owned(), serialized).await?; + + Ok(()) + } + + pub fn restore() -> anyhow::Result { + if !std::path::Path::new(crate::ROOMS_CONFIG_FILE.to_owned().to_str().unwrap()).exists() { + return Ok(Self { rooms: BTreeSet::::new() }); + } + + let serialized = std::fs::read_to_string(crate::ROOMS_CONFIG_FILE.to_owned())?; + let ctrl: Result = serde_json::from_str(&serialized); + if ctrl.is_err() { + return Ok(Self { rooms: BTreeSet::::new() }); + } + Ok(ctrl.unwrap()) + } +} + + #[derive(Debug)] pub struct Config { pub keywords: Value, diff --git a/src/judge.rs b/src/judge.rs index 5838992..445987a 100644 --- a/src/judge.rs +++ b/src/judge.rs @@ -1,5 +1,6 @@ use crate::{ keywords::{KeywordCategory, Keywords}, + config::RoomConfig, CONFIG, }; use matrix_sdk::{ @@ -30,7 +31,7 @@ pub struct Judgement { } impl Judgement { - pub fn judge(&self) -> anyhow::Result { + pub fn judge(&self, config: &RoomConfig) -> anyhow::Result { // Load keywords let mut keywords = CONFIG.keywords.clone(); let keywords = keywords @@ -67,7 +68,7 @@ impl Judgement { Ok(JudgementResult::LikelyScam) } - pub async fn send_debug(&self, room: &Joined) -> anyhow::Result<()> { + pub async fn send_debug(&self, config: &RoomConfig, room: &Joined) -> anyhow::Result<()> { // Load keywords let mut keywords = CONFIG.keywords.clone(); let keywords = keywords @@ -105,19 +106,24 @@ impl Judgement { // Send message let msg = RoomMessageEventContent::text_html( - format!("{counter:?}, {count_all}/{total}, {result:?}"), - format!("{counter:?}
Categories covered: {count_all}/{total}
Verdict: {result:?}")); + format!("{counter:?}\nCategories covered: {count_all}/{total}\nVerdict: {result:?}\n{config:?}"), + format!("{counter:?}
Categories covered: {count_all}/{total}
Verdict: {result:?}
{config:?}")); room.send(msg, None).await.expect("Couldn't send message"); Ok(()) } pub async fn alert( + config: &RoomConfig, room: &Joined, event: &OriginalRoomMessageEvent, result: JudgementResult, is_reply: bool, ) -> anyhow::Result<()> { + if config.no_reply { + return Ok(()); + } + let mut responses = CONFIG.responses.clone(); let responses = responses.as_object_mut().unwrap(); @@ -139,25 +145,28 @@ impl Judgement { // Send message let msg = RoomMessageEventContent::text_html(plain, html); - room.send(msg, None).await.expect("Couldn't send message"); - // Todo: Add room config? -// let reply = msg.make_reply_to(event); -// room.send(reply, None).await.expect("Couldn't send message"); + + if config.reply_to_scam { + let reply = msg.make_reply_to(event); + room.send(reply, None).await.expect("Couldn't send message"); + } else { + room.send(msg, None).await.expect("Couldn't send message"); + } // Send reaction if !is_reply { room.send_raw( json!({ - "m.relates_to": { - "rel_type": "m.annotation", - "event_id": event.event_id.to_string(), - "key": "🚨🚨 SCAM 🚨🚨" - }}), - "m.reaction", - None, - ) - .await - .expect("Couldn't send reaction"); + "m.relates_to": { + "rel_type": "m.annotation", + "event_id": event.event_id.to_string(), + "key": "🚨🚨 SCAM 🚨🚨" + }}), + "m.reaction", + None, + ) + .await + .expect("Couldn't send reaction"); } Ok(()) diff --git a/src/main.rs b/src/main.rs index 3157190..10cf64b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ use matrix_sdk::{ }, Error, LoopCtrl, }; +use std::path::PathBuf; use once_cell::sync::Lazy; pub mod config; @@ -15,10 +16,17 @@ pub mod judge; pub mod keywords; pub mod matrix; +static DATA_DIR: Lazy = Lazy::new(|| dirs::data_dir().expect("No data_dir found").join("scam_police")); +static SESSION_FILE: Lazy = Lazy::new(|| DATA_DIR.join("session")); +static ROOMS_CONFIG_FILE: Lazy = Lazy::new(|| DATA_DIR.join("rooms_config")); + static CONFIG: Lazy = Lazy::new(|| config::Config::load()); async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) -> anyhow::Result<()> { if let Room::Joined(room) = room { + let mut room_ctrl = config::RoomConfigController::restore().expect("Couldn't restore room configs"); + let mut config: config::RoomConfig = room_ctrl.find_or_create(&room, room.room_id().to_owned()).await?.to_owned(); + let orig_event = event .to_owned() .into_full_event(OwnedRoomId::from(room.room_id())); @@ -28,13 +36,30 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) -> any return Ok(()); }; - // Too short to be a scam lol - if text_content.body.chars().count() < 12 { - return Ok(()); + let text_content = text_content.body.to_lowercase(); + let mut debug = false; + if text_content.contains(";sp") { + let command = unsafe { text_content.get_unchecked((text_content.find(";sp").unwrap()+3)..) }; + + match command { + "debug" => debug = true, + + "mute" => config.no_reply = true, + "unmute" => config.no_reply = false, + + "replyto on" => config.reply_to_scam = true, + "replyto off" => config.reply_to_scam = false, + + _ => () + } + + room_ctrl.save(&config).await?; } - let text_content = text_content.body.to_lowercase(); - let debug = text_content.contains(";spdebug"); +// // Too short to be a scam lol +// if text_content.chars().count() < 12 { +// return Ok(()); +// } // Ignore own messages if !debug @@ -83,15 +108,16 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) -> any let reply_judgement = judge::Judgement { text: content }; if debug { - reply_judgement.send_debug(&room).await?; + reply_judgement.send_debug(&config, &room).await?; return Ok(()); } - match reply_judgement.judge()? { + match reply_judgement.judge(&config)? { judge::JudgementResult::Ok => (), judge::JudgementResult::MaybeScam => (), judge::JudgementResult::LikelyScam => { judge::Judgement::alert( + &config, &room, &orig_event, judge::JudgementResult::LikelyScam, @@ -104,11 +130,12 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) -> any } } - match judgement.judge()? { + match judgement.judge(&config)? { judge::JudgementResult::Ok => return Ok(()), judge::JudgementResult::MaybeScam => return Ok(()), judge::JudgementResult::LikelyScam => { judge::Judgement::alert( + &config, &room, &orig_event, judge::JudgementResult::LikelyScam, @@ -141,7 +168,7 @@ async fn main() -> anyhow::Result<()> { ) } else { anyhow::bail!( - "No previous session found, please run \"{} \"", + "No previous session found, please run \"{} \" to login", args.get(0).unwrap() ); }; diff --git a/src/matrix.rs b/src/matrix.rs index ed7e834..60776d7 100644 --- a/src/matrix.rs +++ b/src/matrix.rs @@ -198,18 +198,12 @@ pub async fn sync<'a>( } pub async fn persist_sync_token(sync_token: String) -> anyhow::Result<()> { - let data_dir = dirs::data_dir() - .expect("no data_dir directory found") - .join("scam_police"); - let session_file = data_dir.join("session"); - - let serialized_session = fs::read_to_string(&session_file).await?; + let serialized_session = fs::read_to_string(crate::SESSION_FILE.to_owned()).await?; let mut full_session: FullSession = from_str(&serialized_session)?; full_session.sync_token = Some(sync_token); let serialized_session = serde_json::to_string(&full_session)?; - fs::write(session_file, serialized_session).await?; - + fs::write(crate::SESSION_FILE.to_owned(), serialized_session).await?; Ok(()) }