refactor: move address-related functions (#585)

* refactor: move address-related functions

move functions for address-parsing/-resolution to a submodule.

fix #580

* test: add documentation tests

src/address.rs

@@ -0,0 +1,207 @@
+//! Provides functions to parse input IP addresses, CIDRs or files.
+use std::fs::File;
+use std::io::{prelude::*, BufReader};
+use std::net::{IpAddr, ToSocketAddrs};
+use std::path::Path;
+
+use cidr_utils::cidr::IpCidr;
+use log::debug;
+use trust_dns_resolver::{
+    config::{ResolverConfig, ResolverOpts},
+    Resolver,
+};
+
+use crate::input::Opts;
+use crate::warning;
+
+/// Parses the string(s) into IP addresses.
+///
+/// Goes through all possible IP inputs (files or via argparsing).
+///
+/// ```rust
+/// # use rustscan::input::Opts;
+/// # use rustscan::address::parse_addresses;
+/// let mut opts = Opts::default();
+/// opts.addresses = vec!["192.168.0.0/30".to_owned()];
+///
+/// let ips = parse_addresses(&opts);
+/// ```
+pub fn parse_addresses(input: &Opts) -> Vec<IpAddr> {
+    let mut ips: Vec<IpAddr> = Vec::new();
+    let mut unresolved_addresses: Vec<&str> = Vec::new();
+    let backup_resolver =
+        Resolver::new(ResolverConfig::cloudflare_tls(), ResolverOpts::default()).unwrap();
+
+    for address in &input.addresses {
+        let parsed_ips = parse_address(address, &backup_resolver);
+        if !parsed_ips.is_empty() {
+            ips.extend(parsed_ips);
+        } else {
+            unresolved_addresses.push(address);
+        }
+    }
+
+    // If we got to this point this can only be a file path or the wrong input.
+    for file_path in unresolved_addresses {
+        let file_path = Path::new(file_path);
+
+        if !file_path.is_file() {
+            warning!(
+                format!("Host {file_path:?} could not be resolved."),
+                input.greppable,
+                input.accessible
+            );
+
+            continue;
+        }
+
+        if let Ok(x) = read_ips_from_file(file_path, &backup_resolver) {
+            ips.extend(x);
+        } else {
+            warning!(
+                format!("Host {file_path:?} could not be resolved."),
+                input.greppable,
+                input.accessible
+            );
+        }
+    }
+
+    ips
+}
+
+/// Given a string, parse it as a host, IP address, or CIDR.
+///
+/// This allows us to pass files as hosts or cidr or IPs easily
+/// Call this every time you have a possible IP_or_host
+///
+/// ```rust
+/// # use rustscan::address::parse_address;
+/// # use trust_dns_resolver::Resolver;
+/// let ips = parse_address("127.0.0.1", &Resolver::default().unwrap());
+/// ```
+pub fn parse_address(address: &str, resolver: &Resolver) -> Vec<IpAddr> {
+    IpCidr::from_str(address)
+        .map(|cidr| cidr.iter().collect())
+        .ok()
+        .or_else(|| {
+            format!("{}:{}", &address, 80)
+                .to_socket_addrs()
+                .ok()
+                .map(|mut iter| vec![iter.next().unwrap().ip()])
+        })
+        .unwrap_or_else(|| resolve_ips_from_host(address, resolver))
+}
+
+/// Uses DNS to get the IPS associated with host
+fn resolve_ips_from_host(source: &str, backup_resolver: &Resolver) -> Vec<IpAddr> {
+    let mut ips: Vec<std::net::IpAddr> = Vec::new();
+
+    if let Ok(addrs) = source.to_socket_addrs() {
+        for ip in addrs {
+            ips.push(ip.ip());
+        }
+    } else if let Ok(addrs) = backup_resolver.lookup_ip(source) {
+        ips.extend(addrs.iter());
+    }
+
+    ips
+}
+
+#[cfg(not(tarpaulin_include))]
+/// Parses an input file of IPs and uses those
+fn read_ips_from_file(
+    ips: &std::path::Path,
+    backup_resolver: &Resolver,
+) -> Result<Vec<std::net::IpAddr>, std::io::Error> {
+    let file = File::open(ips)?;
+    let reader = BufReader::new(file);
+
+    let mut ips: Vec<std::net::IpAddr> = Vec::new();
+
+    for address_line in reader.lines() {
+        if let Ok(address) = address_line {
+            ips.extend(parse_address(&address, backup_resolver));
+        } else {
+            debug!("Line in file is not valid");
+        }
+    }
+
+    Ok(ips)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{parse_addresses, Opts};
+    use std::net::Ipv4Addr;
+
+    #[test]
+    fn parse_correct_addresses() {
+        let mut opts = Opts::default();
+        opts.addresses = vec!["127.0.0.1".to_owned(), "192.168.0.0/30".to_owned()];
+        let ips = parse_addresses(&opts);
+
+        assert_eq!(
+            ips,
+            [
+                Ipv4Addr::new(127, 0, 0, 1),
+                Ipv4Addr::new(192, 168, 0, 0),
+                Ipv4Addr::new(192, 168, 0, 1),
+                Ipv4Addr::new(192, 168, 0, 2),
+                Ipv4Addr::new(192, 168, 0, 3)
+            ]
+        );
+    }
+
+    #[test]
+    fn parse_correct_host_addresses() {
+        let mut opts = Opts::default();
+        opts.addresses = vec!["google.com".to_owned()];
+        let ips = parse_addresses(&opts);
+
+        assert_eq!(ips.len(), 1);
+    }
+
+    #[test]
+    fn parse_correct_and_incorrect_addresses() {
+        let mut opts = Opts::default();
+        opts.addresses = vec!["127.0.0.1".to_owned(), "im_wrong".to_owned()];
+        let ips = parse_addresses(&opts);
+
+        assert_eq!(ips, [Ipv4Addr::new(127, 0, 0, 1),]);
+    }
+
+    #[test]
+    fn parse_incorrect_addresses() {
+        let mut opts = Opts::default();
+        opts.addresses = vec!["im_wrong".to_owned(), "300.10.1.1".to_owned()];
+        let ips = parse_addresses(&opts);
+
+        assert!(ips.is_empty());
+    }
+    #[test]
+    fn parse_hosts_file_and_incorrect_hosts() {
+        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
+        let mut opts = Opts::default();
+        opts.addresses = vec!["fixtures/hosts.txt".to_owned()];
+        let ips = parse_addresses(&opts);
+        assert_eq!(ips.len(), 3);
+    }
+
+    #[test]
+    fn parse_empty_hosts_file() {
+        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
+        let mut opts = Opts::default();
+        opts.addresses = vec!["fixtures/empty_hosts.txt".to_owned()];
+        let ips = parse_addresses(&opts);
+        assert_eq!(ips.len(), 0);
+    }
+
+    #[test]
+    fn parse_naughty_host_file() {
+        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
+        let mut opts = Opts::default();
+        opts.addresses = vec!["fixtures/naughty_string.txt".to_owned()];
+        let ips = parse_addresses(&opts);
+        assert_eq!(ips.len(), 0);
+    }
+}

src/lib.rs

@@ -51,3 +51,5 @@ pub mod port_strategy;
 pub mod benchmark;
 
 pub mod scripts;
+
+pub mod address;

src/main.rs

@@ -9,20 +9,14 @@ use rustscan::scanner::Scanner;
 use rustscan::scripts::{init_scripts, Script, ScriptFile};
 use rustscan::{detail, funny_opening, output, warning};
 
-use cidr_utils::cidr::IpCidr;
 use colorful::{Color, Colorful};
 use futures::executor::block_on;
 use std::collections::HashMap;
-use std::fs::File;
-use std::io::{prelude::*, BufReader};
-use std::net::{IpAddr, ToSocketAddrs};
-use std::path::Path;
+use std::net::IpAddr;
 use std::string::ToString;
 use std::time::Duration;
-use trust_dns_resolver::{
-    config::{ResolverConfig, ResolverOpts},
-    Resolver,
-};
+
+use rustscan::address::parse_addresses;
 
 extern crate colorful;
 extern crate dirs;
@@ -231,104 +225,6 @@ fn print_opening(opts: &Opts) {
     );
 }
 
-/// Goes through all possible IP inputs (files or via argparsing)
-/// Parses the string(s) into IPs
-fn parse_addresses(input: &Opts) -> Vec<IpAddr> {
-    let mut ips: Vec<IpAddr> = Vec::new();
-    let mut unresolved_addresses: Vec<&str> = Vec::new();
-    let backup_resolver =
-        Resolver::new(ResolverConfig::cloudflare_tls(), ResolverOpts::default()).unwrap();
-
-    for address in &input.addresses {
-        let parsed_ips = parse_address(address, &backup_resolver);
-        if !parsed_ips.is_empty() {
-            ips.extend(parsed_ips);
-        } else {
-            unresolved_addresses.push(address);
-        }
-    }
-
-    // If we got to this point this can only be a file path or the wrong input.
-    for file_path in unresolved_addresses {
-        let file_path = Path::new(file_path);
-
-        if !file_path.is_file() {
-            warning!(
-                format!("Host {file_path:?} could not be resolved."),
-                input.greppable,
-                input.accessible
-            );
-
-            continue;
-        }
-
-        if let Ok(x) = read_ips_from_file(file_path, &backup_resolver) {
-            ips.extend(x);
-        } else {
-            warning!(
-                format!("Host {file_path:?} could not be resolved."),
-                input.greppable,
-                input.accessible
-            );
-        }
-    }
-
-    ips
-}
-
-/// Given a string, parse it as a host, IP address, or CIDR.
-/// This allows us to pass files as hosts or cidr or IPs easily
-/// Call this every time you have a possible IP_or_host
-fn parse_address(address: &str, resolver: &Resolver) -> Vec<IpAddr> {
-    IpCidr::from_str(address)
-        .map(|cidr| cidr.iter().collect())
-        .ok()
-        .or_else(|| {
-            format!("{}:{}", &address, 80)
-                .to_socket_addrs()
-                .ok()
-                .map(|mut iter| vec![iter.next().unwrap().ip()])
-        })
-        .unwrap_or_else(|| resolve_ips_from_host(address, resolver))
-}
-
-/// Uses DNS to get the IPS associated with host
-fn resolve_ips_from_host(source: &str, backup_resolver: &Resolver) -> Vec<IpAddr> {
-    let mut ips: Vec<std::net::IpAddr> = Vec::new();
-
-    if let Ok(addrs) = source.to_socket_addrs() {
-        for ip in addrs {
-            ips.push(ip.ip());
-        }
-    } else if let Ok(addrs) = backup_resolver.lookup_ip(source) {
-        ips.extend(addrs.iter());
-    }
-
-    ips
-}
-
-#[cfg(not(tarpaulin_include))]
-/// Parses an input file of IPs and uses those
-fn read_ips_from_file(
-    ips: &std::path::Path,
-    backup_resolver: &Resolver,
-) -> Result<Vec<std::net::IpAddr>, std::io::Error> {
-    let file = File::open(ips)?;
-    let reader = BufReader::new(file);
-
-    let mut ips: Vec<std::net::IpAddr> = Vec::new();
-
-    for address_line in reader.lines() {
-        if let Ok(address) = address_line {
-            ips.extend(parse_address(&address, backup_resolver));
-        } else {
-            debug!("Line in file is not valid");
-        }
-    }
-
-    Ok(ips)
-}
-
 #[cfg(unix)]
 fn adjust_ulimit_size(opts: &Opts) -> u64 {
     use rlimit::Resource;
@@ -398,8 +294,7 @@ fn infer_batch_size(opts: &Opts, ulimit: u64) -> u16 {
 mod tests {
     #[cfg(unix)]
     use super::{adjust_ulimit_size, infer_batch_size};
-    use super::{parse_addresses, print_opening, Opts};
-    use std::net::Ipv4Addr;
+    use super::{print_opening, Opts};
 
     #[test]
     #[cfg(unix)]
@@ -442,13 +337,6 @@ mod tests {
 
         assert!(batch_size == 2_000);
     }
-    #[test]
-    fn test_print_opening_no_panic() {
-        let mut opts = Opts::default();
-        opts.ulimit = Some(2_000);
-        // print opening should not panic
-        print_opening(&opts);
-    }
 
     #[test]
     #[cfg(unix)]
@@ -463,73 +351,10 @@ mod tests {
     }
 
     #[test]
-    fn parse_correct_addresses() {
+    fn test_print_opening_no_panic() {
         let mut opts = Opts::default();
-        opts.addresses = vec!["127.0.0.1".to_owned(), "192.168.0.0/30".to_owned()];
-        let ips = parse_addresses(&opts);
-
-        assert_eq!(
-            ips,
-            [
-                Ipv4Addr::new(127, 0, 0, 1),
-                Ipv4Addr::new(192, 168, 0, 0),
-                Ipv4Addr::new(192, 168, 0, 1),
-                Ipv4Addr::new(192, 168, 0, 2),
-                Ipv4Addr::new(192, 168, 0, 3)
-            ]
-        );
-    }
-
-    #[test]
-    fn parse_correct_host_addresses() {
-        let mut opts = Opts::default();
-        opts.addresses = vec!["google.com".to_owned()];
-        let ips = parse_addresses(&opts);
-
-        assert_eq!(ips.len(), 1);
-    }
-
-    #[test]
-    fn parse_correct_and_incorrect_addresses() {
-        let mut opts = Opts::default();
-        opts.addresses = vec!["127.0.0.1".to_owned(), "im_wrong".to_owned()];
-        let ips = parse_addresses(&opts);
-
-        assert_eq!(ips, [Ipv4Addr::new(127, 0, 0, 1),]);
-    }
-
-    #[test]
-    fn parse_incorrect_addresses() {
-        let mut opts = Opts::default();
-        opts.addresses = vec!["im_wrong".to_owned(), "300.10.1.1".to_owned()];
-        let ips = parse_addresses(&opts);
-
-        assert!(ips.is_empty());
-    }
-    #[test]
-    fn parse_hosts_file_and_incorrect_hosts() {
-        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
-        let mut opts = Opts::default();
-        opts.addresses = vec!["fixtures/hosts.txt".to_owned()];
-        let ips = parse_addresses(&opts);
-        assert_eq!(ips.len(), 3);
-    }
-
-    #[test]
-    fn parse_empty_hosts_file() {
-        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
-        let mut opts = Opts::default();
-        opts.addresses = vec!["fixtures/empty_hosts.txt".to_owned()];
-        let ips = parse_addresses(&opts);
-        assert_eq!(ips.len(), 0);
-    }
-
-    #[test]
-    fn parse_naughty_host_file() {
-        // Host file contains IP, Hosts, incorrect IPs, incorrect hosts
-        let mut opts = Opts::default();
-        opts.addresses = vec!["fixtures/naughty_string.txt".to_owned()];
-        let ips = parse_addresses(&opts);
-        assert_eq!(ips.len(), 0);
+        opts.ulimit = Some(2_000);
+        // print opening should not panic
+        print_opening(&opts);
     }
 }