1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::atomic::AtomicBool};
use clap::Parser;
#[derive(Debug, Clone, Parser)]
pub struct Config {
/// tmux socket name
#[arg(short, long, default_value = "ssh")]
pub socket: String,
/// resolve hostnames using the system resolver
#[arg(short, long)]
pub resolve: bool,
/// include <NAME>. If <NAME> is a valid path or '-', hosts with be read from the file or
/// stdin, respectivly.
#[arg(short, long, id = "NAME")]
pub include: Vec<IncludeExclude>,
/// include <NAME>. If <NAME> is a valid path or '-', hosts with be read from the file or
/// stdin, respectivly.
#[arg(short, long, id = "HOST")]
pub exclude: Vec<IncludeExclude>,
}
impl Config {
pub fn included(&self) -> std::io::Result<HashSet<String>> {
Self::collect_values(&self.include)
}
pub fn excluded(&self) -> std::io::Result<HashSet<String>> {
Self::collect_values(&self.exclude)
}
fn collect_values(values: &[IncludeExclude]) -> std::io::Result<HashSet<String>> {
use std::io::BufRead;
values.iter().try_fold(HashSet::new(), |mut acc, item| {
match item {
IncludeExclude::Stdin => {
acc.extend(std::io::stdin().lock().lines().map_while(Result::ok))
}
IncludeExclude::File(filepath) => acc.extend(
std::io::BufReader::new(std::fs::File::open(filepath)?)
.lines()
.map_while(Result::ok),
),
IncludeExclude::Item(s) => {
acc.insert(s.to_owned());
}
}
Ok(acc)
})
}
}
static READ_STDIN: AtomicBool = AtomicBool::new(false);
#[derive(Debug, Clone)]
pub enum IncludeExclude {
Stdin,
Item(String),
File(PathBuf),
}
impl FromStr for IncludeExclude {
type Err = std::io::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
if s == "-" {
if READ_STDIN.load(std::sync::atomic::Ordering::Acquire) {
return Err(std::io::Error::new(
std::io::ErrorKind::Other,
"stdin argument used more than once",
));
}
READ_STDIN.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(Self::Stdin)
} else {
match PathBuf::from_str(s) {
Ok(path) if path.exists() => Ok(Self::File(path)),
_ => Ok(Self::Item(s.to_owned())),
}
}
}
}
|
