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

This commit is contained in:
0xf8 2023-04-12 16:18:38 -04:00
parent 931236b607
commit 69d08d4dfb
Signed by: 0xf8
GPG Key ID: 446580D758689584
4 changed files with 99 additions and 55 deletions

3
Cargo.lock generated
View File

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

View File

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

View File

@ -85,15 +85,19 @@ 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");
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?;

View File

@ -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,10 +69,9 @@ 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)
.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<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))
}