Merge pull request #84 from rusty-sec/match_body

New body matcher
CyAxe · Feb 16, 2023 · 43fcf1f · 43fcf1f
2 parents 132cf3e + c541576
commit 43fcf1f
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 21 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -20,7 +20,7 @@ pub mod cli;
 pub mod lua;
 
 use cli::{
-    bar::{show_msg, MessageLevel, BAR},
+    bar::{show_msg, MessageLevel, BAR, create_progress},
     errors::CliErrors,
     input::load_scripts::{get_scripts, valid_scripts},
 };
@@ -108,6 +108,7 @@ impl Lotus {
                 loaded_scripts
             }
         };
+        {create_progress(loaded_scripts.len() as u64)};
         if self.output.is_none() {
             show_msg("Output argument is missing", MessageLevel::Error);
             std::process::exit(1);

diff --git a/src/lua/network/http.rs b/src/lua/network/http.rs
@@ -93,6 +93,23 @@ impl Sender {
         body: String,
         headers: HeaderMap,
     ) -> Result<HttpResponse, mlua::Error> {
+        {
+            let req_limit = REQUESTS_LIMIT.lock().unwrap();
+            let mut req_sent = REQUESTS_SENT.lock().unwrap();
+            if *req_sent >= *req_limit {
+                let sleep_time = SLEEP_TIME.lock().unwrap();
+                let bar = BAR.lock().unwrap();
+                bar.println(format!(
+                    "The rate limit for requests has been raised, please wait {} seconds ",
+                    *sleep_time
+                ));
+                log::debug!("{}",format!("The rate limit for requests has been raised, please wait {} seconds ",*sleep_time));
+                std::thread::sleep(Duration::from_secs(*sleep_time));
+                *req_sent = 1;
+                bar.println("Continue ...");
+                log::debug!("changing req_sent value to 1");
+            }
+        };
         match self
             .build_client()
             .unwrap()
@@ -105,21 +122,7 @@ impl Sender {
             Ok(resp) => {
                 // Locking Scope
                 {
-                    let req_limit = REQUESTS_LIMIT.lock().unwrap();
                     let mut req_sent = REQUESTS_SENT.lock().unwrap();
-                    if *req_sent >= *req_limit {
-                        let sleep_time = SLEEP_TIME.lock().unwrap();
-                        let bar = BAR.lock().unwrap();
-                        bar.println(format!(
-                            "The rate limit for requests has been raised, please wait {} seconds ",
-                            *sleep_time
-                        ));
-                        log::debug!("{}",format!("The rate limit for requests has been raised, please wait {} seconds ",*sleep_time));
-                        tokio::time::sleep(Duration::from_secs(*sleep_time)).await;
-                        *req_sent = 0;
-                        bar.println("Continue ...");
-                        log::debug!("changing req_sent value to 0");
-                    }
                     *req_sent += 1;
                 };
                 let mut resp_headers: HashMap<String, String> = HashMap::new();

diff --git a/src/lua/parsing/text.rs b/src/lua/parsing/text.rs
@@ -1,13 +1,11 @@
-use crate::lua::network::http::HttpResponse;
 use mlua::UserData;
 use tealr::TypeName;
 
 #[derive(TypeName, Debug)]
 pub struct ResponseMatcher {}
 
 impl ResponseMatcher {
-    pub fn match_and_body(&self, response: HttpResponse, text: Vec<String>) -> bool {
-        let body = response.body;
+    pub fn match_and_body(&self, body: String, text: Vec<String>) -> bool {
         let mut counter = 0;
         text.iter().for_each(|x| {
             if body.contains(x) {
@@ -20,15 +18,31 @@ impl ResponseMatcher {
             false
         }
     }
+    pub fn match_once_body(&self, body: String, text: Vec<String>) -> String {
+        let mut matched_data = "".into();
+        text.iter().for_each(|x| {
+            if body.contains(x) {
+                matched_data = x.to_string();
+            }
+        });
+        matched_data
+    }
 }
 
 impl UserData for ResponseMatcher {
     fn add_methods<'lua, M: mlua::UserDataMethods<'lua, Self>>(methods: &mut M) {
         methods.add_method(
             "match_body",
-            |_, this, (response, text_list): (HttpResponse, Vec<String>)| {
+            |_, this, (response, text_list): (String, Vec<String>)| {
                 Ok(this.match_and_body(response, text_list))
             },
         );
+        methods.add_method(
+            "match_body_once",
+            |_, this, (response, text_list): (String, Vec<String>)| {
+                let is_match = this.match_once_body(response, text_list);
+                Ok(is_match)
+            },
+        )
     }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -19,7 +19,7 @@
 use lotus::{
     cli::{
         args::Opts,
-        bar::create_progress,
+        bar::{BAR,create_progress},
         startup::{new::new_args, urls::args_urls},
     },
     lua::{
@@ -37,7 +37,7 @@ async fn main() -> Result<(), std::io::Error> {
         Opts::URLS { .. } => {
             let opts = args_urls();
             // Open two threads for URL/HOST scanning
-            create_progress((opts.target_data.urls.len() * opts.target_data.hosts.len()) as u64);
+            create_progress((opts.target_data.urls.len() * opts.target_data.hosts.len() * opts.target_data.paths.len()) as u64);
             *SLEEP_TIME.lock().unwrap() = opts.delay;
             *REQUESTS_LIMIT.lock().unwrap() = opts.requests_limit;
             let scan_futures = vec![
@@ -61,6 +61,7 @@ async fn main() -> Result<(), std::io::Error> {
                 ),
             ];
             runner::scan_futures(scan_futures, 3, None).await;
+            BAR.lock().unwrap().finish();
         }
         Opts::NEW {
             scan_type,