From 4fbffb6a0d57cc1d4d3f55d6006b7f443c5023f0 Mon Sep 17 00:00:00 2001 From: Simone Margaritelli Date: Sat, 4 Nov 2023 18:50:18 +0100 Subject: [PATCH] new: implemented --iterate-by argument (ref #7) --- src/creds/combinator.rs | 176 ++++++++++++++++++++++++++++++++++++---- src/creds/mod.rs | 2 +- src/options.rs | 6 +- 3 files changed, 164 insertions(+), 20 deletions(-) diff --git a/src/creds/combinator.rs b/src/creds/combinator.rs index 0144deb..d8d9315 100644 --- a/src/creds/combinator.rs +++ b/src/creds/combinator.rs @@ -1,6 +1,8 @@ use std::time; +use clap::ValueEnum; use itertools::Itertools; +use serde::{Deserialize, Serialize}; use crate::{ creds::{self, expression, iterator, Credentials}, @@ -10,6 +12,13 @@ use crate::{ use super::Expression; +#[derive(ValueEnum, Serialize, Deserialize, Debug, Default, Clone)] +pub(crate) enum IterationStrategy { + #[default] + User, + Password, +} + pub(crate) struct Combinator { options: Options, @@ -32,13 +41,41 @@ impl Combinator { } } + fn combine_iterators( + options: &Options, + targets: Vec, + user_it: Box, + pass_it: Option>, + ) -> Box> { + if let Some(pass_it) = pass_it { + let (outer, inner) = match options.iterate_by { + IterationStrategy::User => (user_it, pass_it), + IterationStrategy::Password => (pass_it, user_it), + }; + + Box::new( + targets + .into_iter() + .cartesian_product(outer) + .cartesian_product(inner) + .map(|((t, out), inn)| (t.to_owned(), out, inn)), + ) + } else { + Box::new( + targets + .into_iter() + .cartesian_product(user_it) + .map(|(t, payload)| (t.to_owned(), payload, "".to_owned())), + ) + } + } + fn for_single_payload( targets: &Vec, options: Options, override_expr: Option, ) -> Result { let dispatched = 0; - let targets: Vec = targets.to_owned(); // get either override, username or password let payload_expr = if let Some(override_expr) = override_expr { override_expr @@ -49,12 +86,7 @@ impl Combinator { }; let payload_it = iterator::new(payload_expr.clone())?; let search_space_size: usize = targets.len() * payload_it.search_space_size(); - let product = Box::new( - targets - .into_iter() - .cartesian_product(payload_it) - .map(|(t, payload)| (t.to_owned(), payload, "".to_owned())), - ); + let product = Self::combine_iterators(&options, targets.to_owned(), payload_it, None); Ok(Self { options, @@ -68,7 +100,6 @@ impl Combinator { fn for_double_payload(targets: &Vec, options: Options) -> Result { let dispatched = 0; - let targets: Vec = targets.to_owned(); // get both let user_expr = expression::parse_expression(options.username.as_ref()); let user_it = iterator::new(user_expr.clone())?; @@ -76,14 +107,7 @@ impl Combinator { let pass_it = iterator::new(pass_expr.clone())?; let search_space_size = targets.len() * user_it.search_space_size() * pass_it.search_space_size(); - - let product = Box::new( - targets - .into_iter() - .cartesian_product(user_it) - .cartesian_product(pass_it) - .map(|((t, u), p)| (t.to_owned(), u, p)), - ); + let product = Self::combine_iterators(&options, targets.to_owned(), user_it, Some(pass_it)); Ok(Self { options, @@ -133,12 +157,17 @@ impl Iterator for Combinator { fn next(&mut self) -> Option { // we're done - if let Some((target, username, password)) = self.product.next() { + if let Some((target, outer, inner)) = self.product.next() { // check if we have to rate limit if self.options.rate_limit > 0 && self.dispatched % self.options.rate_limit == 0 { std::thread::sleep(time::Duration::from_secs(1)); } + let (username, password) = match self.options.iterate_by { + IterationStrategy::User => (outer, inner), + IterationStrategy::Password => (inner, outer), + }; + Some(Credentials { target, username, @@ -155,10 +184,121 @@ mod tests { use std::fs::File; use std::io::Write; - use crate::creds::{Credentials, Expression}; + use crate::creds::{Credentials, Expression, IterationStrategy}; use super::Combinator; + #[test] + fn can_handle_user_iteration_strategy() { + let targets = vec!["foo".to_owned()]; + let mut opts = crate::Options::default(); + + opts.iterate_by = IterationStrategy::User; // default + opts.username = Some("#1-2:u".to_owned()); + opts.password = Some("#1-2:p".to_owned()); + + let comb = Combinator::create(&targets, opts, 0, false, None).unwrap(); + let expected = vec![ + Credentials { + target: "foo".to_owned(), + username: "u".to_owned(), + password: "p".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "u".to_owned(), + password: "pp".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "uu".to_owned(), + password: "p".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "uu".to_owned(), + password: "pp".to_owned(), + }, + ]; + let mut got = vec![]; + for cred in comb { + got.push(cred); + } + + assert_eq!(expected, got); + } + + #[test] + fn can_handle_password_iteration_strategy() { + let targets = vec!["foo".to_owned()]; + let mut opts = crate::Options::default(); + + opts.iterate_by = IterationStrategy::Password; + opts.username = Some("#1-2:u".to_owned()); + opts.password = Some("#1-2:p".to_owned()); + + let comb = Combinator::create(&targets, opts, 0, false, None).unwrap(); + let expected = vec![ + Credentials { + target: "foo".to_owned(), + username: "u".to_owned(), + password: "p".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "uu".to_owned(), + password: "p".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "u".to_owned(), + password: "pp".to_owned(), + }, + Credentials { + target: "foo".to_owned(), + username: "uu".to_owned(), + password: "pp".to_owned(), + }, + ]; + let mut got = vec![]; + for cred in comb { + got.push(cred); + } + + assert_eq!(expected, got); + } + + #[test] + fn iteration_strategies_return_same_results() { + let targets = vec!["foo".to_owned()]; + + let mut by_user_opts = crate::Options::default(); + by_user_opts.iterate_by = IterationStrategy::User; + by_user_opts.username = Some("#1-2:u".to_owned()); + by_user_opts.password = Some("#1-5:p".to_owned()); + + let mut by_pass_opts = crate::Options::default(); + by_pass_opts.iterate_by = IterationStrategy::Password; + by_pass_opts.username = Some("#1-2:u".to_owned()); + by_pass_opts.password = Some("#1-5:p".to_owned()); + + let by_user_comb = Combinator::create(&targets, by_user_opts, 0, false, None).unwrap(); + let by_pass_comb = Combinator::create(&targets, by_pass_opts, 0, false, None).unwrap(); + + assert_eq!( + by_user_comb.search_space_size(), + by_pass_comb.search_space_size() + ); + + let mut by_user: Vec = by_user_comb.collect(); + let mut by_pass: Vec = by_pass_comb.collect(); + + by_user.sort_by(|a, b| a.partial_cmp(b).unwrap()); + by_pass.sort_by(|a, b| a.partial_cmp(b).unwrap()); + + assert_eq!(by_user, by_pass); + } + #[test] fn can_handle_multiple_targets_and_double_credentials() { let targets = vec!["foo".to_owned(), "bar".to_owned()]; diff --git a/src/creds/mod.rs b/src/creds/mod.rs index 135b8b7..1ceede3 100644 --- a/src/creds/mod.rs +++ b/src/creds/mod.rs @@ -2,7 +2,7 @@ mod combinator; mod expression; mod iterator; -pub(crate) use combinator::Combinator; +pub(crate) use combinator::{Combinator, IterationStrategy}; pub(crate) use expression::{parse_expression, Expression}; pub(crate) use iterator::{Iterator, IteratorClone}; diff --git a/src/options.rs b/src/options.rs index 3b60062..1c3a4b7 100644 --- a/src/options.rs +++ b/src/options.rs @@ -1,7 +1,7 @@ use clap::Parser; use serde::{Deserialize, Serialize}; -use crate::session; +use crate::{creds, session}; // NOTE: normally we'd be using clap subcommands, but this approach allows us more flexibility // for plugins registered at runtime, aliases (like ssh/sftp) and so on. @@ -31,6 +31,10 @@ pub(crate) struct Options { #[clap(long, visible_alias = "key")] pub password: Option, + /// Whether to iterate by user or by password. + #[clap(long, value_enum)] + pub iterate_by: creds::IterationStrategy, + /// Save and restore session information to this file. #[clap(short, long)] pub session: Option,