diff --git a/Cargo.toml b/Cargo.toml index 8b580db..88dcad9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ env_logger = "0.10.1" field_accessor_pub = "0.5.2" futures = "0.3.29" indicatif = "0.17.7" +itertools = "0.12.1" lazy_static = "1.4.0" log = "0.4.20" md5 = "0.7.0" diff --git a/src/cli/opts.rs b/src/cli/opts.rs index b2401a9..01828a1 100644 --- a/src/cli/opts.rs +++ b/src/cli/opts.rs @@ -30,12 +30,22 @@ pub struct Opts { )] #[merge(strategy = merge::vec::overwrite_empty)] pub wordlists: Vec, - + /// Crawl mode + #[clap( + short, + long, + default_value = "recursive", + value_name = "MODE", + value_parser = clap::builder::PossibleValuesParser::new(["recursive", "recursion", "r", "permutations", "p", "permutation", "classic", "c"]), + env, + hide_env = true + )] + pub mode: Option, /// Number of threads to use #[clap(short, long, env, hide_env = true)] pub threads: Option, - /// Maximum depth to crawl - #[clap(short, long, default_value = "1", env, hide_env = true)] + /// Crawl recursively until given depth + #[clap(short, long, env, hide_env = true, default_value = "1")] pub depth: Option, /// Output file #[clap(short, long, value_name = "FILE", env, hide_env = true)] @@ -47,7 +57,7 @@ pub struct Opts { #[clap(short, long, env, hide_env = true)] pub user_agent: Option, /// HTTP method - #[clap(short, long, default_value = "GET", value_parser = parse_method, env, hide_env=true)] + #[clap(short = 'X', long, default_value = "GET", value_parser = parse_method, env, hide_env=true)] pub method: Option, /// Data to send with the request #[clap(short = 'D', long, env, hide_env = true)] diff --git a/src/main.rs b/src/main.rs index d5302ca..e66542e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,9 @@ #![allow(dead_code)] use crate::utils::{ - constants::SUCCESS, + constants::{REPLACE_KEYWORD, SUCCESS}, save_to_file, + structs::Mode, tree::{Tree, TreeData}, }; use anyhow::{bail, Result}; @@ -10,9 +11,9 @@ use clap::Parser; use cli::opts::Opts; use colored::Colorize; use futures::future::abortable; -use futures::stream::StreamExt; +use futures::{stream::StreamExt, FutureExt}; use indicatif::HumanDuration; -use log::{error, info}; +use log::{error, info, warn}; use merge::Merge; use parking_lot::Mutex; use ptree::print_tree; @@ -77,6 +78,43 @@ pub async fn _main(opts: Opts) -> Result<()> { error!("Missing URL"); return Ok(()); } + let mode: Mode = if opts.depth.unwrap() > 1 { + Mode::Recursive + } else { + opts.mode.as_deref().unwrap().into() + }; + let mut url = opts.url.clone().unwrap(); + match mode { + Mode::Recursive => { + if opts.depth.is_none() { + error!("Missing depth"); + return Ok(()); + } + if url.matches(REPLACE_KEYWORD).count() > 0 { + warn!( + "URL contains the replace keyword: {}, this is supported with {}", + REPLACE_KEYWORD.bold(), + format!( + "{} {} | {}", + "--mode".dimmed(), + "permutations".bold(), + "classic".bold() + ) + ); + } + } + Mode::Permutations | Mode::Classic => { + if url.matches(REPLACE_KEYWORD).count() == 0 { + url = url.trim_end_matches('/').to_string() + "/" + REPLACE_KEYWORD; + warn!( + "URL does not contain the replace keyword: {}, it will be treated as: {}", + REPLACE_KEYWORD.bold(), + url.bold() + ); + } + } + } + let saved = if opts.resume { let res = tokio::fs::read_to_string(opts.save_file.clone().unwrap()).await; if !res.is_ok() { @@ -128,7 +166,7 @@ pub async fn _main(opts: Opts) -> Result<()> { error!("No words found in wordlists"); return Ok(()); } - let depth = Arc::new(Mutex::new(0)); + let current_depth = Arc::new(Mutex::new(0)); let current_indexes: Arc>>> = Arc::new(Mutex::new(HashMap::new())); @@ -137,7 +175,7 @@ pub async fn _main(opts: Opts) -> Result<()> { Some(json) => Some(utils::tree::from_save( &opts, &json.unwrap(), - depth.clone(), + current_depth.clone(), current_indexes.clone(), words.clone(), )?), @@ -153,11 +191,17 @@ pub async fn _main(opts: Opts) -> Result<()> { saved_tree } else { let t = Arc::new(Mutex::new(Tree::new())); + let cleaned_url = match mode { + Mode::Recursive => url.clone(), + Mode::Permutations | Mode::Classic => { + url.split(REPLACE_KEYWORD).collect::>()[0].to_string() + } + }; t.lock().insert( TreeData { - url: opts.url.clone().unwrap(), + url: cleaned_url.clone(), depth: 0, - path: Url::parse(&opts.url.clone().unwrap())? + path: Url::parse(&cleaned_url.clone())? .path() .to_string() .trim_end_matches('/') @@ -198,20 +242,38 @@ pub async fn _main(opts: Opts) -> Result<()> { let watch = stopwatch::Stopwatch::start_new(); info!("Press {} to save state and exit", "Ctrl+C".bold()); - let chunks = words - .chunks(words.len() / threads) - .map(|x| x.to_vec()) - .collect::>(); - let chunks = Arc::new(chunks); - - let main_fun = runner::start::run( - opts.clone(), - depth.clone(), - tree.clone(), - current_indexes.clone(), - chunks.clone(), - words.clone(), - ); + let main_fun = match mode { + Mode::Recursive => runner::recursive::run( + opts.clone(), + current_depth.clone(), + tree.clone(), + current_indexes.clone(), + Arc::new( + words + .chunks(words.len() / threads) + .map(|x| x.to_vec()) + .collect::>(), + ), + words.clone(), + ) + .boxed(), + Mode::Permutations => runner::permutations::run( + url.clone(), + opts.clone(), + tree.clone(), + words.clone(), + threads, + ) + .boxed(), + Mode::Classic => runner::classic::run( + url.clone(), + opts.clone(), + tree.clone(), + words.clone(), + threads, + ) + .boxed(), + }; let (task, handle) = if let Some(max_time) = opts.max_time { abortable(timeout(Duration::from_secs(max_time as u64), main_fun).into_inner()) } else { @@ -222,7 +284,7 @@ pub async fn _main(opts: Opts) -> Result<()> { let (tx, mut rx) = tokio::sync::mpsc::channel::<()>(1); let ctrlc_tree = tree.clone(); - let ctrlc_depth = depth.clone(); + let ctrlc_depth = current_depth.clone(); let ctrlc_words = words.clone(); let ctrlc_opts = opts.clone(); let ctrlc_aborted = aborted.clone(); @@ -272,10 +334,15 @@ pub async fn _main(opts: Opts) -> Result<()> { "{} Done in {} with an average of {} req/s", SUCCESS.to_string().green(), HumanDuration(watch.elapsed()).to_string().bold(), - ((words.len() * *depth.lock()) as f64 / watch.elapsed().as_secs_f64()) - .round() - .to_string() - .bold() + ((match mode { + Mode::Recursive => words.len() * *current_depth.lock(), + Mode::Classic => words.len(), + Mode::Permutations => words.len().pow(url.matches(REPLACE_KEYWORD).count() as u32), + }) as f64 + / watch.elapsed().as_secs_f64()) + .round() + .to_string() + .bold() ); let root = tree.lock().root.clone().unwrap().clone(); @@ -289,7 +356,7 @@ pub async fn _main(opts: Opts) -> Result<()> { tokio::fs::remove_file(opts.save_file.clone().unwrap()).await?; } if opts.output.is_some() { - let res = save_to_file(&opts, root, depth, tree); + let res = save_to_file(&opts, root, current_depth, tree); match res { Ok(_) => info!("Saved to {}", opts.output.unwrap().bold()), diff --git a/src/runner/classic.rs b/src/runner/classic.rs new file mode 100644 index 0000000..3da6534 --- /dev/null +++ b/src/runner/classic.rs @@ -0,0 +1,198 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use crate::{ + cli::opts::Opts, + utils::{ + constants::{ERROR, PROGRESS_CHARS, PROGRESS_TEMPLATE, REPLACE_KEYWORD, SUCCESS, WARNING}, + tree::{Tree, TreeData}, + }, +}; +use anyhow::Result; +use colored::Colorize; +use indicatif::ProgressBar; +use log::info; +use parking_lot::Mutex; +use serde_json::json; +use tokio::task::JoinHandle; +use url::Url; + +pub async fn run( + url: String, + opts: Opts, // The options passed to the program + tree: Arc>>, // The tree to be populated + words: Vec, // Each chunk is a list of strings to be passed to individual threads + threads: usize, // The number of threads to use +) -> Result<()> { + let spinner = ProgressBar::new_spinner(); + spinner.set_message(format!("Generating URLs...")); + spinner.enable_steady_tick(Duration::from_millis(100)); + + let urls: Vec = words + .clone() + .iter() + .map(|c| { + let mut url = url.clone(); + url = url.replace(REPLACE_KEYWORD, c); + url + }) + .collect(); + spinner.finish_and_clear(); + info!("Generated {} URLs", urls.clone().len().to_string().bold()); + + let mut handles = Vec::>::new(); + let progress = ProgressBar::new(urls.len() as u64).with_style( + indicatif::ProgressStyle::default_bar() + .template(PROGRESS_TEMPLATE)? + .progress_chars(PROGRESS_CHARS), + ); + let chunks = urls + .chunks(urls.clone().len() / threads) + .collect::>(); + + let client = super::client::build(&opts)?; + + for chunk in &chunks { + let chunk = chunk.to_vec(); + let client = client.clone(); + let progress = progress.clone(); + let opts = opts.clone(); + let tree = tree.clone(); + let handle = tokio::spawn(async move { + for url in &chunk { + let sender = super::client::get_sender(&opts, &url, &client); + + let t1 = Instant::now(); + + let response = sender.send().await; + + if let Some(throttle) = opts.throttle { + if throttle > 0 { + let elapsed = t1.elapsed(); + let sleep_duration = Duration::from_secs_f64(1.0 / throttle as f64); + if let Some(sleep) = sleep_duration.checked_sub(elapsed) { + tokio::time::sleep(sleep).await; + } + } + } + match response { + Ok(mut response) => { + let status_code = response.status().as_u16(); + let mut text = String::new(); + while let Ok(chunk) = response.chunk().await { + if let Some(chunk) = chunk { + text.push_str(&String::from_utf8_lossy(&chunk)); + } else { + break; + } + } + let filtered = super::filters::check( + &opts, + &text, + status_code, + t1.elapsed().as_millis(), + ); + + if filtered { + let additions = super::filters::parse_show(&opts, &text, &response); + + progress.println(format!( + "{} {} {} {}{}", + if response.status().is_success() { + SUCCESS.to_string().green() + } else if response.status().is_redirection() { + WARNING.to_string().yellow() + } else { + ERROR.to_string().red() + }, + response.status().as_str().bold(), + url, + format!("{}ms", t1.elapsed().as_millis().to_string().bold()) + .dimmed(), + additions.iter().fold("".to_string(), |acc, addition| { + format!( + "{} | {}: {}", + acc, + addition.key.dimmed().bold(), + addition.value.dimmed() + ) + }) + )); + + let parsed = Url::parse(url).unwrap(); + let mut tree = tree.lock().clone(); + let root_url = tree.root.clone().unwrap().lock().data.url.clone(); + tree.insert( + TreeData { + url: url.clone(), + depth: 0, + path: parsed.path().to_string().replace( + Url::parse(&root_url).unwrap().path().to_string().as_str(), + "", + ), + status_code, + extra: json!(additions), + }, + tree.root.clone(), + ); + } + } + Err(err) => { + if err.is_timeout() { + progress.println(format!( + "{} {} {}", + ERROR.to_string().red(), + "Timeout reached".bold(), + url + )); + } else if err.is_redirect() { + progress.println(format!( + "{} {} {} {}", + WARNING.to_string().yellow(), + "Redirect limit reached".bold(), + url, + "Check --follow-redirects".dimmed() + )); + } else if err.is_connect() { + progress.println(format!( + "{} {} {}", + ERROR.to_string().red(), + "Connection error".bold(), + url + )); + } else if err.is_request() { + progress.println(format!( + "{} {} {} {}", + ERROR.to_string().red(), + "Request error".bold(), + url, + format!("({})", err).dimmed() + )); + } else { + progress.println(format!( + "{} {} {} {}", + ERROR.to_string().red(), + "Unknown Error".bold(), + url, + format!("({})", err).dimmed() + )); + } + } + } + + progress.inc(1); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.await?; + } + + progress.finish_and_clear(); + + Ok(()) +} diff --git a/src/runner/mod.rs b/src/runner/mod.rs index b0c95fc..4425083 100644 --- a/src/runner/mod.rs +++ b/src/runner/mod.rs @@ -1,4 +1,6 @@ +pub mod classic; pub mod client; pub mod filters; -pub mod start; +pub mod permutations; +pub mod recursive; pub mod wordlists; diff --git a/src/runner/permutations.rs b/src/runner/permutations.rs new file mode 100644 index 0000000..ea7dcb8 --- /dev/null +++ b/src/runner/permutations.rs @@ -0,0 +1,201 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use crate::{ + cli::opts::Opts, + utils::{ + constants::{ERROR, PROGRESS_CHARS, PROGRESS_TEMPLATE, REPLACE_KEYWORD, SUCCESS, WARNING}, + tree::{Tree, TreeData}, + }, +}; +use anyhow::Result; +use colored::Colorize; +use indicatif::ProgressBar; +use itertools::Itertools; +use log::info; +use parking_lot::Mutex; +use serde_json::json; +use tokio::task::JoinHandle; +use url::Url; + +pub async fn run( + url: String, + opts: Opts, // The options passed to the program + tree: Arc>>, // The tree to be populated + words: Vec, // Each chunk is a list of strings to be passed to individual threads + threads: usize, // The number of threads to use +) -> Result<()> { + let token_count = url.matches(REPLACE_KEYWORD).count(); + let spinner = ProgressBar::new_spinner(); + spinner.set_message(format!("Generating URLs...")); + spinner.enable_steady_tick(Duration::from_millis(100)); + let combinations: Vec<_> = words.iter().permutations(token_count).collect(); + + let urls: Vec = combinations + .clone() + .iter() + .map(|c| { + let mut url = url.clone(); + for word in c { + url = url.replace(REPLACE_KEYWORD, word); + } + url + }) + .collect(); + spinner.finish_and_clear(); + info!("Generated {} URLs", urls.clone().len().to_string().bold()); + + let mut handles = Vec::>::new(); + let progress = ProgressBar::new(urls.len() as u64).with_style( + indicatif::ProgressStyle::default_bar() + .template(PROGRESS_TEMPLATE)? + .progress_chars(PROGRESS_CHARS), + ); + let chunks = urls + .chunks(urls.clone().len() / threads) + .collect::>(); + let client = super::client::build(&opts)?; + + for chunk in &chunks { + let chunk = chunk.to_vec(); + let tree = tree.clone(); + let client = client.clone(); + let progress = progress.clone(); + let opts = opts.clone(); + let handle = tokio::spawn(async move { + for url in &chunk { + let sender = super::client::get_sender(&opts, &url, &client); + + let t1 = Instant::now(); + + let response = sender.send().await; + + if let Some(throttle) = opts.throttle { + if throttle > 0 { + let elapsed = t1.elapsed(); + let sleep_duration = Duration::from_secs_f64(1.0 / throttle as f64); + if let Some(sleep) = sleep_duration.checked_sub(elapsed) { + tokio::time::sleep(sleep).await; + } + } + } + match response { + Ok(mut response) => { + let status_code = response.status().as_u16(); + let mut text = String::new(); + while let Ok(chunk) = response.chunk().await { + if let Some(chunk) = chunk { + text.push_str(&String::from_utf8_lossy(&chunk)); + } else { + break; + } + } + let filtered = super::filters::check( + &opts, + &text, + status_code, + t1.elapsed().as_millis(), + ); + + if filtered { + let additions = super::filters::parse_show(&opts, &text, &response); + + progress.println(format!( + "{} {} {} {}{}", + if response.status().is_success() { + SUCCESS.to_string().green() + } else if response.status().is_redirection() { + WARNING.to_string().yellow() + } else { + ERROR.to_string().red() + }, + response.status().as_str().bold(), + url, + format!("{}ms", t1.elapsed().as_millis().to_string().bold()) + .dimmed(), + additions.iter().fold("".to_string(), |acc, addition| { + format!( + "{} | {}: {}", + acc, + addition.key.dimmed().bold(), + addition.value.dimmed() + ) + }) + )); + + let parsed = Url::parse(url).unwrap(); + let mut tree = tree.lock().clone(); + let root_url = tree.root.clone().unwrap().lock().data.url.clone(); + tree.insert( + TreeData { + url: url.clone(), + depth: 0, + path: parsed.path().to_string().replace( + Url::parse(&root_url).unwrap().path().to_string().as_str(), + "", + ), + status_code, + extra: json!(additions), + }, + tree.root.clone(), + ); + } + } + Err(err) => { + if err.is_timeout() { + progress.println(format!( + "{} {} {}", + ERROR.to_string().red(), + "Timeout reached".bold(), + url + )); + } else if err.is_redirect() { + progress.println(format!( + "{} {} {} {}", + WARNING.to_string().yellow(), + "Redirect limit reached".bold(), + url, + "Check --follow-redirects".dimmed() + )); + } else if err.is_connect() { + progress.println(format!( + "{} {} {}", + ERROR.to_string().red(), + "Connection error".bold(), + url + )); + } else if err.is_request() { + progress.println(format!( + "{} {} {} {}", + ERROR.to_string().red(), + "Request error".bold(), + url, + format!("({})", err).dimmed() + )); + } else { + progress.println(format!( + "{} {} {} {}", + ERROR.to_string().red(), + "Unknown Error".bold(), + url, + format!("({})", err).dimmed() + )); + } + } + } + + progress.inc(1); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.await?; + } + + progress.finish_and_clear(); + Ok(()) +} diff --git a/src/runner/start.rs b/src/runner/recursive.rs similarity index 100% rename from src/runner/start.rs rename to src/runner/recursive.rs diff --git a/src/utils/constants.rs b/src/utils/constants.rs index 8f5d0d6..97cf632 100644 --- a/src/utils/constants.rs +++ b/src/utils/constants.rs @@ -14,3 +14,4 @@ pub const SAVE_FILE: &str = ".rwalk.json"; pub const STATUS_CODES: &str = "200-299,301,302,307,401,403,405,500"; pub const PROGRESS_TEMPLATE: &str = "{spinner:.blue} (ETA. {eta}) {wide_bar} {pos}/{len} ({per_sec:>11}) | {prefix:>3} {msg:>14.bold}"; pub const PROGRESS_CHARS: &str = "█▉▊▋▌▍▎▏░"; +pub const REPLACE_KEYWORD: &str = "$"; diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 2fcaf46..056f49b 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -11,6 +11,7 @@ use crate::utils::{ pub mod constants; pub mod logger; +pub mod structs; pub mod tree; pub fn banner() { diff --git a/src/utils/structs.rs b/src/utils/structs.rs new file mode 100644 index 0000000..3c61753 --- /dev/null +++ b/src/utils/structs.rs @@ -0,0 +1,16 @@ +pub enum Mode { + Recursive, + Permutations, + Classic, +} + +impl From<&str> for Mode { + fn from(s: &str) -> Self { + match s { + "recursive" | "recursion" | "r" => Mode::Recursive, + "permutations" | "permutation" | "p" => Mode::Permutations, + "classic" | "c" => Mode::Classic, + _ => Mode::Recursive, + } + } +} diff --git a/src/utils/tree.rs b/src/utils/tree.rs index 83c05d0..6ccb424 100644 --- a/src/utils/tree.rs +++ b/src/utils/tree.rs @@ -80,6 +80,14 @@ impl Tree { } } } + + pub fn insert_datas(&mut self, datas: Vec) { + // Insert nodes into the root + let mut previous_node: Option>>> = self.root.clone(); + for data in datas { + previous_node = Some(self.insert(data, previous_node)); + } + } } impl TreeItem for TreeNode {