summaryrefslogtreecommitdiffstats
path: root/src/config.rs
blob: e7ee2e9a6165621764ba5075e4adbc8a4a0ce060 (plain)
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
use std::{
    collections::HashSet, net::IpAddr, path::PathBuf, str::FromStr, sync::atomic::AtomicBool,
};

use clap::Parser;
use ipnet::IpNet;
#[derive(Debug, Clone, Parser)]
pub struct Config {
    /// resolve hostnames using the system resolver
    #[arg(short, long)]
    pub resolve: bool,

    /// filter hosts with open tcp port at <PORT>
    #[arg(short, long)]
    pub port: Option<u16>,

    /// <TIMEOUT> (ms) used for ICMP (or TCP when using --port=<PORT>)
    #[arg(short, long, default_value_t = 250)]
    pub timeout: u64,

    /// Include network neighbours
    #[arg(short, long)]
    pub neigh: bool,

    /// Discover hosts on <NETWORK>
    #[arg(short, long, id = "NETWORK")]
    pub scan: Option<Option<IpNet>>,

    /// 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>,

    /// exclude <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<IpAddr>, HashSet<String>)> {
        Self::collect_values(&self.include)
    }

    pub fn excluded(&self) -> std::io::Result<(HashSet<IpAddr>, HashSet<String>)> {
        Self::collect_values(&self.exclude)
    }

    fn collect_values(
        values: &[IncludeExclude],
    ) -> std::io::Result<(HashSet<IpAddr>, HashSet<String>)> {
        use std::io::BufRead;
        values.iter().try_fold(
            (HashSet::new(), HashSet::new()),
            |(mut ips, mut hosts), item| {
                match item {
                    IncludeExclude::Stdin => std::io::stdin()
                        .lock()
                        .lines()
                        .map_while(Result::ok)
                        .for_each(|s| Self::partition(&mut ips, &mut hosts, s)),
                    IncludeExclude::File(filepath) => {
                        std::io::BufReader::new(std::fs::File::open(filepath)?)
                            .lines()
                            .map_while(Result::ok)
                            .for_each(|s| Self::partition(&mut ips, &mut hosts, s))
                    }
                    IncludeExclude::Item(s) => Self::partition(&mut ips, &mut hosts, s.to_owned()),
                }
                Ok((ips, hosts))
            },
        )
    }

    fn partition(ips: &mut HashSet<std::net::IpAddr>, hosts: &mut HashSet<String>, item: String) {
        match std::net::IpAddr::from_str(&item) {
            Ok(ip) => ips.insert(ip),
            Err(_) => hosts.insert(item),
        };
    }
}

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())),
            }
        }
    }
}