feat: add --distributed to allow scanning multiple similar hosts at…

… the same time
cestef · Oct 9, 2024 · 8446fd7 · 8446fd7
1 parent 079a201
commit 8446fd7
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 5 deletions.
diff --git a/src/cli/opts.rs b/src/cli/opts.rs
@@ -159,6 +159,12 @@ pub struct Opts {
     #[serde(default)]
     pub insecure: bool,
 
+    /// Distribute the requests to multiple hosts
+    #[clap(long, env, hide_env = true, value_delimiter = ',')]
+    #[merge(strategy = merge::vec::overwrite_empty)]
+    #[serde(default)]
+    pub distributed: Vec<String>,
+
     /// Show response additional body information
     #[clap(
         long,

diff --git a/src/runner/classic.rs b/src/runner/classic.rs
@@ -78,10 +78,24 @@ impl Classic {
         opts: Opts,
         engine: Arc<rhai::Engine>,
     ) -> Result<()> {
-        for url in &chunk {
+        for (index, url) in chunk.iter().enumerate() {
+            let mut url = url.clone();
             let t1 = Instant::now();
+            if !opts.distributed.is_empty() {
+                let current = index % (opts.distributed.len() + 1);
+                if current != 0 {
+                    let host_for_this_request = &opts.distributed[current - 1];
 
-            let request = super::client::build_request(&opts, url, &client)?;
+                    let parsed_url = url::Url::parse(&url)?;
+                    url = format!(
+                        "{}://{}{}",
+                        parsed_url.scheme(),
+                        host_for_this_request,
+                        parsed_url.path()
+                    );
+                }
+            }
+            let request = super::client::build_request(&opts, &url, &client)?;
 
             let response = client.execute(request).await;
 
@@ -145,7 +159,7 @@ impl Classic {
                             })
                         ));
 
-                        let parsed = Url::parse(url)?;
+                        let parsed = Url::parse(&url)?;
                         let mut tree = tree.lock().clone();
                         let root_url = tree
                             .root
@@ -206,7 +220,7 @@ impl Classic {
                             url,
                             format!("{}ms", t1.elapsed().as_millis().to_string().bold()).dimmed()
                         ));
-                        let parsed = Url::parse(url)?;
+                        let parsed = Url::parse(&url)?;
                         let mut tree = tree.lock().clone();
                         let root_url = tree
                             .root
@@ -242,7 +256,7 @@ impl Classic {
                                 progress.println(msg);
                                 Ok(())
                             },
-                            url,
+                            &url,
                             err,
                         )?;
                     }

diff --git a/src/runner/recursive.rs b/src/runner/recursive.rs
@@ -181,6 +181,20 @@ impl Recursive {
             let data = previous_node.lock().data.clone();
 
             let mut url = data.url.clone();
+            if !opts.distributed.is_empty() {
+                let current = index % (opts.distributed.len() + 1);
+                if current != 0 {
+                    let host_for_this_request = &opts.distributed[current - 1];
+
+                    let parsed_url = url::Url::parse(&url)?;
+                    url = format!(
+                        "{}://{}{}",
+                        parsed_url.scheme(),
+                        host_for_this_request,
+                        parsed_url.path()
+                    );
+                }
+            }
             match url.ends_with('/') {
                 true => url.push_str(&word),
                 false => url.push_str(&format!("/{}", word)),