added body reading module

This commit is contained in:
delorean 2024-05-23 18:34:03 -05:00
parent cd6e57f15b
commit d514dd5a1c
5 changed files with 52 additions and 12 deletions

View File

@ -37,6 +37,10 @@ pub struct Config {
#[clap(long = "redirects")] #[clap(long = "redirects")]
/// follow redirects /// follow redirects
pub follow: bool, pub follow: bool,
#[clap(long = "body")]
/// read n bytes of the response document body
pub bodysize: usize,
} }
pub fn load() -> Config { pub fn load() -> Config {

View File

@ -8,6 +8,18 @@ pub fn fatal(msg: &str) -> ! {
process::exit(-1); process::exit(-1);
} }
pub fn fmtwhitespace(s: String) -> String {
let mut out = String::with_capacity(s.len());
s.split_whitespace().for_each(|w| {
if !out.is_empty() {
out.push(' ');
}
out.push_str(w);
});
out
}
pub fn fmtcode(code: u16) -> ColoredString { pub fn fmtcode(code: u16) -> ColoredString {
match code { match code {
200..=299 => code.to_string().green(), 200..=299 => code.to_string().green(),
@ -19,17 +31,19 @@ pub fn fmtcode(code: u16) -> ColoredString {
} }
pub fn parsetitle(s: String) -> String { pub fn parsetitle(s: String) -> String {
let mut out = String::with_capacity(s.len()); let title = fmtwhitespace(s);
s.split_whitespace().for_each(|w| {
if !out.is_empty() {
out.push(' ');
}
out.push_str(w);
});
format!("{}{}{}", format!("{}{}{}",
"[".bright_black().bold(), "title[".bright_black().bold(),
out.bright_cyan().bold(), title.bright_cyan().bold(),
"]".bright_black().bold()
)
}
pub fn parsebody(s: String) -> String {
format!("{}{}{}",
"body[".bright_black().bold(),
s.bright_magenta().bold(),
"]".bright_black().bold() "]".bright_black().bold()
) )
} }

View File

@ -30,7 +30,8 @@ pub async fn takeoff(args: Config, params: Params) {
.trim(), .trim(),
scodes, scodes,
params.exclude, params.exclude,
args.pulltitles args.pulltitles,
args.bodysize
) )
.await; .await;
} }

View File

@ -1,6 +1,6 @@
use select::{document::Document, predicate::Name}; use select::{document::Document, predicate::Name};
use super::console::parsetitle; use super::console::{parsetitle, parsebody, fmtwhitespace};
pub fn get_title(body: &String) -> String { pub fn get_title(body: &String) -> String {
let document = Document::from(body.as_str()); let document = Document::from(body.as_str());
@ -12,4 +12,20 @@ pub fn get_title(body: &String) -> String {
.unwrap_or_else(|| "".to_string()); .unwrap_or_else(|| "".to_string());
parsetitle(title) parsetitle(title)
}
pub fn read_body(body: &String, lim: usize) -> String {
let document = Document::from(body.as_str());
let mut text = String::new();
for tag in document.find(Name("body")) {
text.push_str(&tag.text());
}
text = fmtwhitespace(text);
if text.len() > lim {
text.truncate(lim);
}
parsebody(text)
} }

View File

@ -1,7 +1,7 @@
use reqwest::{redirect::Policy, Client}; use reqwest::{redirect::Policy, Client};
use std::time::Duration; use std::time::Duration;
use super::console::{parsehit, parsetitle}; use super::console::parsehit;
use super::modules::*; use super::modules::*;
pub fn mkclient(redir: bool) -> Result<Client, reqwest::Error> { pub fn mkclient(redir: bool) -> Result<Client, reqwest::Error> {
@ -30,6 +30,7 @@ pub async fn query(
codes: Vec<u16>, codes: Vec<u16>,
exclude: bool, exclude: bool,
titles: bool, titles: bool,
bodysize: usize,
) -> Result<(), reqwest::Error> { ) -> Result<(), reqwest::Error> {
let response: reqwest::Response; let response: reqwest::Response;
if let Ok(res) = sendreq(&c, true, url).await { if let Ok(res) = sendreq(&c, true, url).await {
@ -60,6 +61,10 @@ pub async fn query(
out = format!("{} {}", out, get_title(&body)); out = format!("{} {}", out, get_title(&body));
} }
if bodysize > 0 {
out = format!("{} {}", out, read_body(&body, bodysize));
}
println!("{}", out); println!("{}", out);
Ok(()) Ok(())