Enable resume functionality for exit cases.

src/cli/args/scan.rs

@@ -3,7 +3,41 @@ use reqwest::header::{HeaderName, HeaderValue};
 use serde_json::Value;
 use std::collections::HashMap;
 use std::path::PathBuf;
+use std::fs::File;
+use std::io::{BufRead, BufReader};
 use structopt::StructOpt;
+use crate::lua::threads::runner::{LAST_CUSTOM_SCAN_ID, LAST_PATH_SCAN_ID,LAST_HOST_SCAN_ID, LAST_URL_SCAN_ID};
+
+fn read_resume_file(file_path: &str) -> Result<(), std::io::Error> {
+    let file = File::open(file_path)?;
+    let reader = BufReader::new(file);
+
+    for line in reader.lines() {
+        let line = line?;
+        let parts: Vec<&str> = line.split("=").collect();
+        if parts.len() != 2 {
+            continue;
+        }
+
+        match parts[0] {
+            "LAST_URL_SCAN_ID" => {
+                *LAST_URL_SCAN_ID.lock().unwrap() = parts[1].parse().unwrap_or(0);
+            }
+            "LAST_HOST_SCAN_ID" => {
+                *LAST_HOST_SCAN_ID.lock().unwrap() = parts[1].parse().unwrap_or(0);
+            }
+            "LAST_PATH_SCAN_ID" => {
+                *LAST_PATH_SCAN_ID.lock().unwrap() = parts[1].parse().unwrap_or(0);
+            }
+            "LAST_CUSTOM_SCAN_ID" => {
+                *LAST_CUSTOM_SCAN_ID.lock().unwrap() = parts[1].parse().unwrap_or(0);
+            }
+            _ => {}
+        }
+    }
+
+    Ok(())
+}
 
 fn parse_headers(raw_headers: &str) -> Result<HeaderMap, serde_json::Error> {
     let parsed_json = serde_json::from_str::<HashMap<String, String>>(raw_headers);
@@ -140,4 +174,10 @@ pub struct UrlsOpts {
         help = "Create custom input for your scripts"
     )]
     pub input_handler: Option<PathBuf>,
+    #[structopt(
+        long = "resume",
+        parse(try_from_str = read_resume_file),
+        help = "Resume the scan with resume.cfg where your scans progress stopped in the last run"
+    )]
+    _resume: Option<()>
 }

src/cli/logger.rs

@@ -39,12 +39,18 @@ pub fn init_log(log_file: Option<PathBuf>) -> Result<(), std::io::Error> {
 
     if let Some(log_path) = log_file {
         // Disable unwanted loggers
-        logger.chain(fern::log_file(log_path.clone()).unwrap()).apply().unwrap();
-        log::info!("Logging initialized. Writing logs to {}", log_path.to_str().unwrap());
+        logger
+            .chain(fern::log_file(log_path.clone()).unwrap())
+            .apply()
+            .unwrap();
+        log::info!(
+            "Logging initialized. Writing logs to {}",
+            log_path.to_str().unwrap()
+        );
     } else {
         logger.apply().unwrap();
         log::info!("Logging initialized. Writing logs to console.");
     }
-    
+
     Ok(())
 }

src/lib.rs

@@ -23,12 +23,13 @@ use cli::{
 };
 use lua::{
     model::LuaOptions, parsing::files::filename_to_string, run::LuaLoader,
-    threads::runner::iter_futures,
+    threads::runner::{iter_futures, LAST_HOST_SCAN_ID},
 };
 use mlua::ExternalError;
 pub use model::{Lotus, RequestOpts, ScanTypes};
 use std::sync::Arc;
 
+use crate::lua::threads::runner::{LAST_URL_SCAN_ID, LAST_PATH_SCAN_ID, LAST_CUSTOM_SCAN_ID};
 
 impl Lotus {
     /// Run The Lua Script with real target
@@ -47,20 +48,35 @@ impl Lotus {
         if target_data.is_empty() {
             return;
         }
+        let resume_value: usize; 
         let loaded_scripts = match scan_type {
-            ScanTypes::HOSTS => valid_scripts(get_scripts(self.script_path.clone()), 1),
-            ScanTypes::URLS => valid_scripts(get_scripts(self.script_path.clone()), 2),
-            ScanTypes::PATHS => valid_scripts(get_scripts(self.script_path.clone()), 3),
-            ScanTypes::CUSTOM => valid_scripts(get_scripts(self.script_path.clone()), 4),
+            ScanTypes::HOSTS => {
+                resume_value = *LAST_HOST_SCAN_ID.lock().unwrap();
+                valid_scripts(get_scripts(self.script_path.clone()), 1)
+            },
+            ScanTypes::URLS => {
+                resume_value = *LAST_URL_SCAN_ID.lock().unwrap();
+                valid_scripts(get_scripts(self.script_path.clone()), 2)
+            },
+            ScanTypes::PATHS => {
+                resume_value = *LAST_PATH_SCAN_ID.lock().unwrap();
+                valid_scripts(get_scripts(self.script_path.clone()), 3)
+            },
+            ScanTypes::CUSTOM => {
+                resume_value = *LAST_CUSTOM_SCAN_ID.lock().unwrap();
+                valid_scripts(get_scripts(self.script_path.clone()), 4)
+            },
         };
         let lotus_obj = Arc::new(LuaLoader::new(request_option.clone(), self.output.clone()));
         let scan_type = Arc::new(scan_type);
         iter_futures(
+            scan_type.clone(),
             target_data,
             |script_data| async move {
                 let lotus_loader = Arc::clone(&lotus_obj);
                 let scan_type = Arc::clone(&scan_type);
                 iter_futures(
+                    scan_type.clone(),
                     loaded_scripts,
                     |(script_code, script_name)| async move {
                         let lua_opts = LuaOptions {
@@ -87,10 +103,14 @@ impl Lotus {
                         }
                     },
                     self.script_workers,
+                    resume_value,
+                    false,
                 )
                 .await;
             },
             self.workers,
+            resume_value,
+            true,
         )
         .await;
     }

src/lua/parsing/text.rs

@@ -40,7 +40,10 @@ impl ResponseMatcher {
         {
             Ok(re) => Ok(re.is_match(&resp)),
             Err(err) => {
-                log::error!("An error occurred while processing the regular expression pattern: {}", err.to_string());
+                log::error!(
+                    "An error occurred while processing the regular expression pattern: {}",
+                    err.to_string()
+                );
                 Err(CliErrors::RegexPatternError)
             }
         }
@@ -67,7 +70,10 @@ impl ResponseMatcher {
                 Ok(match_iter)
             }
             Err(err) => {
-                log::error!("An error occurred while processing the regular expression pattern: {}", err.to_string());
+                log::error!(
+                    "An error occurred while processing the regular expression pattern: {}",
+                    err.to_string()
+                );
                 Err(CliErrors::RegexPatternError)
             }
         }
@@ -91,7 +97,10 @@ impl ResponseMatcher {
                 Ok(replace_output)
             }
             Err(err) => {
-                log::error!("An error occurred while processing the regular expression pattern: {}", err.to_string());
+                log::error!(
+                    "An error occurred while processing the regular expression pattern: {}",
+                    err.to_string()
+                );
                 Err(CliErrors::RegexPatternError)
             }
         }
@@ -105,24 +114,26 @@ impl ResponseMatcher {
         let mut counter = 0;
         for search_pattern in text.iter() {
             match is_regex.unwrap_or(false) {
-                true => match RegexBuilder::new(&search_pattern)
-                    .multi_line(self.multi_line)
-                    .case_insensitive(self.case_insensitive)
-                    .unicode(self.unicode)
-                    .octal(self.octal)
-                    .dot_matches_new_line(self.dot_matches_new_line)
-                    .build()
-                {
-                    Ok(re_pattern) => {
-                        if re_pattern.is_match(body) {
-                            counter += 1;
+                true => {
+                    match RegexBuilder::new(&search_pattern)
+                        .multi_line(self.multi_line)
+                        .case_insensitive(self.case_insensitive)
+                        .unicode(self.unicode)
+                        .octal(self.octal)
+                        .dot_matches_new_line(self.dot_matches_new_line)
+                        .build()
+                    {
+                        Ok(re_pattern) => {
+                            if re_pattern.is_match(body) {
+                                counter += 1;
+                            }
+                        }
+                        Err(err) => {
+                            log::error!("An error occurred while processing the regular expression pattern: {}",err.to_string());
+                            return Err(CliErrors::RegexPatternError);
                         }
                     }
-                    Err(err) => {
-                        log::error!("An error occurred while processing the regular expression pattern: {}",err.to_string());
-                        return Err(CliErrors::RegexPatternError);
-                    }
-                },
+                }
                 false => {
                     if body.contains(search_pattern) {
                         counter += 1;
@@ -142,24 +153,26 @@ impl ResponseMatcher {
         let mut matched_data = Vec::new();
         for pattern in text {
             match is_regex.unwrap_or(false) {
-                true => match RegexBuilder::new(&pattern)
-                    .multi_line(self.multi_line)
-                    .case_insensitive(self.case_insensitive)
-                    .unicode(self.unicode)
-                    .octal(self.octal)
-                    .dot_matches_new_line(self.dot_matches_new_line)
-                    .build()
-                {
-                    Ok(re) => {
-                        if re.is_match(&body) {
-                            matched_data.push(pattern);
+                true => {
+                    match RegexBuilder::new(&pattern)
+                        .multi_line(self.multi_line)
+                        .case_insensitive(self.case_insensitive)
+                        .unicode(self.unicode)
+                        .octal(self.octal)
+                        .dot_matches_new_line(self.dot_matches_new_line)
+                        .build()
+                    {
+                        Ok(re) => {
+                            if re.is_match(&body) {
+                                matched_data.push(pattern);
+                            }
+                        }
+                        Err(err) => {
+                            log::error!("An error occurred while processing the regular expression pattern: {}",err.to_string());
+                            return Err(CliErrors::RegexPatternError);
                         }
                     }
-                    Err(err) => {
-                        log::error!("An error occurred while processing the regular expression pattern: {}",err.to_string());
-                        return Err(CliErrors::RegexPatternError);
-                    }
-                },
+                }
                 false => {
                     if body.contains(&pattern) {
                         matched_data.push(pattern);

src/lua/threads/runner.rs

@@ -1,20 +1,85 @@
 use futures::stream::{self, StreamExt};
 use futures::Future;
 use futures::{channel::mpsc, sink::SinkExt};
+use lazy_static::lazy_static;
+use std::fs::OpenOptions;
+use std::io::prelude::*;
+use std::sync::{Arc, Mutex};
+use tokio::signal::ctrl_c;
 use tokio::sync::RwLock;
 
+use crate::cli::bar::{show_msg, MessageLevel};
+use crate::ScanTypes;
+
+lazy_static! {
+    pub static ref LAST_URL_SCAN_ID: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
+    pub static ref LAST_HOST_SCAN_ID: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
+    pub static ref LAST_PATH_SCAN_ID: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
+    pub static ref LAST_CUSTOM_SCAN_ID: Arc<Mutex<usize>> = Arc::new(Mutex::new(0));
+}
+
+pub async fn pause_channel() {
+    tokio::spawn(async move {
+        ctrl_c().await.unwrap();
+        if let Err(err) = generate_resume() {
+            show_msg(&err.to_string(), MessageLevel::Error)
+        }
+        std::process::exit(130)
+    });
+}
+
+fn generate_resume() -> Result<(), std::io::Error> {
+    let mut file = OpenOptions::new()
+        .write(true)
+        .create(true)
+        .open("resume.cfg")?;
+
+    let url_scan_id = LAST_URL_SCAN_ID.lock().unwrap();
+    let host_scan_id = LAST_HOST_SCAN_ID.lock().unwrap();
+    let path_scan_id = LAST_PATH_SCAN_ID.lock().unwrap();
+    let custom_scan_id = LAST_CUSTOM_SCAN_ID.lock().unwrap();
+
+    file.write_all(format!("URL_SCAN_ID={}\n", *url_scan_id).as_bytes())?;
+    file.write_all(format!("HOST_SCAN_ID={}\n", *host_scan_id).as_bytes())?;
+    file.write_all(format!("PATH_SCAN_ID={}\n", *path_scan_id).as_bytes())?;
+    file.write_all(format!("CUSTOM_SCAN_ID={}\n", *custom_scan_id).as_bytes())?;
+
+    Ok(())
+}
+
+fn update_index_id(scan_type: Arc<ScanTypes>, index_id: usize) {
+    match *scan_type {
+        ScanTypes::URLS => *LAST_URL_SCAN_ID.lock().unwrap() = index_id,
+        ScanTypes::HOSTS => *LAST_HOST_SCAN_ID.lock().unwrap() = index_id,
+        ScanTypes::PATHS => *LAST_PATH_SCAN_ID.lock().unwrap() = index_id,
+        ScanTypes::CUSTOM => *LAST_CUSTOM_SCAN_ID.lock().unwrap() = index_id,
+    }
+}
+
 // Asynchronous function to iterate over futures concurrently
 // Takes a vector, a function and a number of workers as arguments
 // The function must return a future with no output
 // The vector must implement cloning
-pub async fn iter_futures<F, T, Fut>(target_iter: Vec<T>, target_function: F, workers: usize)
-where
+pub async fn iter_futures<F, T, Fut>(
+    scan_type: Arc<ScanTypes>,
+    target_iter: Vec<T>,
+    target_function: F,
+    workers: usize,
+    skip_index: usize,
+    count_index: bool,
+) where
     F: FnOnce(T) -> Fut + Clone,
     Fut: Future<Output = ()>,
     T: Clone,
 {
     stream::iter(target_iter)
-        .for_each_concurrent(workers, |out| {
+        .enumerate()
+        .skip(skip_index)
+        .for_each_concurrent(workers, |(index_id, out)| {
+            let scan_type = Arc::clone(&scan_type);
+            if count_index == true {
+                update_index_id(scan_type, index_id);
+            }
             let out = out.clone();
             let target_function = target_function.clone();
             async move { target_function(out).await }

src/main.rs

@@ -38,6 +38,8 @@ async fn main() -> Result<(), std::io::Error> {
 }
 
 async fn run_scan() -> Result<(), std::io::Error> {
+    // Spawn a new thread to handle the exit process when the user presses CTRL + C.
+    runner::pause_channel().await;
     let opts = args_scan();
     let fuzz_workers = opts.fuzz_workers;
     show_msg(
@@ -56,7 +58,10 @@ async fn run_scan() -> Result<(), std::io::Error> {
     );
 
     show_msg(
-        &format!("Number of custom entries: {}", opts.target_data.custom.len()),
+        &format!(
+            "Number of custom entries: {}",
+            opts.target_data.custom.len()
+        ),
         MessageLevel::Info,
     );
     // Open two threads for URL/HOST scanning