From d561e59ec9945b0196845f306fc34d96b75be4be Mon Sep 17 00:00:00 2001
From: epi <epibar052@gmail.com>
Date: Sat, 18 Mar 2023 11:44:45 -0500
Subject: [PATCH] added test

---
 src/main.rs        | 16 ++++++++++++++--
 tests/test_main.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 4e876dc..58350dd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -238,20 +238,32 @@ async fn wrapped_main(config: Arc<Configuration>) -> Result<()> {
     }
 
     let words = if config.wordlist.starts_with("http") {
+        // found a url scheme, attempt to download the wordlist
         let response = config.client.get(&config.wordlist).send().await?;
 
+        if !response.status().is_success() {
+            // status code isn't a 200, bail
+            bail!(
+                "[{}] Unable to download wordlist from url: {}",
+                response.status().as_str(),
+                config.wordlist
+            );
+        }
+
+        // attempt to get the filename from the url's path
         let Some(path_segments) = response
             .url()
             .path_segments() else {
-                bail!("Unable to parse path segments from url: {}", config.wordlist);
+                bail!("Unable to parse path from url: {}", response.url());
             };
 
         let Some(filename) = path_segments.last() else {
-            bail!("Unable to parse filename from url: {}", config.wordlist);
+            bail!("Unable to parse filename from url's path: {}", response.url().path());
         };
 
         let filename = filename.to_string();
 
+        // read the body and write it to disk, then use existing code to read the wordlist
         let body = response.text().await?;
 
         std::fs::write(&filename, body)?;
diff --git a/tests/test_main.rs b/tests/test_main.rs
index 6db6a60..11b6089 100644
--- a/tests/test_main.rs
+++ b/tests/test_main.rs
@@ -1,4 +1,5 @@
 mod utils;
+use anyhow::Result;
 use assert_cmd::Command;
 use httpmock::Method::GET;
 use httpmock::{MockServer, Regex};
@@ -218,3 +219,46 @@ fn main_parallel_creates_output_directory() -> Result<(), Box<dyn std::error::Er
 
     Ok(())
 }
+
+#[test]
+/// download a wordlist from a url
+fn main_download_wordlist_from_url() -> Result<(), Box<dyn std::error::Error>> {
+    let srv = MockServer::start();
+
+    let (tmp_dir, _) = setup_tmp_directory(&["a".to_string()], "wordlist")?;
+
+    let mock1 = srv.mock(|when, then| {
+        when.method(GET).path("/derp");
+        then.status(200).body("stuff\nthings");
+    });
+
+    // serve endpoints stuff and things
+    let mock2 = srv.mock(|when, then| {
+        when.method(GET).path("/stuff");
+        then.status(200);
+    });
+
+    let mock3 = srv.mock(|when, then| {
+        when.method(GET).path("/things");
+        then.status(200);
+    });
+
+    Command::cargo_bin("feroxbuster")
+        .unwrap()
+        .current_dir(&tmp_dir)
+        .arg("--url")
+        .arg(srv.url("/"))
+        .arg("--wordlist")
+        .arg(srv.url("/derp"))
+        .assert()
+        .success()
+        .stderr(predicate::str::contains(srv.url("/derp")));
+
+    teardown_tmp_directory(tmp_dir);
+
+    assert_eq!(mock1.hits(), 1); // downloaded wordlist
+    assert_eq!(mock2.hits(), 1); // found stuff from wordlist
+    assert_eq!(mock3.hits(), 1); // found things from wordlist
+
+    Ok(())
+}