diff --git a/Cargo.lock b/Cargo.lock index 934e637..e5aa16f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -540,6 +540,27 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dirs-next" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b98cf8ebf19c3d1b223e151f99a4f9f0690dca41414773390fc824184ac833e1" +dependencies = [ + "cfg-if", + "dirs-sys-next", +] + +[[package]] +name = "dirs-sys-next" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" +dependencies = [ + "libc", + "redox_users", + "winapi", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -1204,6 +1225,16 @@ dependencies = [ "windows-targets 0.53.3", ] +[[package]] +name = "libredox" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "391290121bad3d37fbddad76d8f5d1c1c314cfc646d143d7e07a3086ddff0ce3" +dependencies = [ + "bitflags", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -1850,6 +1881,17 @@ dependencies = [ "bitflags", ] +[[package]] +name = "redox_users" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" +dependencies = [ + "getrandom 0.2.16", + "libredox", + "thiserror 1.0.69", +] + [[package]] name = "regex" version = "1.11.1" @@ -2183,6 +2225,7 @@ dependencies = [ "bytes", "chrono", "crossterm 0.27.0", + "dirs-next", "flate2", "futures", "futures-util", diff --git a/Cargo.toml b/Cargo.toml index fa97049..c3aad7a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ axum = { version = "0.7", features = ["ws"] } prost = "0.13" prost-types = "0.13" bytes = "1" +dirs-next = "2" [profile.release] # Favor smaller, simpler binaries with good runtime perf diff --git a/socktop/Cargo.toml b/socktop/Cargo.toml index 526001a..bdbcb92 100644 --- a/socktop/Cargo.toml +++ b/socktop/Cargo.toml @@ -19,6 +19,7 @@ crossterm = { workspace = true } chrono = { workspace = true } anyhow = { workspace = true } flate2 = { version = "1", default-features = false, features = ["rust_backend"] } +dirs-next = { workspace = true } rustls = "0.23" rustls-pemfile = "2.1" prost = { workspace = true } diff --git a/socktop/src/main.rs b/socktop/src/main.rs index e1caaa7..e2a6f90 100644 --- a/socktop/src/main.rs +++ b/socktop/src/main.rs @@ -2,29 +2,47 @@ mod app; mod history; +mod profiles; mod types; mod ui; mod ws; use app::App; +use profiles::{load_profiles, save_profiles, ProfileEntry, ProfileRequest, ResolveProfile}; use std::env; +use std::io::{self, Write}; -fn parse_args>(args: I) -> Result<(String, Option), String> { +struct ParsedArgs { + url: Option, + tls_ca: Option, + profile: Option, + save: bool, +} + +fn parse_args>(args: I) -> Result { let mut it = args.into_iter(); let prog = it.next().unwrap_or_else(|| "socktop".into()); let mut url: Option = None; let mut tls_ca: Option = None; + let mut profile: Option = None; + let mut save = false; // --save-profile while let Some(arg) = it.next() { match arg.as_str() { "-h" | "--help" => { return Err(format!( - "Usage: {prog} [--tls-ca CERT_PEM|-t CERT_PEM] ws://HOST:PORT/ws" + "Usage: {prog} [--tls-ca CERT_PEM|-t CERT_PEM] [--profile NAME|-P NAME] [--save] [ws://HOST:PORT/ws]" )); } "--tls-ca" | "-t" => { tls_ca = it.next(); } + "--profile" | "-P" => { + profile = it.next(); + } + "--save" => { + save = true; + } _ if arg.starts_with("--tls-ca=") => { if let Some((_, v)) = arg.split_once('=') { if !v.is_empty() { @@ -32,30 +50,36 @@ fn parse_args>(args: I) -> Result<(String, Option } } } + _ if arg.starts_with("--profile=") => { + if let Some((_, v)) = arg.split_once('=') { + if !v.is_empty() { + profile = Some(v.to_string()); + } + } + } _ => { if url.is_none() { url = Some(arg); } else { return Err(format!( - "Unexpected argument. Usage: {prog} [--tls-ca CERT_PEM|-t CERT_PEM] ws://HOST:PORT/ws" + "Unexpected argument. Usage: {prog} [--tls-ca CERT_PEM|-t CERT_PEM] [--profile NAME|-P NAME] [--save] [ws://HOST:PORT/ws]" )); } } } } - - match url { - Some(u) => Ok((u, tls_ca)), - None => Err(format!( - "Usage: {prog} [--tls-ca CERT_PEM|-t CERT_PEM] ws://HOST:PORT/ws" - )), - } + Ok(ParsedArgs { + url, + tls_ca, + profile, + save, + }) } #[tokio::main] async fn main() -> Result<(), Box> { // Reuse the same parsing logic for testability - let (url, tls_ca) = match parse_args(env::args()) { + let parsed = match parse_args(env::args()) { Ok(v) => v, Err(msg) => { eprintln!("{msg}"); @@ -63,6 +87,93 @@ async fn main() -> Result<(), Box> { } }; + let profiles_file = load_profiles(); + let req = ProfileRequest { + profile_name: parsed.profile.clone(), + url: parsed.url.clone(), + tls_ca: parsed.tls_ca.clone(), + }; + let resolved = req.resolve(&profiles_file); + + // Determine final connection parameters (and maybe mutated profiles to persist) + let mut profiles_mut = profiles_file.clone(); + let (url, tls_ca): (String, Option) = match resolved { + ResolveProfile::Direct(u, t) => { + // Possibly save if profile specified and --save or new entry + if let Some(name) = parsed.profile.as_ref() { + let existing = profiles_mut.profiles.get(name); + if existing.is_none() || parsed.save { + // If existing and not forced save, prompt + if existing.is_some() && !parsed.save { + if prompt_yes_no(&format!("Overwrite existing profile '{name}'? [y/N]: ")) { + profiles_mut.profiles.insert( + name.clone(), + ProfileEntry { + url: u.clone(), + tls_ca: t.clone(), + }, + ); + let _ = save_profiles(&profiles_mut); + } + } else { + profiles_mut.profiles.insert( + name.clone(), + ProfileEntry { + url: u.clone(), + tls_ca: t.clone(), + }, + ); + let _ = save_profiles(&profiles_mut); + } + } + } + (u, t) + } + ResolveProfile::Loaded(u, t) => (u, t), + ResolveProfile::PromptSelect(names) => { + eprintln!("Select profile:"); + for (i, n) in names.iter().enumerate() { + eprintln!(" {}. {}", i + 1, n); + } + eprint!("Enter number (or blank to abort): "); + let _ = io::stderr().flush(); + let mut line = String::new(); + if io::stdin().read_line(&mut line).is_ok() { + if let Ok(idx) = line.trim().parse::() { + if idx >= 1 && idx <= names.len() { + let name = &names[idx - 1]; + if let Some(entry) = profiles_mut.profiles.get(name) { + (entry.url.clone(), entry.tls_ca.clone()) + } else { + return Ok(()); + } + } else { + return Ok(()); + } + } else { + return Ok(()); + } + } else { + return Ok(()); + } + } + ResolveProfile::None => { + eprintln!("No URL provided and no profiles to select."); + return Ok(()); + } + }; + let mut app = App::new(); app.run(&url, tls_ca.as_deref()).await } + +fn prompt_yes_no(prompt: &str) -> bool { + eprint!("{prompt}"); + let _ = io::stderr().flush(); + let mut line = String::new(); + if io::stdin().read_line(&mut line).is_ok() { + matches!(line.trim().to_ascii_lowercase().as_str(), "y" | "yes") + } else { + false + } +} diff --git a/socktop/src/profiles.rs b/socktop/src/profiles.rs new file mode 100644 index 0000000..2a75c4d --- /dev/null +++ b/socktop/src/profiles.rs @@ -0,0 +1,96 @@ +//! Connection profiles: load/save simple JSON mapping of profile name -> { url, tls_ca } +//! Stored under XDG config dir: $XDG_CONFIG_HOME/socktop/profiles.json (fallback ~/.config/socktop/profiles.json) + +use serde::{Deserialize, Serialize}; +use std::{collections::BTreeMap, fs, path::PathBuf}; + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ProfileEntry { + pub url: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub tls_ca: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ProfilesFile { + #[serde(default)] + pub profiles: BTreeMap, + #[serde(default)] + pub version: u32, +} + +pub fn config_dir() -> PathBuf { + if let Some(xdg) = std::env::var_os("XDG_CONFIG_HOME") { + PathBuf::from(xdg).join("socktop") + } else { + dirs_next::config_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("socktop") + } +} + +pub fn profiles_path() -> PathBuf { + config_dir().join("profiles.json") +} + +pub fn load_profiles() -> ProfilesFile { + let path = profiles_path(); + match fs::read_to_string(&path) { + Ok(s) => serde_json::from_str(&s).unwrap_or_default(), + Err(_) => ProfilesFile::default(), + } +} + +pub fn save_profiles(p: &ProfilesFile) -> std::io::Result<()> { + let path = profiles_path(); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let data = serde_json::to_vec_pretty(p).expect("serialize profiles"); + fs::write(path, data) +} + +pub enum ResolveProfile { + /// Use the provided runtime inputs (not persisted). (url, tls_ca) + Direct(String, Option), + /// Loaded from existing profile entry (url, tls_ca) + Loaded(String, Option), + /// Should prompt user to select among profile names + PromptSelect(Vec), + /// No profile could be resolved (e.g., missing arguments) + None, +} + +pub struct ProfileRequest { + pub profile_name: Option, + pub url: Option, + pub tls_ca: Option, +} + +impl ProfileRequest { + pub fn resolve(self, pf: &ProfilesFile) -> ResolveProfile { + // Case: only profile name given -> try load + if self.url.is_none() && self.profile_name.is_some() { + let name = self.profile_name.unwrap(); + if let Some(entry) = pf.profiles.get(&name) { + return ResolveProfile::Loaded(entry.url.clone(), entry.tls_ca.clone()); + } else { + return ResolveProfile::None; + } + } + // Both provided -> direct (maybe later saved by caller) + if let Some(u) = self.url { + return ResolveProfile::Direct(u, self.tls_ca); + } + // Nothing provided -> maybe prompt select if profiles exist + if self.url.is_none() && self.profile_name.is_none() { + if pf.profiles.is_empty() { + ResolveProfile::None + } else { + ResolveProfile::PromptSelect(pf.profiles.keys().cloned().collect()) + } + } else { + ResolveProfile::None + } + } +} diff --git a/socktop/tests/cli_args.rs b/socktop/tests/cli_args.rs index 53433e1..6feacb4 100644 --- a/socktop/tests/cli_args.rs +++ b/socktop/tests/cli_args.rs @@ -17,8 +17,8 @@ fn test_help_mentions_short_and_long_flags() { String::from_utf8_lossy(&output.stderr) ); assert!( - text.contains("--tls-ca") && text.contains("-t"), - "help text missing --tls-ca/-t\n{text}" + text.contains("--tls-ca") && text.contains("-t") && text.contains("--profile") && text.contains("-P"), + "help text missing expected flags (--tls-ca/-t, --profile/-P)\n{text}" ); } @@ -53,4 +53,17 @@ fn test_tlc_ca_arg_long_and_short_parsed() { String::from_utf8_lossy(&out2.stderr) ); assert!(text2.contains("Usage:")); + + // Profile flags with help (should not error) + let out3 = Command::new(exe) + .args(["--profile", "dev", "--help"]) + .output() + .expect("run socktop"); + assert!(out3.status.success(), "socktop --profile dev --help did not succeed"); + let text3 = format!( + "{}{}", + String::from_utf8_lossy(&out3.stdout), + String::from_utf8_lossy(&out3.stderr) + ); + assert!(text3.contains("Usage:")); }