Performance improvements.

This commit is contained in:
Eduard Tolosa 2024-10-27 16:26:08 -05:00
parent ede2c96e0f
commit ad731731b7
4 changed files with 31 additions and 54 deletions

View file

@ -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<String>,

View file

@ -15,7 +15,7 @@ use {
};
#[must_use]
pub async fn return_http_data(options: &LibOptions) -> HashMap<String, HttpData> {
pub async fn return_http_data(options: &LibOptions, from_cli: bool) -> HashMap<String, HttpData> {
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<String, HttpData>
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<String, HttpData>
}
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| {

View file

@ -10,19 +10,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// 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<dyn std::error::Error>> {
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(())
}

View file

@ -25,6 +25,7 @@ pub struct HttpData {
pub bad_data: HTTPFilters,
pub html_file_path: String,
pub screenshot_data: Vec<u8>,
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<String>,
pub exclude_codes: Option<String>,
pub show_status_codes: bool,
pub assign_response_data: bool,
pub show_full_data: bool,
pub quiet_flag: bool,
}