main.rs (14364B)
1 use std::convert::TryFrom; 2 use std::fs; 3 use std::fs::File; 4 use std::io::{BufReader, BufWriter, Read, Write}; 5 6 use clap::Clap; 7 use futures_util::stream::FuturesUnordered; 8 use futures_util::stream::TryStreamExt as _; 9 use log::{debug, error, info, warn}; 10 use rand::distributions::Alphanumeric; 11 use rand::{thread_rng, Rng}; 12 use ruma_client::api::r0::membership::invite_user; 13 use ruma_client::api::r0::membership::join_room_by_id; 14 use ruma_client::api::r0::message::create_message_event; 15 use ruma_client::api::r0::sync::sync_events::IncomingResponse; 16 use ruma_client::identifiers::UserId; 17 use ruma_client::{ 18 self, 19 api::r0::{ 20 filter::{FilterDefinition, RoomEventFilter, RoomFilter}, 21 message::create_message_event::Response, 22 room::create_room, 23 sync::sync_events::Filter, 24 }, 25 events::{ 26 collections::all::RoomEvent::{self, RoomMessage}, 27 room::message::{ 28 InReplyTo, MessageEvent, 29 MessageEventContent::{self, Text}, 30 NoticeMessageEventContent, RelatesTo, TextMessageEventContent, 31 }, 32 EventType, 33 }, 34 identifiers::{EventId, RoomId}, 35 HttpsClient, Session, 36 }; 37 use url::Url; 38 39 use crate::config::{load_config, Config, Homeserver}; 40 use crate::logger::setup_logger; 41 42 mod config; 43 mod logger; 44 45 async fn send_notice_reply( 46 client: &HttpsClient, 47 text: String, 48 related_event: EventId, 49 room_id: RoomId, 50 ) -> Result<Response, ruma_client::Error> { 51 let rand_string: String = thread_rng().sample_iter(&Alphanumeric).take(30).collect(); 52 client 53 .request(create_message_event::Request { 54 room_id, 55 event_type: EventType::RoomMessage, 56 txn_id: rand_string, 57 data: MessageEventContent::Notice(NoticeMessageEventContent { 58 body: text, 59 relates_to: Some(RelatesTo { 60 in_reply_to: InReplyTo { 61 event_id: related_event, 62 }, 63 }), 64 }), 65 }) 66 .await 67 } 68 69 async fn send_message( 70 client: &HttpsClient, 71 text: String, 72 room_id: RoomId, 73 ) -> Result<Response, ruma_client::Error> { 74 let rand_string: String = thread_rng().sample_iter(&Alphanumeric).take(30).collect(); 75 client 76 .request(create_message_event::Request { 77 room_id, 78 event_type: EventType::RoomMessage, 79 txn_id: rand_string, 80 data: MessageEventContent::Text(TextMessageEventContent { 81 body: text, 82 format: None, 83 formatted_body: None, 84 relates_to: None, 85 }), 86 }) 87 .await 88 } 89 90 async fn get_control_room( 91 client: &HttpsClient, 92 target_user: String, 93 ) -> Result<RoomId, failure::Error> { 94 fs::create_dir_all("./tmp/")?; 95 let mut room_id_read = String::new(); 96 match File::open(format!( 97 "./tmp/control_room_{}", 98 client.session().unwrap().user_id.hostname() 99 )) { 100 Ok(f) => { 101 let mut br = BufReader::new(f); 102 match br.read_to_string(&mut room_id_read) { 103 Ok(_) => { 104 let room_id = RoomId::try_from(room_id_read.as_str()).unwrap(); 105 Ok(room_id) 106 } 107 Err(e) => Err(failure::Error::try_from(e)?), 108 } 109 } 110 Err(e) => { 111 warn!("Unable to open control_room_tmp: {}", e); 112 let resp = client 113 .request(create_room::Request { 114 invite: vec![UserId::try_from(target_user.as_str())?], 115 name: Some(format!( 116 "AutoInviteBot Control Room {}", 117 client.session().unwrap().user_id.hostname() 118 )), 119 preset: Some(create_room::RoomPreset::TrustedPrivateChat), 120 creation_content: None, 121 room_alias_name: None, 122 topic: None, 123 visibility: None, 124 }) 125 .await?; 126 let room_id: RoomId = resp.room_id; 127 let f = File::create(format!( 128 "./tmp/control_room_{}", 129 client.session().unwrap().user_id.hostname() 130 ))?; 131 let mut f = BufWriter::new(f); 132 f.write_all(room_id.to_string().as_bytes())?; 133 134 Ok(room_id) 135 } 136 } 137 } 138 139 async fn get_client( 140 server: &Homeserver, 141 target_user: String, 142 ) -> Result<(HttpsClient, RoomId), failure::Error> { 143 info!("Starting session as {}", server.mxid); 144 145 let session: Session; 146 let mut client: HttpsClient = HttpsClient::https( 147 Url::parse(server.address.as_str()).expect("unable to parse url"), 148 None, 149 ); 150 if server.access_token.is_some() { 151 session = Session { 152 device_id: "Ruma Bot".to_string(), 153 access_token: server 154 .access_token 155 .as_ref() 156 .expect("missing access_token") 157 .to_owned(), 158 user_id: UserId::try_from(server.mxid.as_str()).unwrap(), 159 }; 160 161 client = HttpsClient::https( 162 Url::parse(server.address.as_str()).expect("unable to parse url"), 163 Some(session), 164 ); 165 } else if server.password.is_some() { 166 client 167 .log_in( 168 server.mxid.clone(), 169 server 170 .password 171 .as_ref() 172 .expect("failed to read password") 173 .to_owned(), 174 None, 175 None, 176 ) 177 .await?; 178 } else { 179 error!("Please provide either a password or an access_token!"); 180 } 181 182 let control_room_id = get_control_room(&client, target_user).await?; 183 184 info!("Started session as {}", server.mxid); 185 186 Ok((client, control_room_id)) 187 } 188 189 async fn parse_invites( 190 client: &HttpsClient, 191 room_id: RoomId, 192 target_user: String, 193 message: String, 194 ) -> Result<(), ruma_client::Error> { 195 // Auto join rooms 196 debug!("Invited to {:?}", room_id); 197 client 198 .request(join_room_by_id::Request { 199 room_id: room_id.clone(), 200 third_party_signed: None, 201 }) 202 .await?; 203 debug!("Joined {:?}", room_id.clone()); 204 let response = client 205 .request(invite_user::Request { 206 room_id: room_id.clone(), 207 user_id: UserId::try_from(target_user.as_str()).unwrap(), 208 }) 209 .await?; 210 debug!("Invited correct user {:?}", response); 211 212 let rand_string: String = thread_rng().sample_iter(&Alphanumeric).take(30).collect(); 213 client 214 .request(create_message_event::Request { 215 room_id: room_id.clone(), 216 event_type: EventType::RoomMessage, 217 txn_id: rand_string, 218 data: MessageEventContent::Notice(NoticeMessageEventContent { 219 body: message, 220 relates_to: None, 221 }), 222 }) 223 .await?; 224 225 debug!("Sent a message about what happened"); 226 Ok(()) 227 } 228 229 async fn handle_mention( 230 client: &HttpsClient, 231 config: &Config, 232 control_room_id: RoomId, 233 room_id: RoomId, 234 sender: UserId, 235 event_id: EventId, 236 body: String, 237 ) -> Result<(), failure::Error> { 238 if config.debug { 239 send_notice_reply( 240 &client, 241 format!( 242 "> <{}> {}\n\n[DEBUG] Mentioned main account about this mention", 243 sender, body 244 ), 245 event_id.clone(), 246 room_id.clone(), 247 ) 248 .await?; 249 } 250 251 // Send mention to control room 252 send_message( 253 &client, 254 format!( 255 "> <{}> {}\n\n Mention in {} by {}", 256 sender, body, room_id, sender 257 ), 258 control_room_id.clone(), 259 ) 260 .await?; 261 262 Ok(()) 263 } 264 265 async fn parse_joins( 266 client: &HttpsClient, 267 config: &Config, 268 server: &Homeserver, 269 event: RoomEvent, 270 room_id: RoomId, 271 control_room_id: RoomId, 272 ) -> Result<(), failure::Error> { 273 if let RoomMessage(msg) = event { 274 let content = msg.content; 275 let event_id = msg.event_id; 276 let sender = msg.sender; 277 if let Text(msg_content) = content { 278 let formatted_body: Option<String> = msg_content.formatted_body; 279 let body: String = msg_content.body; 280 281 if sender.to_string() != client.session().unwrap().user_id.to_string() 282 && formatted_body.clone().is_some() 283 && formatted_body 284 .clone() 285 .unwrap() 286 .to_lowercase() 287 .contains(&server.mxid.clone().to_lowercase()) 288 { 289 handle_mention( 290 client, 291 config, 292 control_room_id.clone(), 293 room_id.clone(), 294 sender.clone(), 295 event_id.clone(), 296 body.clone(), 297 ) 298 .await?; 299 } else if sender.to_string() != client.session().unwrap().user_id.to_string() 300 && formatted_body.clone().is_none() 301 && body.to_lowercase().contains( 302 server 303 .mxid 304 .clone() 305 .to_lowercase() 306 .split(':') 307 .collect::<Vec<_>>()[0], 308 ) 309 { 310 handle_mention( 311 client, 312 config, 313 control_room_id.clone(), 314 room_id.clone(), 315 sender.clone(), 316 event_id.clone(), 317 body.clone(), 318 ) 319 .await?; 320 } 321 // Todo make a smarter command handler 322 /*if body.starts_with("!test") { 323 send_notice_reply( 324 &client, 325 format!("> <{}> {}\n\ntest resp", sender, body), 326 event_id.clone(), 327 room_id.clone(), 328 ) 329 .await?; 330 }*/ 331 } 332 } 333 Ok(()) 334 } 335 336 async fn load_next_batch() -> Result<String, failure::Error> { 337 fs::create_dir_all("./tmp/")?; 338 let mut next_batch = String::new(); 339 let f = File::open("./tmp/next_batch")?; 340 let mut br = BufReader::new(f); 341 br.read_to_string(&mut next_batch)?; 342 Ok(next_batch) 343 } 344 345 async fn save_next_batch(next_batch: String) -> Result<(), failure::Error> { 346 let f = File::create("./tmp/next_batch")?; 347 let mut f = BufWriter::new(f); 348 f.write_all(next_batch.as_bytes())?; 349 Ok(()) 350 } 351 352 fn generate_filter() -> Result<Filter, failure::Error> { 353 let filter = FilterDefinition { 354 event_fields: None, 355 event_format: None, 356 account_data: None, 357 room: Some(RoomFilter { 358 include_leave: None, 359 account_data: None, 360 timeline: Some(RoomEventFilter { 361 not_types: vec![], 362 not_rooms: vec![], 363 limit: None, 364 rooms: None, 365 not_senders: vec![], 366 senders: None, 367 types: Some(vec!["m.room.message".to_owned()]), 368 contains_url: None, 369 }), 370 ephemeral: None, 371 state: None, 372 not_rooms: vec![], 373 rooms: None, 374 }), 375 presence: None, 376 }; 377 Ok(Filter::FilterDefinition(filter)) 378 } 379 380 async fn do_stuff(config: &Config, server: &Homeserver) -> Result<(), failure::Error> { 381 let target_user = config.target_user.clone(); 382 let (client, control_room_id) = get_client(server, target_user.clone()).await?; 383 384 let since = match load_next_batch().await { 385 Ok(v) => Some(v), 386 Err(e) => { 387 error!("{:?}", e); 388 None 389 } 390 }; 391 392 let filter = match generate_filter() { 393 Ok(v) => Some(v), 394 Err(e) => { 395 error!("{:?}", e); 396 None 397 } 398 }; 399 400 let mut sync_stream = Box::pin(client.sync(filter, since, false)); 401 let message = config.message.clone(); 402 while let Some(response) = sync_stream.try_next().await? { 403 let res: IncomingResponse = response; 404 let next_batch: String = res.next_batch; 405 save_next_batch(next_batch).await?; 406 407 for (room_id, _room) in res.rooms.invite { 408 parse_invites( 409 &client, 410 room_id.clone(), 411 target_user.clone(), 412 message.clone(), 413 ) 414 .await?; 415 } 416 for (room_id, room) in res.rooms.join { 417 let events = room.timeline.events; 418 419 for event in events.into_iter().flat_map(|r| r.into_result()) { 420 if let RoomEvent::RoomMessage(MessageEvent { .. }) = event { 421 parse_joins( 422 &client, 423 config, 424 server, 425 event.clone(), 426 room_id.clone(), 427 control_room_id.clone(), 428 ) 429 .await?; 430 } 431 } 432 } 433 } 434 435 Ok(()) 436 } 437 438 #[derive(Clap)] 439 #[clap(version = "0.1.0", author = "MTRNord")] 440 struct Opts { 441 #[clap(short = 'c', long = "config", default_value = "config.yaml")] 442 config: String, 443 /// A level of verbosity, and can be used multiple times 444 #[clap(short = 'v', long = "verbose", parse(from_occurrences))] 445 verbose: i32, 446 } 447 448 #[tokio::main] 449 async fn main() -> Result<(), failure::Error> { 450 let opts: Opts = Opts::parse(); 451 452 let log_level = match opts.verbose { 453 0 => log::LevelFilter::Info, // default 454 1 => log::LevelFilter::Error, // -v 455 2 => log::LevelFilter::Warn, // -vv 456 3 => log::LevelFilter::Debug, // -vvv 457 _ => log::LevelFilter::Trace, // -vvvv and above 458 }; 459 460 setup_logger(log_level).expect("unable to setup logger"); 461 462 let config = load_config(opts.config).expect("unable to read config"); 463 464 let mut futures = FuturesUnordered::new(); 465 466 config.servers.iter().for_each(|x| { 467 futures.push(do_stuff(&config, x)); 468 }); 469 470 while let Some(x) = futures.try_next().await? { 471 error!("{:?}", x); 472 } 473 474 Ok(()) 475 }