misc: refactored credentials generators to use traits

src/creds/combinator.rs

@@ -1,18 +1,17 @@
 use std::time;
 
 use crate::{
-    creds::{self, expression, Credentials, Generator},
+    creds::{self, expression, generator, Credentials},
     options::Options,
     session::Error,
 };
 
-#[derive(Default, Debug)]
 pub(crate) struct Combinator {
     options: Options,
     user_expr: creds::Expression,
-    user_it: creds::Generator,
+    user_it: Box<dyn creds::Generator>,
     current_user: Option<String>,
-    pass_it: creds::Generator,
+    pass_it: Box<dyn creds::Generator>,
     pass_expr: creds::Expression,
     dispatched: usize,
     total: usize,
@@ -28,25 +27,25 @@ impl Combinator {
             } else {
                 expression::parse_expression(options.password.as_ref())
             };
-            let user_it = Generator::new(user_expr.clone())?;
+            let user_it = generator::new(user_expr.clone())?;
 
             (
                 user_expr,
                 user_it,
                 creds::Expression::default(),
-                creds::Generator::default(),
+                generator::empty()?,
             )
         } else {
             // get both
             let user_expr = expression::parse_expression(options.username.as_ref());
-            let user_it = Generator::new(user_expr.clone())?;
+            let user_it = generator::new(user_expr.clone())?;
 
             let expr = expression::parse_expression(options.password.as_ref());
             (
                 user_expr,
                 user_it,
                 expr.clone(),
-                Generator::new(expr.clone())?,
+                generator::new(expr.clone())?,
             )
         };
 
@@ -96,7 +95,7 @@ impl Combinator {
             (false, next_pass)
         } else {
             // reset internal iterator
-            self.pass_it = Generator::new(self.pass_expr.clone()).unwrap();
+            self.pass_it = generator::new(self.pass_expr.clone()).unwrap();
             (true, self.pass_it.next().unwrap())
         }
     }

src/creds/generator.rs

@@ -5,30 +5,96 @@ use crate::creds::expression::Expression;
 use crate::creds::permutator::Permutator;
 use crate::session::Error;
 
-#[derive(Default, Debug)]
-pub(crate) struct Generator {
-    constant: Option<String>,
-    lines: Option<Lines<BufReader<File>>>,
-    permutator: Option<Permutator>,
-    paths: Option<glob::Paths>,
-    current: usize,
-    elements: usize,
+pub(crate) trait Generator: Iterator<Item = String> {
+    fn search_space_size(&self) -> usize;
 }
 
-// TODO: this should probably be refactored into a trait
-impl Generator {
-    fn from_constant_value(value: String) -> Self {
-        Generator {
-            constant: Some(value),
-            permutator: None,
-            paths: None,
-            elements: 1,
-            current: 0,
-            lines: None,
+pub(crate) fn new(expr: Expression) -> Result<Box<dyn Generator>, Error> {
+    match expr {
+        Expression::Constant { value } => {
+            let gen = Constant::new(value)?;
+            Ok(Box::new(gen))
+        }
+        Expression::Wordlist { filename } => {
+            let gen = Wordlist::new(filename)?;
+            Ok(Box::new(gen))
         }
+        Expression::Range { min, max, charset } => {
+            let gen = Range::new(charset, min, max)?;
+            Ok(Box::new(gen))
+        }
+        Expression::Glob { pattern } => {
+            let gen = Glob::new(pattern)?;
+            Ok(Box::new(gen))
+        }
+    }
+}
+
+pub(crate) fn empty() -> Result<Box<dyn Generator>, Error> {
+    Ok(Box::new(Empty::new()))
+}
+
+struct Empty {}
+
+impl Empty {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl Generator for Empty {
+    fn search_space_size(&self) -> usize {
+        0
     }
+}
 
-    fn from_wordlist(path: String) -> Result<Self, Error> {
+impl Iterator for Empty {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        None
+    }
+}
+
+struct Constant {
+    done: bool,
+    value: String,
+}
+
+impl Constant {
+    pub fn new(value: String) -> Result<Self, Error> {
+        let done = false;
+        Ok(Self { done, value })
+    }
+}
+
+impl Generator for Constant {
+    fn search_space_size(&self) -> usize {
+        1
+    }
+}
+
+impl Iterator for Constant {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            None
+        } else {
+            self.done = true;
+            Some(self.value.to_owned())
+        }
+    }
+}
+
+struct Wordlist {
+    lines: Lines<BufReader<File>>,
+    current: usize,
+    elements: usize,
+}
+
+impl Wordlist {
+    pub fn new(path: String) -> Result<Self, Error> {
         log::debug!("loading wordlist from {} ...", &path);
 
         // count the number of lines first
@@ -40,101 +106,109 @@ impl Generator {
         let file = File::open(path).map_err(|e| e.to_string())?;
         let reader = BufReader::new(file);
 
-        Ok(Generator {
-            constant: None,
-            permutator: None,
-            paths: None,
+        Ok(Self {
             elements,
             current: 0,
-            lines: Some(reader.lines()),
+            lines: reader.lines(),
         })
     }
+}
+
+impl Generator for Wordlist {
+    fn search_space_size(&self) -> usize {
+        self.elements
+    }
+}
+
+impl Iterator for Wordlist {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current < self.elements {
+            self.current += 1;
+            if let Some(res) = self.lines.next() {
+                if let Ok(line) = res {
+                    return Some(line);
+                } else {
+                    log::error!("could not read line: {:?}", res.err());
+                }
+            }
+        }
+        None
+    }
+}
+
+struct Range {
+    permutator: Permutator,
+    elements: usize,
+}
 
-    fn from_range(charset: String, min_length: usize, max_length: usize) -> Result<Self, Error> {
+impl Range {
+    pub fn new(charset: String, min_length: usize, max_length: usize) -> Result<Self, Error> {
         if min_length == 0 {
             return Err("min length can't be zero".to_owned());
         } else if min_length > max_length {
             return Err("min length can't be greater than max length".to_owned());
         }
 
-        let gen = Permutator::new(charset.chars().collect(), min_length, max_length);
-        let tot = gen.search_space_size();
-        let generator = Some(gen);
+        let permutator = Permutator::new(charset.chars().collect(), min_length, max_length);
+        let elements = permutator.search_space_size();
 
-        Ok(Generator {
-            lines: None,
-            permutator: generator,
-            paths: None,
-            current: 0,
-            elements: tot,
-            constant: None,
+        Ok(Self {
+            permutator,
+            elements,
         })
     }
+}
 
-    fn from_glob(pattern: String) -> Result<Self, Error> {
+impl Generator for Range {
+    fn search_space_size(&self) -> usize {
+        self.elements
+    }
+}
+
+impl Iterator for Range {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.permutator.next()
+    }
+}
+
+struct Glob {
+    paths: glob::Paths,
+    elements: usize,
+}
+
+impl Glob {
+    pub fn new(pattern: String) -> Result<Self, Error> {
         // validate the pattern and count the elements first
         let paths = match glob::glob(&pattern) {
             Err(e) => return Err(e.to_string()),
             Ok(paths) => paths,
         };
+        let elements = paths.count();
+        let paths = glob::glob(&pattern).unwrap();
 
-        Ok(Generator {
-            constant: None,
-            lines: None,
-            permutator: None,
-            paths: Some(glob::glob(&pattern).unwrap()),
-            current: 0,
-            elements: paths.count(),
-        })
-    }
-
-    pub fn new(expr: Expression) -> Result<Self, Error> {
-        match expr {
-            Expression::Constant { value } => Ok(Self::from_constant_value(value)),
-            Expression::Wordlist { filename } => Self::from_wordlist(filename),
-            Expression::Range { min, max, charset } => Self::from_range(charset, min, max),
-            Expression::Glob { pattern } => Self::from_glob(pattern),
-        }
+        Ok(Self { paths, elements })
     }
+}
 
-    pub fn search_space_size(&self) -> usize {
+impl Generator for Glob {
+    fn search_space_size(&self) -> usize {
         self.elements
     }
-
-    fn next_line(&mut self) -> Option<String> {
-        if let Some(lines) = &mut self.lines {
-            if let Some(res) = lines.next() {
-                if let Ok(line) = res {
-                    return Some(line);
-                } else {
-                    log::error!("could not read line: {:?}", res.err());
-                }
-            }
-        }
-        None
-    }
 }
 
-impl Iterator for Generator {
+impl Iterator for Glob {
     type Item = String;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.current < self.elements {
-            self.current += 1;
-            if let Some(value) = &self.constant {
-                return Some(value.to_owned());
-            } else if self.lines.is_some() {
-                return self.next_line();
-            } else if let Some(permutator) = self.permutator.as_mut() {
-                return permutator.next();
-            } else if let Some(paths) = &mut self.paths {
-                if let Some(next) = paths.next() {
-                    if let Ok(path) = next {
-                        return Some(path.to_str().unwrap().to_owned());
-                    } else {
-                        log::error!("glob error: {:?}", next);
-                    }
-                }
+        if let Some(next) = self.paths.next() {
+            if let Ok(path) = next {
+                return Some(path.to_str().unwrap().to_owned());
+            } else {
+                log::error!("glob error: {:?}", next);
             }
         }
         None
@@ -146,12 +220,11 @@ mod tests {
     use std::fs::File;
     use std::io::Write;
 
-    use super::Generator;
-    use crate::creds::Expression;
+    use crate::creds::{generator, Expression};
 
     #[test]
     fn can_handle_constant() {
-        let gen = Generator::new(Expression::Constant {
+        let gen = generator::new(Expression::Constant {
             value: "hi".to_owned(),
         })
         .unwrap();
@@ -177,7 +250,7 @@ mod tests {
         tmpwordlist.flush().unwrap();
         drop(tmpwordlist);
 
-        let gen = Generator::new(Expression::Wordlist {
+        let gen = generator::new(Expression::Wordlist {
             filename: tmppath.to_str().unwrap().to_owned(),
         })
         .unwrap();
@@ -193,7 +266,7 @@ mod tests {
         let expected = vec![
             "a", "b", "c", "aa", "ab", "ac", "ba", "bb", "bc", "ca", "cb", "cc",
         ];
-        let gen = Generator::new(Expression::Range {
+        let gen = generator::new(Expression::Range {
             min: 1,
             max: 2,
             charset: "abc".to_owned(),
@@ -225,7 +298,7 @@ mod tests {
             expected.push(format!("{}/{}", &tmpdirname, filename));
         }
 
-        let gen = Generator::new(Expression::Glob {
+        let gen = generator::new(Expression::Glob {
             pattern: format!("{}/*.txt", tmpdirname),
         })
         .unwrap();