diff --git a/Cargo.lock b/Cargo.lock index d2bb106..6272896 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1182,6 +1182,12 @@ dependencies = [ "serde", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf-8" version = "0.7.6" @@ -1289,6 +1295,8 @@ dependencies = [ "tokio-tungstenite", "tracing", "tracing-subscriber", + "url", + "urlencoding", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 68b2051..0112038 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,3 +11,5 @@ tracing-subscriber = { version = "0.3", features = ["fmt"] } clap = { version = "4", features = ["derive"] } futures-util = "0.3" chrono = "0.4" +urlencoding = "2" +url = "2" diff --git a/src/main.rs b/src/main.rs index 0d3c8c1..36c37ae 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use chrono::Local; use clap::Parser; +use url::Url; use futures_util::StreamExt; use tokio::task::JoinSet; use tokio_tungstenite::{ @@ -30,6 +31,10 @@ struct Args { /// Enable debug logging (shows request/response headers) #[arg(long)] debug: bool, + + /// Query string parameters to add to all URLs (pre-encoded, e.g., "name=First%20Last&key=value") + #[arg(short = 'q', long = "query-string-all")] + query_string_all: Option, } #[tokio::main] @@ -53,9 +58,46 @@ async fn main() -> Result<(), Box> { fs::create_dir_all(&session_dir)?; info!("Created session directory: {}", session_dir.display()); + // Process URLs and add query string parameters if specified + let mut processed_urls = Vec::new(); + for url_str in &args.urls { + let mut url = Url::parse(url_str)?; + + if let Some(ref extra_params) = args.query_string_all { + // Parse existing query string and extra params + let existing: Vec<(String, String)> = url + .query_pairs() + .map(|(k, v)| (k.into_owned(), v.into_owned())) + .collect(); + + let extra: Vec<(&str, &str)> = extra_params + .split('&') + .filter(|s| !s.is_empty()) + .filter_map(|pair| pair.split_once('=')) + .collect(); + + // Rebuild query string with existing + extra params + { + let mut query_pairs = url.query_pairs_mut(); + query_pairs.clear(); + for (k, v) in &existing { + query_pairs.append_pair(k, v); + } + for (k, v) in &extra { + // Decode first since append_pair will re-encode + let k_decoded = urlencoding::decode(k).unwrap_or_else(|_| (*k).into()); + let v_decoded = urlencoding::decode(v).unwrap_or_else(|_| (*v).into()); + query_pairs.append_pair(&k_decoded, &v_decoded); + } + } + } + + processed_urls.push(url.to_string()); + } + // Connect to all URLs simultaneously let mut connect_futures = Vec::new(); - for (idx, url) in args.urls.iter().enumerate() { + for (idx, url) in processed_urls.iter().enumerate() { let letter = (b'A' + idx as u8) as char; let bearer_token = args.bearer_token.clone(); let url = url.clone();