Merge pull request #77 from rusty-sec/http_limit

adding requests limit and delay

src/cli/args.rs

@@ -127,6 +127,14 @@ pub enum Opts {
             help = "Set http proxy for all connections"
         )]
         proxy: Option<String>,
+        #[structopt(
+            long = "requests-limit",
+            help = "requests limit",
+            default_value = "2000"
+            )]
+        requests_limit: i32,
+        #[structopt(long = "delay", help = "sleeping dalay", default_value = "5")]
+        delay: u64,
 
         #[structopt(long = "log", help = "Saving Lotus Logs for debugging")]
         log: Option<PathBuf>,

src/cli/startup/urls.rs

@@ -14,6 +14,8 @@ pub struct UrlArgs {
     pub exit_after: i32,
     pub req_opts: RequestOpts,
     pub lotus_obj: Lotus,
+    pub requests_limit: i32,
+    pub delay: u64
 }
 
 pub struct TargetData {
@@ -22,7 +24,7 @@ pub struct TargetData {
 }
 
 pub fn args_urls() -> UrlArgs {
-    let (urls, hosts, exit_after, req_opts, lotus_obj) = match Opts::from_args() {
+    let (urls, hosts, exit_after, req_opts, lotus_obj, requests_limit, delay) = match Opts::from_args() {
         Opts::URLS {
             redirects,
             workers,
@@ -35,6 +37,8 @@ pub fn args_urls() -> UrlArgs {
             urls,
             headers,
             exit_after,
+            requests_limit,
+            delay
         } => {
             // setup logger
             init_log(log).unwrap();
@@ -70,7 +74,7 @@ pub fn args_urls() -> UrlArgs {
                 .map(|url| url.to_string())
                 .collect::<Vec<String>>();
             let hosts = get_target_hosts(urls_vec.clone());
-            (urls_vec, hosts, exit_after, req_opts, lotus_obj)
+            (urls_vec, hosts, exit_after, req_opts, lotus_obj, requests_limit, delay)
         }
         _ => {
             std::process::exit(1);
@@ -82,5 +86,7 @@ pub fn args_urls() -> UrlArgs {
         exit_after,
         req_opts,
         lotus_obj,
+        requests_limit,
+        delay
     }
 }

src/lua/network/http.rs

@@ -16,22 +16,22 @@
  * limitations under the license.
  */
 
+use crate::BAR;
 use reqwest::{header::HeaderMap, redirect, Client, Method, Proxy};
-use std::{collections::HashMap, time::Duration};
+use std::collections::HashMap;
 mod http_lua_api;
 pub use http_lua_api::Sender;
 use mlua::ExternalError;
+use lazy_static::lazy_static;
+use std::thread::sleep;
+use std::time::Duration;
+use std::sync::{Arc,Mutex};
 use tealr::{mlu::FromToLua, TypeName};
 
-/// RespType for lua userdata
-#[derive(FromToLua, Clone, Debug, TypeName)]
-pub enum RespType {
-    NoErrors,
-    Emtpy,
-    Str(String),
-    Int(i32),
-    Headers(HashMap<String, String>),
-    Error(String),
+lazy_static!{
+    pub static ref REQUESTS_LIMIT: Arc<Mutex<i32>> = Arc::new(Mutex::new(5));
+    pub static ref REQUESTS_SENT: Arc<Mutex<i32>> = Arc::new(Mutex::new(0));
+    pub static ref SLEEP_TIME: Arc<Mutex<u64>> = Arc::new(Mutex::new(5));
 }
 
 #[derive(Debug, FromToLua, TypeName)]
@@ -104,6 +104,19 @@ impl Sender {
             .await
         {
             Ok(resp) => {
+                let req_limit = REQUESTS_LIMIT.lock().unwrap();
+                let mut req_sent = REQUESTS_SENT.lock().unwrap();
+                *req_sent += 1;
+                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));
+                    sleep(Duration::from_secs(*sleep_time));
+                    *req_sent = 0;
+                    bar.println("Continue ...");
+                    log::debug!("changing req_sent value to 0");
+                }
                 let mut resp_headers: HashMap<String, String> = HashMap::new();
                 resp.headers()
                     .iter()
@@ -121,9 +134,12 @@ impl Sender {
                     body,
                     headers: resp_headers,
                 };
+
                 Ok(resp_data_struct)
             }
-            Err(err) => Err(err.to_lua_err()),
+            Err(err) => {
+                Err(err.to_lua_err())
+            },
         }
     }
 }

src/main.rs

@@ -22,7 +22,10 @@ use lotus::{
         bar::create_progress,
         startup::{new::new_args, urls::args_urls},
     },
-    lua::threads::runner,
+    lua::{
+        threads::runner,
+        network::http::{SLEEP_TIME, REQUESTS_LIMIT}
+    },
     ScanTypes,
 };
 use structopt::StructOpt;
@@ -36,6 +39,8 @@ async fn main() -> Result<(), std::io::Error> {
             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);
+            *SLEEP_TIME.lock().unwrap() = opts.delay;
+            *REQUESTS_LIMIT.lock().unwrap() = opts.requests_limit;
             let scan_futures = vec![
                 opts.lotus_obj.start(
                     opts.target_data.urls,