Add Homeserver and MXID resolve functions; Make login use MXID in the command line instead of stdin

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",

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"

src/main.rs

@@ -85,6 +85,8 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    let args: Vec<String> = std::env::args().collect();
+
     let data_dir = dirs::data_dir()
         .expect("no data_dir directory found")
         .join("scam_police");
@@ -92,8 +94,10 @@ async fn main() -> anyhow::Result<()> {
 
     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)

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<String>,
 }
 
-pub async fn restore_session(session_file: &Path) -> anyhow::Result<(Client, Option<String>)> {
-    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<Client> {
+    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<Client> {
-    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,8 +69,7 @@ pub async fn login(data_dir: &Path, session_file: &Path) -> anyhow::Result<Clien
     Ok(client)
 }
 
-/// Build a new client.
-pub async fn build_client(data_dir: &Path) -> 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)
@@ -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()
+            .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,8 +106,7 @@ 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());
@@ -145,8 +114,56 @@ pub async fn build_client(data_dir: &Path) -> anyhow::Result<(Client, ClientSess
             },
         }
 }
+
+//
+// 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<String> {
+    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<String>,
@@ -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<String>)> {
+    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))
+}