From 69d08d4dfb35ef5598c4f808b0ca3942dde15c7e Mon Sep 17 00:00:00 2001 From: 0xf8 <0xf8.dev@proton.me> Date: Wed, 12 Apr 2023 16:18:38 -0400 Subject: [PATCH] Add Homeserver and MXID resolve functions; Make login use MXID in the command line instead of stdin --- Cargo.lock | 3 +- Cargo.toml | 3 +- src/main.rs | 10 ++-- src/matrix.rs | 138 ++++++++++++++++++++++++++++++++------------------ 4 files changed, 99 insertions(+), 55 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2719afe..33c2f24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2067,13 +2067,14 @@ checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041" [[package]] name = "scam-police" -version = "0.2.0" +version = "0.3.0" dependencies = [ "anyhow", "dirs", "matrix-sdk", "once_cell", "rand 0.8.5", + "reqwest", "rpassword", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 3db19b5..134156d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "scam-police" -version = "0.2.0" +version = "0.3.0" edition = "2021" authors = [ "@0xf8:projectsegfau.lt", "@jjj333:pain.agency" ] @@ -12,6 +12,7 @@ dirs = "5.0.0" matrix-sdk = "0.6.2" once_cell = "1.17.1" rand = "0.8.5" +reqwest = "0.11.16" rpassword = "7.2.0" serde = "1.0.160" serde_json = "1.0.95" diff --git a/src/main.rs b/src/main.rs index 1e268c2..dbc7e0f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -85,15 +85,19 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) { #[tokio::main] async fn main() -> anyhow::Result<()> { + let args: Vec = std::env::args().collect(); + let data_dir = dirs::data_dir() .expect("no data_dir directory found") .join("scam_police"); let session_file = data_dir.join("session"); - + let (client, sync_token) = if session_file.exists() { matrix::restore_session(&session_file).await? + } else if args.len() > 1 { + (matrix::login(&data_dir, &session_file, args.get(1).unwrap().to_owned()).await?, None) } else { - (matrix::login(&data_dir, &session_file).await?, None) + anyhow::bail!("No previous session found, please run {} MXID", args.get(1).unwrap()); }; let (client, sync_settings) = matrix::sync(client, sync_token) @@ -108,7 +112,7 @@ async fn main() -> anyhow::Result<()> { matrix::persist_sync_token(response.next_batch) .await .map_err(|err| Error::UnknownError(err.into()))?; - + Ok(LoopCtrl::Continue) }).await?; diff --git a/src/matrix.rs b/src/matrix.rs index ed0f263..057880c 100644 --- a/src/matrix.rs +++ b/src/matrix.rs @@ -3,10 +3,12 @@ use matrix_sdk::{ }; use rand::{distributions::Alphanumeric, thread_rng, Rng}; use serde::{Deserialize, Serialize}; +use serde_json::{Value, from_str}; use std::{ io::{self, Write}, path::{Path, PathBuf}, }; +use reqwest::Client as http; use tokio::fs; #[derive(Debug, Serialize, Deserialize)] @@ -24,39 +26,19 @@ pub struct FullSession { sync_token: Option, } -pub async fn restore_session(session_file: &Path) -> anyhow::Result<(Client, Option)> { - let serialized_session = fs::read_to_string(session_file).await?; - let FullSession { - client_session, - user_session, - sync_token, - } = serde_json::from_str(&serialized_session)?; - let client = Client::builder() - .homeserver_url(client_session.homeserver) - .sled_store(client_session.db_path, Some(&client_session.passphrase))? - .build() - .await?; +// +// Matrix Login & Init +// - println!("[*] Restoring session for {}…", user_session.user_id); +pub async fn login(data_dir: &Path, session_file: &Path, mxid: String) -> anyhow::Result { + println!("[*] No previous session, loggin in with mxid..."); - client.restore_login(user_session).await?; - - Ok((client, sync_token)) -} - -pub async fn login(data_dir: &Path, session_file: &Path) -> anyhow::Result { - println!("[*] No previous session found, logging in…"); - - let (client, client_session) = build_client(data_dir).await?; + let (user, hs) = resolve_mxid(mxid).await?; + let (client, client_session) = build_client(data_dir, hs).await?; loop { - let mut user = String::new(); - io::stdout().flush().expect("Unable to write to stdout"); - io::stdin() - .read_line(&mut user) - .expect("Unable to read user input"); - let password = rpassword::prompt_password("Password> ").unwrap(); + let password = rpassword::prompt_password("Password\n> ").unwrap(); match client .login_username(&user, &password) @@ -87,10 +69,9 @@ pub async fn login(data_dir: &Path, session_file: &Path) -> anyhow::Result anyhow::Result<(Client, ClientSession)> { +pub async fn build_client(data_dir: &Path, hs: String) -> anyhow::Result<(Client, ClientSession)> { let mut rng = thread_rng(); - + let db_subfolder: String = (&mut rng) .sample_iter(Alphanumeric) .take(7) @@ -104,29 +85,18 @@ pub async fn build_client(data_dir: &Path) -> anyhow::Result<(Client, ClientSess .map(char::from) .collect(); - loop { - let mut homeserver = String::new(); - - print!("Homeserver> "); - io::stdout().flush().expect("Unable to write to stdout"); - io::stdin() - .read_line(&mut homeserver) - .expect("Unable to read user input"); - - println!("[*] Checking homeserver…"); - - match Client::builder() - .homeserver_url(&homeserver) - .sled_store(&db_path, Some(&passphrase)) - .unwrap() + match Client::builder() + .homeserver_url(&hs) + .sled_store(&db_path, Some(&passphrase))? .build() .await { Ok(client) => { + println!("[*] Homeserver OK"); return Ok(( client, ClientSession { - homeserver, + homeserver: hs, db_path, passphrase, }, @@ -136,17 +106,64 @@ pub async fn build_client(data_dir: &Path) -> anyhow::Result<(Client, ClientSess matrix_sdk::ClientBuildError::AutoDiscovery(_) | matrix_sdk::ClientBuildError::Url(_) | matrix_sdk::ClientBuildError::Http(_) => { - println!("[!] Error checking the homeserver: {error}"); - println!("[!] Please try again\n"); + anyhow::bail!("[!] {error:?}"); } _ => { return Err(error.into()); } }, } +} + +// +// Helper Functions +// + +// Resolve mxid into user and hs +pub async fn resolve_mxid(mxid: String) -> anyhow::Result<(String, String)> { + if mxid.get(0..1).unwrap() != "@" || !mxid.contains(":") { + anyhow::bail!("Invalid mxid"); + } + + let sep = mxid.find(":").unwrap(); + let user = mxid.get(1..sep).unwrap().to_string(); + let hs = resolve_homeserver(mxid.get((sep + 1)..).unwrap().to_string()).await?; + + Ok((user, hs)) +} + +// Resolve homeserver +pub async fn resolve_homeserver(homeserver: String) -> anyhow::Result { + let mut hs = homeserver; + if !hs.contains("://") { + hs = format!("https://{hs}"); + } + + if hs.chars().last().unwrap().to_string() == "/" { + hs.pop(); + } + + let ident = http::new().get(format!("{hs}/.well-known/matrix/client")).send().await; + match ident { + Ok(r) => { + let body = r.text().await?; + let json: Value = from_str(&body)?; + + let discovered = json["m.homeserver"]["base_url"].as_str().unwrap(); + + Ok(discovered.to_string()) + }, + Err(e) => { + Err(e.into()) + } } } + +// +// Persistence +// + pub async fn sync<'a>( client: Client, initial_sync_token: Option, @@ -187,7 +204,7 @@ pub async fn persist_sync_token(sync_token: String) -> anyhow::Result<()> { let session_file = data_dir.join("session"); let serialized_session = fs::read_to_string(&session_file).await?; - let mut full_session: FullSession = serde_json::from_str(&serialized_session)?; + 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)?; @@ -195,3 +212,24 @@ pub async fn persist_sync_token(sync_token: String) -> anyhow::Result<()> { Ok(()) } + +pub async fn restore_session(session_file: &Path) -> anyhow::Result<(Client, Option)> { + let serialized_session = fs::read_to_string(session_file).await?; + let FullSession { + client_session, + user_session, + sync_token, + } = from_str(&serialized_session)?; + + let client = Client::builder() + .homeserver_url(client_session.homeserver) + .sled_store(client_session.db_path, Some(&client_session.passphrase))? + .build() + .await?; + + println!("[*] Restoring session for {}…", user_session.user_id); + + client.restore_login(user_session).await?; + + Ok((client, sync_token)) +} \ No newline at end of file