From ad731731b772abd1f3a1759d739d7d21c5b87a1f Mon Sep 17 00:00:00 2001 From: Eduard Tolosa Date: Sun, 27 Oct 2024 16:26:08 -0500 Subject: [PATCH] Performance improvements. --- src/args.rs | 4 ++-- src/httplib.rs | 41 +++++++++++++++++------------------------ src/main.rs | 36 ++++++++++-------------------------- src/structs.rs | 4 ++-- 4 files changed, 31 insertions(+), 54 deletions(-) diff --git a/src/args.rs b/src/args.rs index 23d35f8..4066099 100644 --- a/src/args.rs +++ b/src/args.rs @@ -10,8 +10,8 @@ pub struct Cli { /// Timeout in seconds pub timeout: u64, #[clap(short, long)] - /// Show HTTP status codes - pub show_codes: bool, + /// Show HTTP status codes, final URL and domain + pub show_full_data: bool, #[clap(short, long)] /// Domain to check - can be omitted if using stdin pub domain: Option, diff --git a/src/httplib.rs b/src/httplib.rs index f78ab81..ce35d45 100644 --- a/src/httplib.rs +++ b/src/httplib.rs @@ -15,7 +15,7 @@ use { }; #[must_use] -pub async fn return_http_data(options: &LibOptions) -> HashMap { +pub async fn return_http_data(options: &LibOptions, from_cli: bool) -> HashMap { let threads = options.hosts.len().min(options.threads); let filter_codes = options.filter_codes.as_deref().unwrap_or_default(); @@ -25,7 +25,10 @@ pub async fn return_http_data(options: &LibOptions) -> HashMap let user_agent = utils::return_random_user_agent(&options.user_agents); async move { - let mut http_data = HttpData::default(); + let mut http_data = HttpData { + checked_host: host.clone(), + ..Default::default() + }; let mut response = None; // Attempt both HTTPS and HTTP requests, retry if necessary @@ -50,33 +53,32 @@ pub async fn return_http_data(options: &LibOptions) -> HashMap } if let Some(resp) = response { - http_data.checked_host = host.clone(); + // Those are always set + http_data.protocol = resp.url().scheme().to_string(); + http_data.status_code = resp.status().as_u16(); http_data.final_url = resp.url().to_string(); - if options.assign_response_data { - http_data = assign_response_data(http_data, resp, options.return_filters).await; - } else { - http_data.status_code = resp.status().as_u16(); - http_data.http_status = "ACTIVE".to_string(); + if !from_cli { + assign_response_data(&mut http_data, resp, options.return_filters).await; } } else { http_data.http_status = "INACTIVE".to_string(); } if !options.quiet_flag - && !http_data.checked_host.is_empty() + && !http_data.final_url.is_empty() && (filter_codes.is_empty() || filter_codes.contains(&http_data.status_code.to_string())) && (exclude_codes.is_empty() || !exclude_codes.contains(&http_data.status_code.to_string())) { - if options.show_status_codes { + if options.show_full_data { println!( "{},[{}],[{}]", http_data.checked_host, http_data.final_url, http_data.status_code ); } else { - println!("{},[{}]", http_data.checked_host, http_data.final_url); + println!("{}://{}", http_data.protocol, http_data.checked_host); } } (host, http_data) @@ -105,18 +107,11 @@ pub fn return_http_client(timeout: u64, max_redirects: usize) -> Client { } #[allow(clippy::field_reassign_with_default)] -pub async fn assign_response_data( - mut http_data: HttpData, - resp: Response, - return_filters: bool, -) -> HttpData { +pub async fn assign_response_data(http_data: &mut HttpData, resp: Response, return_filters: bool) { let headers = resp.headers().clone(); let url = resp.url().clone(); http_data.http_status = "ACTIVE".to_string(); - http_data.status_code = resp.status().as_u16(); - http_data.final_url = url.to_string(); - http_data.protocol = url.scheme().to_string(); http_data.content_type = headers .get(CONTENT_TYPE) .and_then(|v| v.to_str().ok()) @@ -132,10 +127,11 @@ pub async fn assign_response_data( http_data.headers = format!("{headers:?}"); - return_title_and_body(&mut http_data, &full_body); + return_title_and_body(http_data, &full_body); http_data.words_count = full_body.split_whitespace().count(); http_data.lines = full_body.lines().count() + 1; + http_data.points_to_another_host = url.host_str() != Some(&http_data.checked_host); if return_filters { let host = url.host_str().unwrap_or_default(); @@ -143,8 +139,6 @@ pub async fn assign_response_data( let user_agents = utils::user_agents(); http_data.bad_data = return_filters_data(host, client, user_agents).await; } - - http_data } pub fn return_title_and_body(http_data: &mut HttpData, body: &str) { @@ -211,12 +205,11 @@ pub async fn return_filters_data( user_agents: user_agents_list, retries: 1, threads, - assign_response_data: true, quiet_flag: true, ..Default::default() }; - let data = return_http_data(&lib_options).await; + let data = return_http_data(&lib_options, false).await; data.values() .map(|http_data| { diff --git a/src/main.rs b/src/main.rs index 788e81c..ab91ddb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,19 +10,7 @@ async fn main() -> Result<(), Box> { // Eval args let args = args::Cli::parse(); - let filter_codes = args.filter_codes; - let exclude_codes = args.exclude_codes; - let threads = args.threads; - let retries = args.retries; - let timeout = args.timeout; - let max_redirects = args.max_redirects; - let user_agents_list = utils::user_agents(); - let show_status_codes = args.show_codes; - - let client = httplib::return_http_client(timeout, max_redirects); - let mut buffer = String::new(); - let mut hosts = HashSet::new(); if args.domain.is_some() { @@ -43,25 +31,21 @@ async fn main() -> Result<(), Box> { let lib_options = LibOptions { hosts, - client, - user_agents: user_agents_list, - retries, - threads, - filter_codes, - exclude_codes, - show_status_codes, + client: httplib::return_http_client(args.timeout, args.max_redirects), + user_agents: utils::user_agents(), + retries: args.retries, + threads: args.threads, + filter_codes: args.filter_codes, + exclude_codes: args.exclude_codes, + show_full_data: args.show_full_data, ..Default::default() }; - if !args.quiet { - if show_status_codes { - println!("DOMAIN,[FINAL_URL],[STATUS_CODE]"); - } else { - println!("DOMAIN,[FINAL_URL]"); - } + if !args.quiet && args.show_full_data { + println!("DOMAIN,[FINAL_URL],[STATUS_CODE]"); } - let _ = httplib::return_http_data(&lib_options).await; + let _ = httplib::return_http_data(&lib_options, true).await; Ok(()) } diff --git a/src/structs.rs b/src/structs.rs index 9228890..f16a277 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -25,6 +25,7 @@ pub struct HttpData { pub bad_data: HTTPFilters, pub html_file_path: String, pub screenshot_data: Vec, + pub points_to_another_host: bool, } #[derive(Clone, Debug, Default)] @@ -37,7 +38,6 @@ pub struct LibOptions { pub return_filters: bool, pub filter_codes: Option, pub exclude_codes: Option, - pub show_status_codes: bool, - pub assign_response_data: bool, + pub show_full_data: bool, pub quiet_flag: bool, }