Refactor RangeIterator to not require Option types

There's no need to deal with `unwrap` for this specific algorithm,
instead we can rely on simple types to achieve the same result.

README.md

@@ -281,7 +281,7 @@ In this case, I would say it is safe to increase the open file limit. As most Li
 Although, if this breaks anything, please don't blame me. 
 
 **Windows Subsystem for Linux**
-Windows Subsystem for Linux does not support ulimt (see issue #39). 
+Windows Subsystem for Linux does not support ulimit (see issue #39). 
 
 The best way is to use it on a host computer, in Docker, or in a VM that isn't WSL. 
 

src/main.rs

@@ -5,6 +5,9 @@ mod tui;
 mod scanner;
 use scanner::Scanner;
 
+mod port_strategy;
+use port_strategy::PortStrategy;
+
 use colorful::Color;
 use colorful::Colorful;
 use futures::executor::block_on;
@@ -13,7 +16,7 @@ use rlimit::{getrlimit, setrlimit};
 use std::collections::HashMap;
 use std::process::Command;
 use std::{net::IpAddr, time::Duration};
-use structopt::StructOpt;
+use structopt::{clap::arg_enum, StructOpt};
 
 extern crate colorful;
 extern crate dirs;
@@ -28,6 +31,14 @@ const AVERAGE_BATCH_SIZE: rlimit::rlim = 3000;
 #[macro_use]
 extern crate log;
 
+arg_enum! {
+    #[derive(Debug, StructOpt)]
+    pub enum ScanOrder {
+        Serial,
+        Random,
+    }
+}
+
 #[derive(StructOpt, Debug)]
 #[structopt(name = "rustscan", setting = structopt::clap::AppSettings::TrailingVarArg)]
 /// Fast Port Scanner built in Rust.
@@ -44,6 +55,10 @@ struct Opts {
     #[structopt(short, long)]
     quiet: bool,
 
+    //Accessible mode. Turns off features which negatively affect screen readers.
+    #[structopt(short, long)]
+    accessible: bool,
+
     /// The batch size for port scanning, it increases or slows the speed of
     /// scanning. Depends on the open file limit of your OS.  If you do 65535
     /// it will do every port at the same time. Although, your OS may not
@@ -59,6 +74,12 @@ struct Opts {
     #[structopt(short, long)]
     ulimit: Option<rlimit::rlim>,
 
+    /// The order of scanning to be performed. The "serial" option will
+    /// scan ports in ascending order while the "random" option will scan
+    /// ports randomly.
+    #[structopt(long, possible_values = &ScanOrder::variants(), case_insensitive = true, default_value = "serial")]
+    scan_order: ScanOrder,
+
     /// The Nmap arguments to run.
     /// To use the argument -A, end RustScan's args with '-- -A'.
     /// Example: 'rustscan -T 1500 127.0.0.1 -- -A -sC'.
@@ -78,7 +99,7 @@ fn main() {
     let opts = Opts::from_args();
     info!("Mains() `opts` arguments are {:?}", opts);
 
-    if !opts.quiet {
+    if !opts.quiet && !opts.accessible {
         print_opening();
     }
 
@@ -87,11 +108,10 @@ fn main() {
 
     let scanner = Scanner::new(
         &opts.ips,
-        LOWEST_PORT_NUMBER,
-        TOP_PORT_NUMBER,
         batch_size,
         Duration::from_millis(opts.timeout.into()),
         opts.quiet,
+        PortStrategy::pick(LOWEST_PORT_NUMBER, TOP_PORT_NUMBER, opts.scan_order),
     );
 
     let scan_result = block_on(scanner.run());
@@ -234,10 +254,7 @@ fn infer_batch_size(opts: &Opts, ulimit: rlimit::rlim) -> u16 {
 
     // Adjust the batch size when the ulimit value is lower than the desired batch size
     if ulimit < batch_size {
-        warning!(
-            "File limit is lower than default batch size.
-         Consider upping with --ulimt. 
-         May cause harm to sensitive servers",
+        warning!("File limit is lower than default batch size. Consider upping with --ulimt. May cause harm to sensitive servers",
             opts.quiet
         );
 
@@ -248,7 +265,7 @@ fn infer_batch_size(opts: &Opts, ulimit: rlimit::rlim) -> u16 {
             // ulimit is smaller than aveage batch size
             // user must have very small ulimit
             // decrease batch size to half of ulimit
-            warning!("Your file limit is very small, which negatively impacts RustScan's speed. Use the Docker image, or up the Ulimt with '--ulimt 5000'. ");
+            warning!("Your file limit is very small, which negatively impacts RustScan's speed. Use the Docker image, or up the Ulimit with '--ulimit 5000'. ");
             info!("Halving batch_size because ulimit is smaller than average batch size");
             batch_size = ulimit / 2
         } else if ulimit > DEFAULT_FILE_DESCRIPTORS_LIMIT {
@@ -272,7 +289,7 @@ fn infer_batch_size(opts: &Opts, ulimit: rlimit::rlim) -> u16 {
 
 #[cfg(test)]
 mod tests {
-    use crate::{adjust_ulimit_size, infer_batch_size, print_opening, Opts};
+    use crate::{adjust_ulimit_size, infer_batch_size, print_opening, Opts, ScanOrder};
     use std::{net::IpAddr, str::FromStr};
 
     #[test]
@@ -284,6 +301,8 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            accessible: false,
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 120);
 
@@ -299,6 +318,8 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            accessible: false,
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 9_000);
 
@@ -315,6 +336,8 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            accessible: false,
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 5_000);
 
@@ -330,6 +353,8 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            accessible: false,
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = adjust_ulimit_size(&opts);
 
@@ -350,6 +375,8 @@ mod tests {
             timeout: 1_000,
             ulimit: None,
             command: Vec::new(),
+            accessible: true,
+            scan_order: ScanOrder::Serial,
         };
 
         infer_batch_size(&opts, 1_000_000);

src/port_strategy/mod.rs

@@ -0,0 +1,74 @@
+mod range_iterator;
+use super::ScanOrder;
+use range_iterator::RangeIterator;
+
+pub enum PortStrategy {
+    Serial(SerialRange),
+    Random(RandomRange),
+}
+
+impl PortStrategy {
+    pub fn pick(start: u16, end: u16, scan_type: ScanOrder) -> Self {
+        match scan_type {
+            ScanOrder::Serial => PortStrategy::Serial(SerialRange { start, end }),
+            ScanOrder::Random => PortStrategy::Random(RandomRange { start, end }),
+        }
+    }
+
+    pub fn order(&self) -> Vec<u16> {
+        match self {
+            PortStrategy::Serial(range) => range.generate(),
+            PortStrategy::Random(range) => range.generate(),
+        }
+    }
+}
+
+trait RangeOrder {
+    fn generate(&self) -> Vec<u16>;
+}
+
+pub struct SerialRange {
+    start: u16,
+    end: u16,
+}
+
+impl RangeOrder for SerialRange {
+    fn generate(&self) -> Vec<u16> {
+        (self.start..self.end).collect()
+    }
+}
+
+pub struct RandomRange {
+    start: u16,
+    end: u16,
+}
+
+impl RangeOrder for RandomRange {
+    fn generate(&self) -> Vec<u16> {
+        RangeIterator::new(self.start.into(), self.end.into()).collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::PortStrategy;
+    use crate::ScanOrder;
+
+    #[test]
+    fn serial_strategy() {
+        let strategy = PortStrategy::pick(1, 100, ScanOrder::Serial);
+        let result = strategy.order();
+        let expected_range = (1..100).into_iter().collect::<Vec<u16>>();
+        assert_eq!(expected_range, result);
+    }
+    #[test]
+    fn random_strategy() {
+        let strategy = PortStrategy::pick(1, 100, ScanOrder::Random);
+        let mut result = strategy.order();
+        let expected_range = (1..100).into_iter().collect::<Vec<u16>>();
+        assert_ne!(expected_range, result);
+
+        result.sort();
+        assert_eq!(expected_range, result);
+    }
+}

src/port_strategy/range_iterator.rs

@@ -2,9 +2,10 @@ use gcd::Gcd;
 use rand::Rng;
 
 pub struct RangeIterator {
+    active: bool,
     normalized_end: u32,
-    normalized_first_pick: Option<u32>,
-    normalized_pick: Option<u32>,
+    normalized_first_pick: u32,
+    normalized_pick: u32,
     actual_start: u32,
     step: u32,
 }
@@ -23,11 +24,17 @@ impl RangeIterator {
         let normalized_end = end - start;
         let step = pick_random_coprime(normalized_end);
 
+        // Randomly choose a number within the range to be the first
+        // and assign it as a pick.
+        let mut rng = rand::thread_rng();
+        let normalized_first_pick = rng.gen_range(0, normalized_end);
+
         Self {
+            active: true,
             normalized_end,
             step,
-            normalized_first_pick: None,
-            normalized_pick: None,
+            normalized_first_pick,
+            normalized_pick: normalized_first_pick,
             actual_start: start,
         }
     }
@@ -39,28 +46,21 @@ impl Iterator for RangeIterator {
     // The next step is always bound by the formula: N+1 = (N + STEP) % TOP_OF_THE_RANGE
     // It will only stop once we generate a number equal to the first generated number.
     fn next(&mut self) -> Option<Self::Item> {
-        // Randomly choose a number within the range to be the first
-        // and assign it as a pick.
-        if self.normalized_first_pick.is_none() {
-            let mut rng = rand::thread_rng();
-            let normalized_first_pick = rng.gen_range(0, self.normalized_end);
-
-            self.normalized_first_pick = Some(normalized_first_pick);
-            self.normalized_pick = Some(normalized_first_pick);
-            return Some((self.actual_start + normalized_first_pick) as u16);
+        if !self.active {
+            return None;
         }
 
-        let current_pick = self.normalized_pick.unwrap();
+        let current_pick = self.normalized_pick;
         let next_pick = (current_pick + self.step) % self.normalized_end;
 
         // If the next pick is equal to the first pick this means that
         // we have iterated through the entire range.
-        if next_pick == self.normalized_first_pick.unwrap() {
-            return None;
+        if next_pick == self.normalized_first_pick {
+            self.active = false;
         }
 
-        self.normalized_pick = Some(next_pick);
-        Some((self.actual_start + next_pick) as u16)
+        self.normalized_pick = next_pick;
+        Some((self.actual_start + current_pick) as u16)
     }
 }
 

src/scanner/mod.rs

@@ -1,5 +1,4 @@
-mod range_iterator;
-use range_iterator::RangeIterator;
+use super::PortStrategy;
 
 use async_std::io;
 use async_std::net::TcpStream;
@@ -21,28 +20,25 @@ use std::{
 #[cfg(not(tarpaulin_include))]
 pub struct Scanner {
     hosts: Vec<IpAddr>,
-    start: u16,
-    end: u16,
     batch_size: u16,
     timeout: Duration,
     quiet: bool,
+    port_strategy: PortStrategy,
 }
 
 impl Scanner {
     pub fn new(
         hosts: &[IpAddr],
-        start: u16,
-        end: u16,
         batch_size: u16,
         timeout: Duration,
         quiet: bool,
+        port_strategy: PortStrategy,
     ) -> Self {
         Self {
-            start,
-            end,
             batch_size,
             timeout,
             quiet,
+            port_strategy,
             hosts: hosts.iter().map(|host| host.to_owned()).collect(),
         }
     }
@@ -51,27 +47,13 @@ impl Scanner {
     /// If you want to run RustScan normally, this is the entry point used
     /// Returns all open ports as Vec<u16>
     pub async fn run(&self) -> Vec<SocketAddr> {
-        let range_iterator = RangeIterator::new(self.start.into(), self.end.into());
-        let mut sockets: Vec<SocketAddr> = Vec::new();
+        let ports: Vec<u16> = self.port_strategy.order();
+        let batch_per_host: usize = self.batch_size as usize / self.hosts.len();
         let mut open_sockets: Vec<SocketAddr> = Vec::new();
 
-        for port in range_iterator {
-            for host in &self.hosts {
-                sockets.push(SocketAddr::new(*host, port));
-            }
-
-            if sockets.len() >= self.batch_size.into() {
-                let mut results = self.scan_range(&sockets).await;
-                open_sockets.append(&mut results);
-                sockets = Vec::new();
-            }
-        }
-
-        // This will happen if we have a number of sockets remaining
-        // that is lower than the batch size.
-        if !sockets.is_empty() {
-            let mut results = self.scan_range(&sockets).await;
-            open_sockets.append(&mut results);
+        for batch in ports.chunks(batch_per_host) {
+            let mut sockets = self.scan_ports(batch).await;
+            open_sockets.append(&mut sockets);
         }
 
         open_sockets
@@ -79,10 +61,12 @@ impl Scanner {
 
     /// Given a slice of sockets, scan them all.
     /// Returns a vector of open sockets.
-    async fn scan_range(&self, sockets: &[SocketAddr]) -> Vec<SocketAddr> {
+    async fn scan_ports(&self, ports: &[u16]) -> Vec<SocketAddr> {
         let mut ftrs = FuturesUnordered::new();
-        for socket in sockets {
-            ftrs.push(self.scan_socket(socket));
+        for port in ports {
+            for host in &self.hosts {
+                ftrs.push(self.scan_socket(SocketAddr::new(*host, *port)));
+            }
         }
 
         let mut open_sockets: Vec<SocketAddr> = Vec::new();
@@ -108,8 +92,8 @@ impl Scanner {
     ///     self.scan_port(10:u16)
     ///
     /// Note: `self` must contain `self.host`.
-    async fn scan_socket(&self, socket: &SocketAddr) -> io::Result<SocketAddr> {
-        match self.connect(*socket).await {
+    async fn scan_socket(&self, socket: SocketAddr) -> io::Result<SocketAddr> {
+        match self.connect(socket).await {
             Ok(x) => {
                 // match stream_result.shutdown(Shutdown::Both)
                 info!("Shutting down stream");
@@ -120,7 +104,7 @@ impl Scanner {
                     println!("Open {}", socket.to_string().purple());
                 }
 
-                Ok(*socket)
+                Ok(socket)
             }
             Err(e) => match e.kind() {
                 ErrorKind::Other => {
@@ -155,6 +139,7 @@ impl Scanner {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::ScanOrder;
     use async_std::task::block_on;
     use std::{net::IpAddr, time::Duration};
 
@@ -162,7 +147,8 @@ mod tests {
     fn scanner_runs() {
         // Makes sure the program still runs and doesn't panic
         let addrs = vec!["127.0.0.1".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(100), true);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         // if the scan fails, it wouldn't be able to assert_eq! as it panicked!
         assert_eq!(1, 1);
@@ -171,7 +157,8 @@ mod tests {
     fn ipv6_scanner_runs() {
         // Makes sure the program still runs and doesn't panic
         let addrs = vec!["::1".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(100), false);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         // if the scan fails, it wouldn't be able to assert_eq! as it panicked!
         assert_eq!(1, 1);
@@ -179,14 +166,16 @@ mod tests {
     #[test]
     fn quad_zero_scanner_runs() {
         let addrs = vec!["0.0.0.0".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(500), true);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }
     #[test]
     fn google_dns_runs() {
         let addrs = vec!["8.8.8.8".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 400, 445, 10, Duration::from_millis(1_500), true);
+        let strategy = PortStrategy::pick(400, 445, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }
@@ -196,7 +185,8 @@ mod tests {
         let addrs = vec!["8.8.8.8".parse::<IpAddr>().unwrap()];
 
         // mac should have this automatically scaled down
-        let scanner = Scanner::new(&addrs, 400, 600, 10_000, Duration::from_millis(1500), true);
+        let strategy = PortStrategy::pick(400, 600, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }