Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: custom ssh config parser #63

Merged
merged 12 commits into from
Feb 21, 2024
21 changes: 2 additions & 19 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ ratatui = "0.26.1"
regex = { version = "1.10.3", default-features = false, features = ["std"] }
shellexpand = "3.1.0"
shlex = "1.3.0"
ssh2-config = "0.2.3"
strum = "0.26.1"
strum_macros = "0.26.1"
tui-input = "0.8.0"
unicode-width = "0.1.11"
1 change: 1 addition & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod ssh;
pub mod ssh_config;
pub mod ui;

use clap::Parser;
Expand Down
67 changes: 38 additions & 29 deletions src/ssh.rs
Original file line number Diff line number Diff line change
@@ -1,30 +1,37 @@
use itertools::Itertools;
use regex::Regex;
use ssh2_config::{ParseRule, SshConfig};
use std::collections::VecDeque;
use std::error::Error;
use std::fs::File;
use std::{io::BufReader, process::Command};
use std::process::Command;

use crate::ssh_config::{self, HostVecExt};

#[derive(Debug, Clone)]
pub struct Host {
pub hostname: String,
pub aliases: String,
pub user: String,
pub user: Option<String>,
pub target: String,
pub port: String,
pub port: Option<String>,
}

/// # Errors
///
/// Will return `Err` if the SSH command cannot be executed.
pub fn connect(host: &Host) -> Result<(), Box<dyn Error>> {
Command::new("ssh")
.arg(format!("{}@{}", host.user, host.target))
.arg("-p")
.arg(&host.port)
.spawn()?
.wait()?;
let mut command = Command::new("ssh");

if let Some(user) = &host.user {
command.arg(format!("{}@{}", user, host.target));
} else {
command.arg(host.target.clone());
}

if let Some(port) = &host.port {
command.arg("-p").arg(port);
}

command.spawn()?.wait()?;

Ok(())
}
Expand All @@ -39,9 +46,9 @@ pub fn connect(host: &Host) -> Result<(), Box<dyn Error>> {
/// # Errors
///
/// Will return `Err` if the command cannot be executed.
///
///
/// # Panics
///
///
/// Will panic if the regex cannot be compiled.
pub fn run_with_pattern(pattern: &str, host: &Host) -> Result<(), Box<dyn Error>> {
let re = Regex::new(r"(?P<skip>%%)|(?P<h>%h)|(?P<u>%u)|(?P<p>%p)").unwrap();
Expand All @@ -51,9 +58,9 @@ pub fn run_with_pattern(pattern: &str, host: &Host) -> Result<(), Box<dyn Error>
} else if caps.name("h").is_some() {
host.hostname.clone()
} else if caps.name("u").is_some() {
host.user.clone()
host.user.clone().unwrap_or_default()
} else if caps.name("p").is_some() {
host.port.clone()
host.port.clone().unwrap_or_default()
} else {
String::new()
}
Expand All @@ -80,22 +87,24 @@ pub fn parse_config(raw_path: &String) -> Result<Vec<Host>, Box<dyn Error>> {
.ok_or("Failed to convert path to string")?
.to_string();

let mut reader = BufReader::new(File::open(path)?);
let config = SshConfig::default().parse(&mut reader, ParseRule::STRICT)?;

let hosts = config
.get_hosts()
let hosts = ssh_config::Parser::new()
.parse_file(path)?
.apply_patterns()
.merge_same_hosts()
.iter()
.filter(|host| host.params.host_name.is_some())
.filter(|host| host.get(&ssh_config::EntryType::Hostname).is_some())
.map(|host| Host {
hostname: host.pattern[0].pattern.clone(),
aliases: host.pattern[1..]
.iter()
.map(|p| p.pattern.clone())
.join(", "),
user: host.params.user.clone().unwrap_or_default(),
target: host.params.host_name.clone().unwrap_or_default(),
port: host.params.port.unwrap_or(22).to_string(),
hostname: host
.get_patterns()
.first()
.unwrap_or(&String::new())
.clone(),
aliases: host.get_patterns().iter().skip(1).join(", "),
user: host.get(&ssh_config::EntryType::User),
target: host
.get(&ssh_config::EntryType::Hostname)
.unwrap_or_default(),
port: host.get(&ssh_config::EntryType::Port),
})
.collect();

Expand Down
223 changes: 223 additions & 0 deletions src/ssh_config/host.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,223 @@
use regex::Regex;
use std::collections::HashMap;

use super::EntryType;

pub(crate) type Entry = (EntryType, String);

#[derive(Debug, Clone)]
pub struct Host {
patterns: Vec<String>,
entries: HashMap<EntryType, String>,
}

impl Host {
#[must_use]
pub fn new(patterns: Vec<String>) -> Host {
Host {
patterns,
entries: HashMap::new(),
}
}

pub fn update(&mut self, entry: Entry) {
self.entries.insert(entry.0, entry.1);
}

pub(crate) fn extend_patterns(&mut self, host: &Host) {
self.patterns.extend(host.patterns.clone());
}

pub(crate) fn extend_entries(&mut self, host: &Host) {
self.entries.extend(host.entries.clone());
}

pub(crate) fn extend_if_not_contained(&mut self, host: &Host) {
for (key, value) in &host.entries {
if !self.entries.contains_key(key) {
self.entries.insert(key.clone(), value.clone());
}
}
}

#[allow(clippy::must_use_candidate)]
pub fn get_patterns(&self) -> &Vec<String> {
&self.patterns
}

/// # Panics
///
/// Will panic if the regex cannot be compiled.
#[allow(clippy::must_use_candidate)]
pub fn matching_pattern_regexes(&self) -> Vec<(Regex, bool)> {
if self.patterns.is_empty() {
return Vec::new();
}

self.patterns
.iter()
.filter_map(|pattern| {
let contains_wildcard =
pattern.contains('*') || pattern.contains('?') || pattern.contains('!');
if !contains_wildcard {
return None;
}

let mut pattern = pattern
.replace('.', r"\.")
.replace('*', ".*")
.replace('?', ".");

let is_negated = pattern.starts_with('!');
if is_negated {
pattern.remove(0);
}

pattern = format!("^{pattern}$");
Some((Regex::new(&pattern).unwrap(), is_negated))
})
.collect()
}

#[allow(clippy::must_use_candidate)]
pub fn get(&self, entry: &EntryType) -> Option<String> {
self.entries.get(entry).cloned()
}

#[allow(clippy::must_use_candidate)]
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}

#[allow(clippy::module_name_repetitions)]
pub trait HostVecExt {
/// Merges the hosts with the same entries into one host.
fn merge_same_hosts(&mut self) -> &mut Self;

/// Spreads the hosts with multiple patterns into multiple hosts with one pattern.
fn spread(&mut self) -> &mut Self;

/// Apply patterns entries to non-pattern hosts and remove the pattern hosts.
fn apply_patterns(&mut self) -> &mut Self;
}

impl HostVecExt for Vec<Host> {
fn merge_same_hosts(&mut self) -> &mut Self {
for i in (0..self.len()).rev() {
for j in (0..i).rev() {
if self[i].entries != self[j].entries {
continue;
}

let host = self[i].clone();
self[j].extend_patterns(&host);
self[j].extend_entries(&host);
self.remove(i);
break;
}
}

self
}

fn spread(&mut self) -> &mut Self {
let mut hosts = Vec::new();

for host in self.iter_mut() {
let patterns = host.get_patterns();
if patterns.is_empty() {
hosts.push(host.clone());
continue;
}

for pattern in patterns {
let mut new_host = host.clone();
new_host.patterns = vec![pattern.clone()];
hosts.push(new_host);
}
}

self
}

/// Apply patterns entries to non-pattern hosts and remove the pattern hosts.
///
/// You might want to call [`HostVecExt::merge_same_hosts`] after this.
fn apply_patterns(&mut self) -> &mut Self {
let hosts = self.spread();
let mut pattern_indexes = Vec::new();

for i in 0..hosts.len() {
let matching_pattern_regexes = hosts[i].matching_pattern_regexes();
if matching_pattern_regexes.is_empty() {
continue;
}

pattern_indexes.push(i);

for j in (i + 1)..hosts.len() {
if !hosts[j].matching_pattern_regexes().is_empty() {
continue;
}

for (regex, is_negated) in &matching_pattern_regexes {
if regex.is_match(&hosts[j].patterns[0]) == *is_negated {
continue;
}

let host = hosts[i].clone();
hosts[j].extend_if_not_contained(&host);
break;
}
}
}

for i in pattern_indexes.into_iter().rev() {
hosts.remove(i);
}

hosts
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_apply_patterns() {
let mut hosts = Vec::new();

let mut host = Host::new(vec!["*".to_string()]);
host.update((EntryType::Hostname, "example.com".to_string()));
hosts.push(host);

let mut host = Host::new(vec!["!example.com".to_string()]);
host.update((EntryType::User, "hello".to_string()));
hosts.push(host);

let mut host = Host::new(vec!["example.com".to_string()]);
host.update((EntryType::Port, "22".to_string()));
hosts.push(host);

let mut host = Host::new(vec!["hello.com".to_string()]);
host.update((EntryType::Port, "22".to_string()));
hosts.push(host);

let hosts = hosts.apply_patterns();

assert_eq!(hosts.len(), 2);

assert_eq!(hosts[0].patterns[0], "example.com");
assert_eq!(hosts[0].entries.len(), 2);
assert_eq!(hosts[0].entries[&EntryType::Hostname], "example.com");
assert_eq!(hosts[0].entries[&EntryType::Port], "22");

assert_eq!(hosts[1].patterns[0], "hello.com");
assert_eq!(hosts[1].entries.len(), 3);
assert_eq!(hosts[1].entries[&EntryType::Hostname], "example.com");
assert_eq!(hosts[1].entries[&EntryType::User], "hello");
assert_eq!(hosts[1].entries[&EntryType::Port], "22");
}
}
Loading