diff --git a/Cargo.lock b/Cargo.lock index b5441e1..e2cb601 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -55,6 +55,12 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" +[[package]] +name = "argparse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f8ebf5827e4ac4fd5946560e6a99776ea73b596d80898f357007317a7141e47" + [[package]] name = "arrayref" version = "0.3.7" @@ -387,6 +393,31 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossterm" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64e6c0fbe2c17357405f7c758c1ef960fce08bdfb2c03d88d2a18d7e09c4b67" +dependencies = [ + "bitflags", + "crossterm_winapi", + "libc", + "mio", + "parking_lot 0.12.1", + "signal-hook", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ae1b35a484aa10e07fe0638d02301c5ad24de82d310ccbd2f3693da5f09bf1c" +dependencies = [ + "winapi", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -1466,6 +1497,16 @@ dependencies = [ "parking_lot_core 0.8.6", ] +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core 0.9.7", +] + [[package]] name = "parking_lot_core" version = "0.8.6" @@ -2070,6 +2111,7 @@ name = "scam-police" version = "0.6.0" dependencies = [ "anyhow", + "argparse", "dirs", "matrix-sdk", "once_cell", @@ -2079,6 +2121,7 @@ dependencies = [ "serde", "serde_json", "strfmt", + "terminal-menu", "tokio", "url", ] @@ -2197,6 +2240,36 @@ dependencies = [ "digest 0.10.6", ] +[[package]] +name = "signal-hook" +version = "0.3.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "732768f1176d21d09e076c23a93123d40bba92d50c4058da34d45c8de8e682b9" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29ad2e15f37ec9a6cc544097b78a1ec90001e9f71b81338ca39f430adaca99af" +dependencies = [ + "libc", + "mio", + "signal-hook", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + [[package]] name = "signature" version = "1.6.4" @@ -2225,7 +2298,7 @@ dependencies = [ "fxhash", "libc", "log", - "parking_lot", + "parking_lot 0.11.2", ] [[package]] @@ -2317,6 +2390,16 @@ dependencies = [ "windows-sys 0.45.0", ] +[[package]] +name = "terminal-menu" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df4f5aa03b86607186c90883734a1c5751e18828f7c2f96f94a282ec1a1bd7e5" +dependencies = [ + "crossterm", + "lazy_static", +] + [[package]] name = "thiserror" version = "1.0.40" @@ -2691,7 +2774,7 @@ checksum = "be0ecb0db480561e9a7642b5d3e4187c128914e58aa84330b9493e3eb68c5e7f" dependencies = [ "futures", "js-sys", - "parking_lot", + "parking_lot 0.11.2", "pin-utils", "wasm-bindgen", "wasm-bindgen-futures", diff --git a/Cargo.toml b/Cargo.toml index 4119bf7..9c72c90 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ authors = [ "@0xf8:projectsegfau.lt", "@jjj333:pain.agency" ] [dependencies] anyhow = "1.0.70" +argparse = "0.2.2" dirs = "5.0.0" matrix-sdk = "0.6.2" once_cell = "1.17.1" @@ -17,5 +18,6 @@ rpassword = "7.2.0" serde = "1.0.160" serde_json = "1.0.95" strfmt = "0.2.4" +terminal-menu = "2.0.5" tokio = { version = "1.27.0", features = ["macros", "rt-multi-thread"] } url = "2.3.1" diff --git a/src/matrix.rs b/src/matrix.rs index 60776d7..1c78cf0 100644 --- a/src/matrix.rs +++ b/src/matrix.rs @@ -1,5 +1,5 @@ use matrix_sdk::{ - config::SyncSettings, ruma::api::client::filter::FilterDefinition, Client, Session, + config::SyncSettings, ruma::api::client::{filter::FilterDefinition, session::get_login_types::v3::{IdentityProvider, LoginType}}, Client, Session, }; use rand::{distributions::Alphanumeric, thread_rng, Rng}; use reqwest::Client as http; @@ -24,6 +24,77 @@ pub struct FullSession { sync_token: Option, } +#[derive(Debug, Clone)] +enum LoginChoice { + Password, + Sso, + SsoIdp(IdentityProvider), +} + +impl LoginChoice { + pub async fn login(&self, client: &Client, user: String, hs: String) -> anyhow::Result { + match self { + LoginChoice::Password => Self::login_password(client, user, hs).await, + LoginChoice::Sso => Self::login_sso(client, None).await, + LoginChoice::SsoIdp(idp) => Self::login_sso(client, Some(idp.to_owned())).await, + } + } + + async fn login_password(client: &Client, user: String, _hs: String) -> anyhow::Result { + loop { + let password = prompt_password("Password\n> ")?; + + match client + .login_username(&user, &password) + .initial_device_display_name("scam-police") + .send() + .await + { + Ok(_) => { + println!("[*] Logged in as {user}"); + return Ok(client.to_owned()); + } + Err(e) => { + println!("[!] Error logging in: {e}"); + println!("[!] Please try again\n"); + } + } + } + } + + async fn login_sso(client: &Client, idp: Option) -> anyhow::Result { + let redirect_url = String::new(); + + let token = if let Some(idp) = idp { + client.get_sso_login_url(&redirect_url, Some(&idp.id)).await + } else { + client.get_sso_login_url(&redirect_url, None).await + }; + + println!("{redirect_url:?}, {token:?}"); + + match token { + Ok(t) => { + match client + .login_token(&t) + .initial_device_display_name("scam-police") + .send() + .await { + Ok(_) => { + Ok(client.to_owned()) + } + Err(e) => { + anyhow::bail!("{e}") + } + } + }, + Err(e) => { + anyhow::bail!("Failed to get SSO token: {e}") + } + } + } +} + // // Matrix Login & Init // @@ -32,28 +103,66 @@ pub async fn login(data_dir: &Path, session_file: &Path, mxid: String) -> anyhow println!("[*] No previous session, logging in with mxid..."); let (user, hs) = resolve_mxid(mxid).await?; - let (client, client_session) = build_client(data_dir, hs).await?; + let (client, client_session) = build_client(data_dir, hs.to_owned()).await?; - loop { - let password = prompt_password("Password\n> ")?; - - match client - .login_username(&user, &password) - .initial_device_display_name("scam-police") - .send() - .await - { - Ok(_) => { - println!("[*] Logged in as {user}"); - break; - } - Err(error) => { - println!("[!] Error logging in: {error}"); - println!("[!] Please try again\n"); - } + let mut login_choices = Vec::new(); + for login_type in client.get_login_types().await?.flows { + match login_type { + LoginType::Password(_) => { + login_choices.push(LoginChoice::Password); + }, + LoginType::Sso(sso) => { + if sso.identity_providers.is_empty() { + login_choices.push(LoginChoice::Sso); + } else { + login_choices.extend(sso.identity_providers.into_iter().map(LoginChoice::SsoIdp)); + } + }, + // Ignore all other types + _ => {}, } } + let client = match login_choices.to_owned().len() { + 0 => anyhow::bail!("No supported login types"), + 1 => login_choices.to_owned().get(0).unwrap().login(&client, user, hs.to_owned()).await, + _ => { + use terminal_menu::*; + + let mut menu_items = vec![label("-- Scam Police Login --")]; + let choices: Vec<(LoginChoice, String)> = login_choices.into_iter().map(|a| (a.to_owned(), match a { + LoginChoice::Password => format!("Password"), + LoginChoice::Sso => format!("SSO"), + LoginChoice::SsoIdp(idp) => format!("SSO ({})", idp.name), + })).collect(); + + for choice in choices.to_owned() { + menu_items.push(button(choice.1)); + } + let menu = menu(menu_items); + run(&menu); + + let menu = mut_menu(&menu); + let mut selected: Option = None; + let name = menu.selected_item_name().to_string(); + for c in choices { + if c.1 == name { + selected = Some(c.0.to_owned()); + } + } + + match selected { + Some(s) => s.login(&client, user, hs.to_owned()).await, + None => anyhow::bail!("Invalid selection, aborting login") + } + } + }; + + let client = match client { + Ok(client) => client, + Err(e) => anyhow::bail!("{e}") + }; + let user_session = client .session() .expect("A logged-in client should have a session"); @@ -66,6 +175,7 @@ pub async fn login(data_dir: &Path, session_file: &Path, mxid: String) -> anyhow Ok(client) } + pub async fn build_client(data_dir: &Path, hs: String) -> anyhow::Result<(Client, ClientSession)> { let mut rng = thread_rng();