Skip to content

Commit

Permalink
Improve tree report to include url types
Browse files Browse the repository at this point in the history
  • Loading branch information
cestef committed Apr 17, 2024
1 parent 1122c66 commit 97f6f48
Show file tree
Hide file tree
Showing 5 changed files with 71 additions and 15 deletions.
4 changes: 2 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use parking_lot::Mutex;
use ptree::print_tree;
use tokio::{io::AsyncWriteExt, task::JoinHandle, time::timeout};
use url::Url;
use utils::structs::FuzzMatch;
use utils::{structs::FuzzMatch, tree::UrlType};

use crate::utils::{
constants::SUCCESS,
Expand Down Expand Up @@ -258,7 +258,7 @@ pub async fn _main(opts: Opts) -> Result<()> {
.to_string(),
status_code: 0,
extra: serde_json::Value::Null,
is_dir: true,
url_type: UrlType::Dir,
},
None,
);
Expand Down
24 changes: 19 additions & 5 deletions src/runner/classic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::{
cli::opts::Opts,
utils::{
constants::{ERROR, PROGRESS_CHARS, PROGRESS_TEMPLATE, SUCCESS, WARNING},
tree::{Tree, TreeData},
tree::{Tree, TreeData, UrlType},
},
};
use anyhow::{anyhow, Result};
Expand All @@ -21,7 +21,7 @@ use reqwest::Client;
use serde_json::json;
use url::Url;

use super::{wordlists::ParsedWordlist, Runner};
use super::{filters::is_directory, wordlists::ParsedWordlist, Runner};

pub struct Classic {
url: String,
Expand Down Expand Up @@ -153,6 +153,15 @@ impl Classic {
.data
.url
.clone();
let maybe_content_type = response.headers().get("content-type").map(|x| {
x.to_str()
.unwrap_or_default()
.split(';')
.next()
.unwrap_or_default()
.to_string()
});
let is_dir = is_directory(&response);
tree.insert(
TreeData {
url: url.clone(),
Expand All @@ -163,7 +172,13 @@ impl Classic {
),
status_code,
extra: json!(additions),
is_dir: false,
url_type: if is_dir {
UrlType::Dir
} else if let Some(content_type) = maybe_content_type {
UrlType::File(content_type)
} else {
UrlType::Unknown
},
},
tree.root.clone(),
);
Expand Down Expand Up @@ -200,8 +215,7 @@ impl Classic {
),
status_code: 0,
extra: json!([]),
//TODO: is_dir
is_dir: false,
url_type: UrlType::Unknown,
},
tree.root.clone(),
);
Expand Down
24 changes: 20 additions & 4 deletions src/runner/recursive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use crate::{
cli::opts::Opts,
utils::{
constants::{DEFAULT_DEPTH, ERROR, PROGRESS_CHARS, PROGRESS_TEMPLATE, SUCCESS, WARNING},
tree::{Tree, TreeData, TreeNode},
tree::{Tree, TreeData, TreeNode, UrlType},
},
};

Expand All @@ -40,7 +40,8 @@ impl super::Runner for Recursive {
let root_progress = MultiProgress::new();
// Spawn a thread for each previous node
for previous_node in &previous_nodes {
if !previous_node.lock().data.is_dir && !self.opts.force_recursion {
if previous_node.lock().data.url_type != UrlType::Dir && !self.opts.force_recursion
{
log::debug!("Skipping not-directory {}", previous_node.lock().data.url);
continue;
}
Expand Down Expand Up @@ -234,14 +235,29 @@ impl Recursive {
.iter()
.any(|child| child.lock().data.path == *word)
{
let maybe_content_type =
response.headers().get("content-type").map(|x| {
x.to_str()
.unwrap_or_default()
.split(';')
.next()
.unwrap_or_default()
.to_string()
});
tree.lock().insert(
TreeData {
url: url.clone(),
depth: data.depth + 1,
path: word.clone(),
status_code,
extra: json!(additions),
is_dir,
url_type: if is_dir {
UrlType::Dir
} else if let Some(content_type) = maybe_content_type {
UrlType::File(content_type)
} else {
UrlType::Unknown
},
},
Some(previous_node.clone()),
);
Expand Down Expand Up @@ -277,7 +293,7 @@ impl Recursive {
path: word.clone(),
status_code: 0,
extra: json!([]),
is_dir: false,
url_type: UrlType::Unknown,
},
Some(previous_node.clone()),
);
Expand Down
7 changes: 7 additions & 0 deletions src/utils/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,13 @@ pub fn build_opts_table(
]);
}

if opts.depth.is_some() {
builder.push_record(vec![
"Depth",
&opts.depth.as_ref().unwrap().to_string().bold(),
]);
}

if !opts.transform.is_empty() {
let mut wordlist_filters_builder = Builder::default();
wordlist_filters_builder.push_record(vec!["Wordlist", "Filter", "Value"]);
Expand Down
27 changes: 23 additions & 4 deletions src/utils/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,14 @@ pub struct TreeData {
pub path: String,
pub status_code: u16,
pub extra: Value,
pub is_dir: bool,
pub url_type: UrlType,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum UrlType {
Dir,
File(String),
Unknown,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
Expand Down Expand Up @@ -191,7 +198,7 @@ impl TreeItem for TreeNode<TreeData> {
let emoji = get_emoji_for_status_code_colored(self.data.status_code);
write!(
f,
"{}{} /{}",
"{}{} /{} ({})",
if self.data.status_code == 0 {
style.paint("🔍".to_string())
} else {
Expand All @@ -202,7 +209,19 @@ impl TreeItem for TreeNode<TreeData> {
} else {
style.paint(format!(" {}", self.data.status_code.to_string().dimmed()))
},
style.paint(&self.data.path.trim_start_matches('/'))
style.paint(&self.data.path.trim_start_matches('/')),
style.paint(
(if self.data.url_type == UrlType::Dir {
"dir".to_string()
} else {
match &self.data.url_type {
UrlType::File(ext) => format!("file.{}", ext),
UrlType::Unknown => "unknown".to_string(),
_ => "unknown".to_string(),
}
})
.dimmed()
)
)?;
Ok(())
}
Expand Down Expand Up @@ -316,7 +335,7 @@ mod tests {
path: "/test".to_string(),
status_code: 200,
extra: Value::Null,
is_dir: false,
url_type: UrlType::Dir,
},
children: vec![],
};
Expand Down

0 comments on commit 97f6f48

Please sign in to comment.