implementing input-handler option

src/cli/args.rs

@@ -6,6 +6,7 @@ use std::collections::HashMap;
 use std::path::PathBuf;
 use structopt::StructOpt;
 
+
 fn parse_headers(raw_headers: &str) -> Result<HeaderMap, serde_json::Error> {
     let parsed_json = serde_json::from_str::<HashMap<String, String>>(raw_headers);
 
@@ -173,5 +174,7 @@ pub enum Opts {
         exit_after: i32,
         #[structopt(long = "env-vars",parse(try_from_str = get_env_vars),default_value="{}",help = "Set Global Vars for scripts")]
         env_vars: Value,
+        #[structopt(long = "input-handler", parse(from_os_str), help = "Create custom input for your scripts")]
+        input_handler: Option<PathBuf>
     },
 }

src/cli/input.rs

@@ -59,6 +59,22 @@ pub fn get_target_paths(urls: Vec<String>) -> Result<Vec<String>, String> {
     Ok(paths)
 }
 
+pub fn get_stdin_input() -> Result<Vec<String>, CliErrors> {
+    if atty::is(atty::Stream::Stdin) {
+        Err(CliErrors::EmptyStdin)
+    } else {
+        let stdin = io::stdin();
+        let mut input_lines: Vec<String> = Vec::new();
+        stdin.lock().lines().for_each(|x| {
+            let the_line = x.unwrap();
+            input_lines.push(the_line);
+        });
+        input_lines.sort();
+        input_lines.dedup();
+        Ok(input_lines)
+    }
+
+}
 pub fn get_target_urls(url_file: Option<PathBuf>) -> Result<Vec<String>, CliErrors> {
     if url_file.is_some() {
         let urls = filename_to_string(url_file.unwrap().to_str().unwrap());
@@ -72,25 +88,9 @@ pub fn get_target_urls(url_file: Option<PathBuf>) -> Result<Vec<String>, CliErro
             Err(CliErrors::ReadingError)
         }
     } else {
-        if atty::is(atty::Stream::Stdin) {
-            Err(CliErrors::EmptyStdin)
-        } else {
-            let stdin = io::stdin();
-            let mut urls: Vec<String> = Vec::new();
-            stdin.lock().lines().for_each(|x| {
-                let the_url = x.unwrap();
-                match url::Url::parse(&the_url) {
-                    Ok(..) => {
-                        urls.push(the_url);
-                    }
-                    Err(..) => {
-                        log::error!("Cannot Parse {} url, ignoring ..", the_url);
-                    }
-                };
-            });
-            urls.sort();
-            urls.dedup();
-            Ok(urls)
+        let mut urls = get_stdin_input()?;
+        urls.sort();
+        urls.dedup();
+        Ok(urls)
         }
-    }
 }

src/cli/startup/scan/input_handler.rs

@@ -0,0 +1,17 @@
+use mlua::Lua;
+use serde_json::Value as JsonValue;
+
+pub fn custom_input_lua(input_data: Vec<String>, code: &str) -> Result<Vec<JsonValue>, mlua::Error> {
+    let lua_env = Lua::new();
+    log::debug!("Loading Input handler lua code");
+    lua_env.load(code).exec()?;
+    let parse_input = lua_env.globals().get::<_, mlua::Function>("parse_input")?;
+    log::debug!("Calling parse_input function");
+    let output_data = parse_input.call::<_, Vec<mlua::Value>>(input_data)?;
+    let final_output: Vec<JsonValue> = output_data
+        .into_iter()
+        .map(|table| serde_json::to_value(&table).unwrap())
+        .collect();
+    log::debug!("parse_input function returned {} item",final_output.len());
+    Ok(final_output)
+}

src/cli/startup/scan/mod.rs

@@ -0,0 +1 @@
+pub mod scan;

src/cli/startup/scan/scan.rs

@@ -0,0 +1,165 @@
+use crate::cli::args::Opts;
+use crate::cli::input::{get_stdin_input, get_target_paths, get_target_hosts};
+use crate::cli::logger::init_log;
+use crate::lua::parsing::files::filename_to_string;
+use crate::{show_msg, MessageLevel, Lotus, RequestOpts};
+use std::sync::{Arc, Mutex};
+use structopt::StructOpt;
+use url::Url;
+#[path ="input_handler.rs"]
+mod input_handler;
+use input_handler::custom_input_lua;
+
+pub struct ScanArgs {
+    pub target_data: TargetData,
+    pub exit_after: i32,
+    pub is_request: bool,
+    pub req_opts: RequestOpts,
+    pub lotus_obj: Lotus,
+    pub requests_limit: i32,
+    pub delay: u64,
+    pub fuzz_workers: usize,
+    pub verbose: bool,
+}
+
+pub struct TargetData {
+    pub urls: Vec<String>,
+    pub hosts: Vec<String>,
+    pub paths: Vec<String>,
+    pub custom: Vec<serde_json::Value>
+}
+
+pub fn args_scan() -> ScanArgs {
+    let (
+        urls,
+        hosts,
+        paths,
+        custom,
+        exit_after,
+        is_request,
+        req_opts,
+        lotus_obj,
+        requests_limit,
+        delay,
+        fuzz_workers,
+        verbose,
+    ) = match Opts::from_args() {
+        Opts::SCAN {
+            redirects,
+            workers,
+            scripts_workers,
+            timeout,
+            script_path,
+            output,
+            proxy,
+            log,
+            urls,
+            headers,
+            exit_after,
+            requests_limit,
+            delay,
+            fuzz_workers,
+            verbose,
+            is_request,
+            env_vars,
+            input_handler
+        } => {
+            // setup logger
+            init_log(log).unwrap();
+            let req_opts = RequestOpts {
+                headers,
+                proxy,
+                timeout,
+                redirects,
+            };
+            let lotus_obj = Lotus {
+                script_path,
+                output,
+                workers,
+                script_workers: scripts_workers,
+                stop_after: Arc::new(Mutex::new(1)),
+                env_vars,
+            };
+            let input_data = if let Some(urls_file) = urls {
+                let read_file = filename_to_string(urls_file.to_str().unwrap());
+                if let Err(..) = read_file {
+                    show_msg("Cannot Read the urls file", MessageLevel::Error);
+                    std::process::exit(1);
+                }
+                read_file.unwrap().lines().map(|line| line.to_string()).collect::<Vec<String>>()
+            } else {
+                match get_stdin_input() {
+                    Ok(input_data) => input_data,
+                    Err(..) => {
+                        show_msg("No input in Stdin", MessageLevel::Error);
+                        std::process::exit(1);
+                    }
+                }
+            };
+            let input_handler = if let Some(custom_input) = input_handler {
+                let lua_code = filename_to_string(custom_input.to_str().unwrap());
+                if let Err(err) = lua_code {
+                    show_msg(&format!("Unable to read custom input lua script: {}",err), MessageLevel::Error);
+                    vec![]
+                } else {
+                    let lua_output = custom_input_lua(input_data.clone(),&lua_code.unwrap());
+                    if let Ok(lua_output) = lua_output {
+                        lua_output
+                    } else {
+                        show_msg(&format!("{}", lua_output.unwrap_err()), MessageLevel::Error);
+                        vec![]
+                    }
+                }
+            } else {
+                vec![]
+            };
+            log::debug!("{:?}", input_handler);
+            let mut urls = vec![];
+            let mut paths = vec![];
+            let mut hosts = vec![];
+            if input_handler.len() == 0 {
+                urls = input_data.iter().filter_map(|target_url| Url::parse(target_url).ok()).map(|url| url.to_string()).collect();
+                paths = match get_target_paths(urls.clone()) {
+                    Ok(paths) => paths,
+                    Err(err) => {
+                        show_msg(
+                            &format!("Failed to get target paths: {}", err),
+                            MessageLevel::Error,
+                        );
+                        vec![]
+                    }
+                };
+                hosts = get_target_hosts(urls.clone());
+            }
+            (
+                urls,
+                hosts,
+                paths,
+                input_handler,
+                exit_after,
+                is_request,
+                req_opts,
+                lotus_obj,
+                requests_limit,
+                delay,
+                fuzz_workers,
+                verbose,
+            )
+        }
+        _ => {
+            std::process::exit(1);
+        }
+    };
+
+    ScanArgs {
+        target_data: TargetData { urls, hosts, paths, custom },
+        exit_after,
+        is_request,
+        req_opts,
+        lotus_obj,
+        requests_limit,
+        delay,
+        fuzz_workers,
+        verbose,
+    }
+}