From d514dd5a1ceb41d415da8cb59d641cae860b2d2f Mon Sep 17 00:00:00 2001
From: delorean <jackdelmar@protonmail.com>
Date: Thu, 23 May 2024 18:34:03 -0500
Subject: [PATCH] added body reading module

---
 src/common/conf.rs    |  4 ++++
 src/common/console.rs | 32 +++++++++++++++++++++++---------
 src/common/exec.rs    |  3 ++-
 src/common/modules.rs | 18 +++++++++++++++++-
 src/common/net.rs     |  7 ++++++-
 5 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/src/common/conf.rs b/src/common/conf.rs
index b2d366c..c596494 100644
--- a/src/common/conf.rs
+++ b/src/common/conf.rs
@@ -37,6 +37,10 @@ pub struct Config {
     #[clap(long = "redirects")]
     /// follow redirects
     pub follow: bool,
+
+    #[clap(long = "body")]
+    /// read n bytes of the response document body
+    pub bodysize: usize,
 }
 
 pub fn load() -> Config {
diff --git a/src/common/console.rs b/src/common/console.rs
index 3757aee..1ecef7e 100644
--- a/src/common/console.rs
+++ b/src/common/console.rs
@@ -8,6 +8,18 @@ pub fn fatal(msg: &str) -> ! {
     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 {
     match code {
         200..=299 => code.to_string().green(),
@@ -19,17 +31,19 @@ pub fn fmtcode(code: u16) -> ColoredString {
 }
 
 pub fn parsetitle(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);
-    });
+    let title = fmtwhitespace(s);
 
     format!("{}{}{}",
-    "[".bright_black().bold(),
-    out.bright_cyan().bold(),
+    "title[".bright_black().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()
     )
 }
diff --git a/src/common/exec.rs b/src/common/exec.rs
index 182f4de..b682ac1 100644
--- a/src/common/exec.rs
+++ b/src/common/exec.rs
@@ -30,7 +30,8 @@ pub async fn takeoff(args: Config, params: Params) {
                         .trim(),
                     scodes,
                     params.exclude,
-                    args.pulltitles
+                    args.pulltitles,
+                    args.bodysize
                 )
                 .await;
             }
diff --git a/src/common/modules.rs b/src/common/modules.rs
index 144710f..886bbd2 100644
--- a/src/common/modules.rs
+++ b/src/common/modules.rs
@@ -1,6 +1,6 @@
 use select::{document::Document, predicate::Name};
 
-use super::console::parsetitle;
+use super::console::{parsetitle, parsebody, fmtwhitespace};
 
 pub fn get_title(body: &String) -> String {
     let document = Document::from(body.as_str());
@@ -12,4 +12,20 @@ pub fn get_title(body: &String) -> String {
     .unwrap_or_else(|| "".to_string());
     
     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)
 }
\ No newline at end of file
diff --git a/src/common/net.rs b/src/common/net.rs
index b2cf913..6c368cf 100644
--- a/src/common/net.rs
+++ b/src/common/net.rs
@@ -1,7 +1,7 @@
 use reqwest::{redirect::Policy, Client};
 use std::time::Duration;
 
-use super::console::{parsehit, parsetitle};
+use super::console::parsehit;
 use super::modules::*;
 
 pub fn mkclient(redir: bool) -> Result<Client, reqwest::Error> {
@@ -30,6 +30,7 @@ pub async fn query(
     codes: Vec<u16>,
     exclude: bool,
     titles: bool,
+    bodysize: usize,
 ) -> Result<(), reqwest::Error> {
     let response: reqwest::Response;
     if let Ok(res) = sendreq(&c, true, url).await {
@@ -60,6 +61,10 @@ pub async fn query(
         out = format!("{} {}", out, get_title(&body));
     }
 
+    if bodysize > 0 {
+        out = format!("{} {}", out, read_body(&body, bodysize));
+    }
+
     println!("{}", out);
 
     Ok(())