From 3c1dd898d5fe9483a9e020f87241c10918435f6c Mon Sep 17 00:00:00 2001 From: 0xf8 <0xf8.dev@proton.me> Date: Wed, 31 May 2023 23:10:10 -0400 Subject: [PATCH] separate ssh platform code into ssh.rs | extend ssh config | improve tui --- src/config.rs | 76 +++++++++++++-------- src/input.rs | 31 +++++++++ src/main.rs | 9 +-- src/platform.rs | 173 ++---------------------------------------------- src/ssh.rs | 166 ++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 256 insertions(+), 199 deletions(-) create mode 100644 src/ssh.rs diff --git a/src/config.rs b/src/config.rs index 8a15ac6..870eee4 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use colored::Colorize; use crate::{ platform::Platform, input }; use flexbuffers::{ FlexbufferSerializer, Reader }; use std::collections::BTreeSet; @@ -19,27 +20,31 @@ pub struct Host { impl std::fmt::Display for Host { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_fmt(std::format_args!("{} via {}@{}", self.platform, self.user, self.host)) + f.write_fmt(std::format_args!("[{}] {}@{}", self.platform, self.user, self.host)) } } impl Host { - pub fn edit(&mut self) -> Result { - let mut data = toml::to_string_pretty(self)?; - - { - let t = mktemp::Temp::new_file().unwrap(); - fs::write(t.to_owned(), data)?; + pub fn edit(&mut self, config: &mut ConfigManager) -> Result { + let data = toml::to_string_pretty(self)?; + let t = input::edit_temp_from_string(&data)?; - let editor = std::env::var("EDITOR").expect("EDITOR is not set"); - std::process::Command::new(editor) - .arg(t.to_owned().to_string_lossy().to_string()) - .spawn()?.wait()?; + match toml::from_str::(&fs::read_to_string(t)?) { + Ok(d) => { + println!(" ## {} ## ", "Definition OK".green()); - data = fs::read_to_string(t)?; + // Replace old config with new config + config.configs.remove(self); + config.configs.insert(d.to_owned()); + config.save(); + + Ok(d) + }, + Err(e) => { + println!(" ## {}: {} ## ", "Rejecting definition".red(), e); + Ok(self.to_owned()) + }, } - - Ok(toml::from_str(&data)?) } } @@ -53,36 +58,53 @@ pub struct ConfigManager { pub configs: BTreeSet, } + impl ConfigManager { pub fn new(given_config_root: Option, search_path: Option) -> Result { + // Resolve config root let config_root = match given_config_root { Some(c) => path::Path::new(&c).to_owned(), None => dirs::data_local_dir().unwrap().join("keyman/"), }; + // Make config root if it doesn't exist (eg: first run) if !fs::try_exists(config_root.to_owned()).unwrap_or(false) { fs::create_dir_all(config_root.to_owned())?; } + // Resolve search path let search_path = match search_path { Some(s) => std::path::Path::new(&s).to_owned(), None => dirs::home_dir().unwrap().join(".ssh/"), }; - match fs::read(config_root.join("config")) { - Ok(c) => { - let root = Reader::get_root(c.as_slice())?; - match ConfigManager::deserialize(root) { - Ok(mut mgr) => { - mgr.config_dir = config_root; - mgr.search_path = search_path; - Ok(mgr) - }, - Err(_) => Ok(Self { config_dir: config_root, search_path, configs: BTreeSet::new() }), - } - } - Err(_) => Ok(Self { config_dir: config_root, search_path, configs: BTreeSet::new() }), + // Read config + let config = fs::read(config_root.join("config")); + if config.is_err() { + return Ok(Self { config_dir: config_root.to_owned(), search_path, configs: BTreeSet::new() }) } + let config = config.unwrap(); + + // Deserialize + let root = Reader::get_root(config.as_slice())?; + let mgr = ConfigManager::deserialize(root); + + if mgr.is_err() { + return Ok(Self { config_dir: config_root, search_path, configs: BTreeSet::new() }) + } + + let mgr = mgr.unwrap(); + Ok(Self { config_dir: config_root, search_path, configs: mgr.configs.to_owned() }) + } + + pub fn remove(&mut self, host: &Host) { + self.configs.remove(host); + self.save(); + } + + pub fn insert(&mut self, host: Host) { + self.configs.insert(host); + self.save(); } pub fn save(&mut self) { diff --git a/src/input.rs b/src/input.rs index 866f4e0..d0c55c7 100644 --- a/src/input.rs +++ b/src/input.rs @@ -32,6 +32,7 @@ pub fn get_host(config: &mut ConfigManager) -> Result { }) } + pub fn get_platform() -> Result { let platforms = Platform::all(); let c = Select::with_theme(&ColorfulTheme::default()) @@ -43,6 +44,7 @@ pub fn get_platform() -> Result { Ok(platforms.get(c).unwrap().to_owned()) } + pub fn get_key(config: &ConfigManager, platform: Platform) -> Result { let keys: Vec = platform.get_keys(config); let c = FuzzySelect::with_theme(&ColorfulTheme::default()) @@ -58,6 +60,7 @@ pub fn get_key(config: &ConfigManager, platform: Platform) -> Result Result { let theme = ColorfulTheme::default(); let mut input = Input::with_theme(&theme); @@ -81,4 +84,32 @@ pub fn get_custom_key() -> Result { println!("File doesn't exist") } }) +} + + +fn edit_temp(temp: &mktemp::Temp) -> Result<()> { + let editor = std::env::var("EDITOR").expect("EDITOR is not set"); + std::process::Command::new(editor) + .arg(&temp.to_str().unwrap()) + .spawn()?.wait()?; + + Ok(()) +} + + +pub fn edit_temp_from_string(data: &String) -> Result { + let t = mktemp::Temp::new_file()?; + fs::write(&t, data)?; + edit_temp(&t)?; + + Ok(t) +} + + +pub fn edit_temp_from_file(file: &String) -> Result { + let t = mktemp::Temp::new_file()?; + fs::copy(file, &t)?; + edit_temp(&t)?; + + Ok(t) } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 7268668..604f97c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,12 +3,12 @@ use anyhow::Result; use config::ConfigManager; use std::io; -use termion::screen::IntoAlternateScreen; use tui::{ backend::TermionBackend, Terminal }; pub mod config; pub mod input; pub mod platform; +pub mod ssh; fn main() -> Result<()> { let mut search_path: Option = None; @@ -31,16 +31,17 @@ fn main() -> Result<()> { let mut config: ConfigManager = ConfigManager::new(config_root, search_path)?; - let backend = TermionBackend::new(io::stdout().into_alternate_screen()?); + let backend = TermionBackend::new(io::stdout()); // let backend = TermionBackend::new(io::stdout()); // used for debugging let mut term = Terminal::new(backend)?; - + term.hide_cursor()?; + loop { term.clear()?; let mut host = input::get_host(&mut config)?; let platform = host.platform.to_owned(); - platform.run(&mut host, &mut config)?; + platform.run(&mut host, &mut config, &mut term)?; } } diff --git a/src/platform.rs b/src/platform.rs index 6bcc6f2..80549c8 100644 --- a/src/platform.rs +++ b/src/platform.rs @@ -1,11 +1,8 @@ use anyhow::Result; -use colored::Colorize; -use crate::{ ConfigManager, config::Host, input }; -use dialoguer::{ Select, Input, Confirm, Editor, theme::ColorfulTheme }; +use crate::{ ConfigManager, config::Host }; use serde::{ Serialize, Deserialize }; -use ssh2_config::{ SshConfig, ParseRule }; -use ssh_key::PrivateKey; -use std::fs; +use std::io; +use tui::{ backend::TermionBackend, Terminal }; #[derive(Clone, Debug)] @@ -53,172 +50,12 @@ impl Platform { Platform::SSH => self.ssh_get_keys(config), } } - pub fn run(&self, host: &mut Host, config: &mut ConfigManager) -> Result<()> { + pub fn run(&self, host: &mut Host, config: &mut ConfigManager, term: &mut Terminal>) -> Result<()> { match self { - Platform::SSH => self.ssh_run(host, config), + Platform::SSH => self.ssh_run(host, config, term), } } - - // --- --- --- - // SSH - // --- --- --- - - fn ssh_new_config(&self, host: String, port: i16, keyid: PlatformKey) -> String { - format!(r#"# Generated by Keyman - -Host {host} - HostName {host} - Port {port} - IdentityFile {} -"#, keyid.file) - } - - // New host - fn ssh_new_host(&self, config: &mut ConfigManager) -> Result { - println!(" ## Creating new SSH host entry ## "); - let keyid = input::get_key(config, self.to_owned())?; - - let userhost = loop { - let inp: String = Input::with_theme(&ColorfulTheme::default()) - .with_prompt("User@Host") - .interact_text()?; - - let split = inp.split("@"); - let split: Vec = split.map(|a| a.to_string()).collect(); - - if split.len() < 2 { - println!("{} Example: paul@example.com:34 - @[:port]", "Incorrect user@host given.".red()); - } else { - break split - } - }; - - let user: String = userhost.get(0).unwrap().to_string(); - let mut host: String = userhost.get(1).unwrap().to_string(); - let mut port: i16 = 22; - - let _host = host.to_owned(); - let h: Vec<&str> = _host.split(":").collect(); - - if h.len() > 1 { - let portstr = h.get(1).unwrap(); - port = i16::from_str_radix(portstr, 10)?; - host = h.get(0).unwrap().to_string(); - } - - let config_dir = config.config_dir.join("hosts/").join(host.to_owned()); - fs::create_dir_all(config_dir.to_owned()).expect(&format!("Couldnt create_dir_all on {:?}", config_dir.to_owned())); - let config_file = config_dir.join("config"); - - fs::write(config_file.to_owned(), self.ssh_new_config(host.to_owned(), port, keyid.to_owned())).expect(&format!("Couldn't write config at {:?}", config_file.to_owned())); - - Ok(Host { - platform: Platform::SSH, - id: keyid.file, - user, - host, - port, - config: config_file.to_owned().to_string_lossy().to_string() - }) - } - - // Get keys - fn ssh_get_keys(&self, config: &ConfigManager) -> Vec { - let mut keys = vec![]; - let Ok(dir) = std::fs::read_dir(config.search_path.to_owned()) else { panic!("Couldn't read {:?}", config.search_path.to_owned()) }; - - for file in dir { - let Ok(f) = file else { continue }; - - match PrivateKey::read_openssh_file(&f.path()) { - Ok(p) => keys.push(PlatformKey { - file: format!("{}", f.path().to_string_lossy()), - display: format!("{}: {}", p.fingerprint(ssh_key::HashAlg::Sha256), f.file_name().to_string_lossy()) - }), - // Not a private key - Err(_) => continue - }; - } - - keys - } - - // Run - fn ssh_run(&self, host: &mut Host, config: &mut ConfigManager) -> Result<()> { - Ok(loop { - let action = Select::with_theme(&ColorfulTheme::default()) - .with_prompt(format!("{}@{}", host.user.to_owned(), host.host.to_owned())) - .item("Connect") - .item("Edit config") - .item("Edit definition") - .item("Delete") - .item("Back") - .default(0) - .interact()?; - - match action { - 0 => { // Connect - // TODO: clear screen before connect - std::process::Command::new("ssh") - .arg("-F") - .arg(host.config.to_owned()) - .arg("-l") - .arg(host.user.to_owned()) - .arg(host.host.to_owned()).spawn()?.wait()?; - } - 1 => { // Edit config - // TODO: replace this with configuer::editor -// if let Some(cnf) = Editor::new().edit(fs::read_to_string(config.config_dir.join("config"))?.as_str())? { -// fs::write(config.config_dir.join("config"), cnf)?; -// } - let t = mktemp::Temp::new_file()?; - fs::copy(&host.config, &t)?; - - let editor = std::env::var("EDITOR").expect("EDITOR is not set"); - std::process::Command::new(editor) - .arg(&t.to_str().unwrap()) - .spawn()?.wait()?; - - let mut file = std::io::BufReader::new(fs::File::open(&t).expect(&format!("Couldn't open config ({})", &host.config))); - match SshConfig::default().parse(&mut file, ParseRule::STRICT) { - Ok(_) => { - println!(" ## {} ## ", "Config OK".green()); - fs::copy(&t, &host.config)?; - }, - Err(e) => println!(" ## {}: {} ## ", "Rejecting config".red(), e), - } - } - 2 => { // Edit definition - let edited_host = host.edit()?; - config.configs.remove(host); - config.configs.insert(edited_host.to_owned()); - host.clone_from(&edited_host); - config.save(); - } - 3 => { // Delete - let confirm = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Are you sure you want to delete this definition?") - .default(false) - .interact()?; - - if confirm { - let mut conf = std::path::PathBuf::from(host.config.to_owned()); - conf.pop(); - - config.configs.remove(host); - config.save(); - fs::remove_dir_all(conf)?; - break; - } - } // Back - 4 => break, - _ => () - } - }) - } - - // --- --- --- // GPG // --- --- --- diff --git a/src/ssh.rs b/src/ssh.rs new file mode 100644 index 0000000..473349c --- /dev/null +++ b/src/ssh.rs @@ -0,0 +1,166 @@ +use anyhow::Result; +use colored::Colorize; +use crate::{ ConfigManager, config::Host, input, platform::{ Platform, PlatformKey } }; +use dialoguer::{ Select, Input, Confirm, theme::ColorfulTheme }; +use ssh2_config::{ SshConfig, ParseRule }; +use ssh_key::PrivateKey; +use std::fs; +use std::io; +use tui::{ backend::TermionBackend, Terminal }; + + +fn new_config(host: String, user: String, port: i16, keyid: PlatformKey) -> String { + format!(r#"# Generated by Keyman + +Host {host} + HostName {host} + User {user} + Port {port} + IdentityFile {} +"#, keyid.file) +} + +impl Platform { + // New host + pub fn ssh_new_host(&self, config: &mut ConfigManager) -> Result { + println!(" ## Creating new SSH host entry ## "); + let keyid = input::get_key(config, self.to_owned())?; + + let userhost = loop { + let inp: String = Input::with_theme(&ColorfulTheme::default()) + .with_prompt("User@Host") + .interact_text()?; + + let split = inp.split("@"); + let split: Vec = split.map(|a| a.to_string()).collect(); + + if split.len() < 2 { + println!("{} Example: paul@example.com:34 - @[:port]", "Incorrect user@host given.".red()); + } else { + break split + } + }; + + let user: String = userhost.get(0).unwrap().to_string(); + let mut host: String = userhost.get(1).unwrap().to_string(); + let mut port: i16 = 22; + + let _host = host.to_owned(); + let h: Vec<&str> = _host.split(":").collect(); + + if h.len() > 1 { + let portstr = h.get(1).unwrap(); + port = i16::from_str_radix(portstr, 10)?; + host = h.get(0).unwrap().to_string(); + } + + let config_dir = config.config_dir.join("hosts/").join(host.to_owned()); + fs::create_dir_all(config_dir.to_owned()).expect(&format!("Couldnt create_dir_all on {:?}", config_dir.to_owned())); + let config_file = config_dir.join("config"); + + fs::write(config_file.to_owned(), new_config(host.to_owned(), user.to_owned(), port, keyid.to_owned())).expect(&format!("Couldn't write config at {:?}", config_file.to_owned())); + + Ok(Host { + platform: Platform::SSH, + id: keyid.file, + user, + host, + port, + config: config_file.to_owned().to_string_lossy().to_string() + }) + } + + // Get keys + pub fn ssh_get_keys(&self, config: &ConfigManager) -> Vec { + let mut keys = vec![]; + let Ok(dir) = std::fs::read_dir(config.search_path.to_owned()) else { panic!("Couldn't read {:?}", config.search_path.to_owned()) }; + + for file in dir { + let Ok(f) = file else { continue }; + + match PrivateKey::read_openssh_file(&f.path()) { + Ok(p) => keys.push(PlatformKey { + file: format!("{}", f.path().to_string_lossy()), + display: format!("{}: {}", p.fingerprint(ssh_key::HashAlg::Sha256), f.file_name().to_string_lossy()) + }), + // Not a private key + Err(_) => continue + }; + } + + keys + } + + // Run + pub fn ssh_run(&self, host: &mut Host, config: &mut ConfigManager, term: &mut Terminal>) -> Result<()> { + Ok(loop { + let action = Select::with_theme(&ColorfulTheme::default()) + .with_prompt(format!("{}@{}", host.user.to_owned(), host.host.to_owned())) + .item("Connect") + .item("Edit config") + .item("Edit definition") + .item("Delete") + .item("Back") + .default(0) + .interact()?; + + term.clear()?; + match action { + + // Connect + 0 => { + // TODO: clear screen before connect + std::process::Command::new("ssh") + .arg("-F") + .arg(host.config.to_owned()) + .arg(host.host.to_owned()).spawn()?.wait()?; + } + + + // Edit config + 1 => { + let t = input::edit_temp_from_file(&host.config)?; + + // Validate SSH config + let mut file = std::io::BufReader::new(fs::File::open(&t)?); + if let Err(e) = SshConfig::default().parse(&mut file, ParseRule::STRICT) { + println!(" ## {}: {} ## ", "Rejecting config".red(), e) + } else { + println!(" ## {} ## ", "Config OK".green()); + fs::copy(&t, &host.config)?; + } + } + + + // Edit definition + 2 => { + let edited_host = host.edit(config)?; + host.clone_from(&edited_host); + } + + + + // Delete + 3 => { + let confirm = Confirm::with_theme(&ColorfulTheme::default()) + .with_prompt("Are you sure you want to delete this definition?") + .default(false) + .interact()?; + + if confirm { + let mut conf = std::path::PathBuf::from(host.config.to_owned()); + conf.pop(); + + config.remove(host); + fs::remove_dir_all(conf)?; + break; + } + } + + + // Back + _ => break + } + }) + } +} \ No newline at end of file