Merge pull request #61 from rusty-sec/change_report

MindPatch · web-flow · commit 4f09a94962fd · 2022-12-31T15:39:24.000+02:00
Change report
diff --git a/src/cli/args.rs b/src/cli/args.rs
@@ -68,4 +68,6 @@ pub struct Opts {
 
     #[structopt(long = "headers", parse(try_from_str = parse_headers), required = false, default_value = "{}")]
     pub headers: HeaderMap,
+    #[structopt(long = "exit-after-errors", default_value = "2000")]
+    pub exit_after: i32,
 }
diff --git a/src/cli/bar.rs b/src/cli/bar.rs
@@ -16,8 +16,15 @@
  * limitations under the License.
  */
 
+use console::Style;
 use indicatif::{ProgressBar, ProgressStyle};
 
+pub enum MessageLevel {
+    Info,
+    Warn,
+    Error,
+}
+
 /// Lotus ProgressBar based on the length of `bar` parameter
 pub fn create_progress(bar: u64) -> ProgressBar {
     let bar = ProgressBar::new(bar);
@@ -32,3 +39,25 @@ pub fn create_progress(bar: u64) -> ProgressBar {
     );
     bar
 }
+
+pub fn show_msg(message: &str, msglevel: MessageLevel) {
+    let print_level = match msglevel {
+        MessageLevel::Info => {
+            log::info!("{}", message);
+            format!("[{}]", Style::new().blue().apply_to("INFO"))
+        }
+        MessageLevel::Warn => {
+            log::warn!("{}", message);
+            format!("[{}]", Style::new().yellow().apply_to("WARN"))
+        }
+        MessageLevel::Error => {
+            log::error!("{}", message);
+            format!("[{}]", Style::new().red().apply_to("ERROR"))
+        }
+    };
+    if let MessageLevel::Error = msglevel {
+        eprintln!("{print_level} {message}");
+    } else {
+        println!("{print_level} {message}");
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -16,7 +16,6 @@
  * limitations under the License.
  */
 
-#![allow(unused_imports)]
 mod cli;
 mod lua_api;
 mod network;
@@ -26,14 +25,15 @@ mod payloads;
 mod scan;
 
 use cli::bar::create_progress;
+use cli::bar::{show_msg, MessageLevel};
 use cli::errors::CliErrors;
+use futures::{stream, StreamExt};
 use glob::glob;
 use log::error;
 use parsing::files::filename_to_string;
 use reqwest::header::HeaderMap;
 use std::path::PathBuf;
-use futures::{stream, StreamExt};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct RequestOpts {
@@ -47,11 +47,12 @@ pub struct Lotus {
     pub script_path: PathBuf,
     pub output: Option<PathBuf>,
     pub workers: usize,
-    pub script_workers: usize
+    pub script_workers: usize,
+    pub stop_after: Arc<Mutex<i32>>,
 }
 
 impl Lotus {
-    pub async fn start(&self, urls: Vec<String>, request_option: RequestOpts) {
+    pub async fn start(&self, urls: Vec<String>, request_option: RequestOpts, exit_after: i32) {
         let loaded_scripts = {
             if self.script_path.is_dir() {
                 self.load_scripts()
@@ -60,39 +61,60 @@ impl Lotus {
             }
         };
         if loaded_scripts.is_err() {
-            eprintln!("Reading errors"); // TODO
+            show_msg(&format!("Loading scripts error: {}",loaded_scripts.unwrap_err()), MessageLevel::Error);
             std::process::exit(1);
         }
         let bar =
             create_progress(urls.len() as u64 * loaded_scripts.as_ref().unwrap().len() as u64);
-        if loaded_scripts.is_err() {
-            eprintln!("Reading error bruh"); // TODO
+        let loaded_scripts = loaded_scripts.unwrap();
+        if self.output.is_none() {
+            show_msg("Output argument is missing", MessageLevel::Error);
             std::process::exit(1);
         }
-        let loaded_scripts = loaded_scripts.unwrap();
-
         let lotus_obj = Arc::new(scan::LuaLoader::new(
             &bar,
             request_option.clone(),
-            self.output.as_ref().unwrap().to_str().unwrap().to_string()));
+            self.output.as_ref().unwrap().to_str().unwrap().to_string(),
+        ));
         stream::iter(urls)
             .map(move |url| {
                 let loaded_scripts = loaded_scripts.clone();
                 let lotus_loader = Arc::clone(&lotus_obj);
                 stream::iter(loaded_scripts.into_iter())
                     .map(move |(script_out, script_name)| {
                         let url = url.clone();
-                        log::debug!("Running {} script on {} url",script_name, url);
                         let lotus_loader = Arc::clone(&lotus_loader);
+                        let error_check = {
+                            if *self.stop_after.lock().unwrap() == exit_after {
+                                log::debug!("Ignoring scripts");
+                                false
+                            } else {
+                                log::debug!("Running {} script on {} url", script_name, url);
+                                true
+                            }
+                        };
                         async move {
-                            lotus_loader.run_scan(url.as_str(),None,&script_out, &script_name).await
+                            if error_check == false {
+                                // Nothing
+                            } else {
+                                let run_scan = lotus_loader
+                                    .run_scan(url.as_str(), None, &script_out, &script_name)
+                                    .await;
+                                if run_scan.is_err() {
+                                    log::error!("Script is raising error");
+                                    let mut a = self.stop_after.lock().unwrap();
+                                    log::debug!("Errors Counter: {}", a);
+                                    *a += 1;
+                                }
+                            }
                         }
                     })
                     .buffer_unordered(self.script_workers)
                     .collect::<Vec<_>>()
             })
             .buffer_unordered(self.workers)
-            .collect::<Vec<_>>().await;
+            .collect::<Vec<_>>()
+            .await;
     }
 
     fn load_scripts(&self) -> Result<Vec<(String, String)>, CliErrors> {
diff --git a/src/lua_api.rs b/src/lua_api.rs
@@ -16,7 +16,8 @@
  * limitations under the License.
  */
 
-use crate::output::report::{AllReports, OutReport};
+use crate::output::vuln::{AllReports, OutReport};
+use crate::output::cve::CveReport;
 use crate::parsing::html::{css_selector, html_parse, html_search, Location};
 use crate::parsing::url::HttpMessage;
 use crate::payloads;
@@ -93,11 +94,12 @@ pub fn http_func(target_url: &str, lua: &Lua) {
         .set(
             "Reports",
             AllReports {
-                reports: Vec::new(),
+                reports: Vec::new()
             },
         )
         .unwrap();
-    lua.globals().set("NewReport", OutReport::init()).unwrap();
+    lua.globals().set("VulnReport", OutReport::init()).unwrap();
+    lua.globals().set("CveReport", CveReport::init()).unwrap();
 }
 
 pub fn get_utilsfunc<'prog>(the_bar: &'prog indicatif::ProgressBar, lua: &Lua) {
diff --git a/src/main.rs b/src/main.rs
@@ -20,9 +20,11 @@ mod cli;
 use cli::args::Opts;
 use cli::errors::CliErrors;
 use cli::logger::init_log;
+use cli::bar::{show_msg, MessageLevel};
 use lotus::RequestOpts;
 use std::io;
 use std::io::BufRead;
+use std::sync::{Arc, Mutex};
 use structopt::StructOpt;
 
 #[tokio::main]
@@ -33,7 +35,7 @@ async fn main() -> Result<(), std::io::Error> {
 
     let urls = get_target_urls();
     if urls.is_err() {
-        eprintln!("EmptyStdin");
+        show_msg("No input in Stdin",MessageLevel::Error);
         std::process::exit(1);
     }
     // default request options
@@ -47,7 +49,8 @@ async fn main() -> Result<(), std::io::Error> {
         script_path: args.script_path,
         output: args.output,
         workers: args.workers,
-        script_workers: args.scripts_workers
+        script_workers: args.scripts_workers,
+        stop_after: Arc::new(Mutex::new(1)),
     };
     lotus_obj
         .start(
@@ -56,6 +59,7 @@ async fn main() -> Result<(), std::io::Error> {
                 .map(|url| url.to_string())
                 .collect::<Vec<String>>(),
             req_opts,
+            args.exit_after,
         )
         .await;
     Ok(())
diff --git a/src/network/http.rs b/src/network/http.rs
@@ -137,3 +137,26 @@ impl Sender {
         }
     }
 }
+
+pub trait SenderExt {
+    fn make_curl(&self) -> String;
+}
+
+impl SenderExt for Sender {
+    fn make_curl(&self) -> String {
+        let mut curl_command = "curl ".to_string();
+        self.headers.iter().for_each(|(header_name, header_value)| {
+            let header_command = format!(
+                "-H '{}: {}'",
+                header_name.as_str(),
+                header_value.to_str().unwrap()
+            );
+            curl_command.push_str(&header_command);
+        });
+        if self.proxy.is_none() {
+            curl_command.push_str(&format!("-x {}",self.proxy.clone().unwrap()));
+        }
+        curl_command.push_str(&format!("--connect-timeout {}",self.timeout));
+        curl_command
+    }
+}
diff --git a/src/output/cve.rs b/src/output/cve.rs
@@ -0,0 +1,71 @@
+/*
+ * This file is part of Lotus Project, an Web Security Scanner written in Rust based on Lua Scripts
+ * For details, please see https://github.com/rusty-sec/lotus/
+ *
+ * Copyright (c) 2022 - Khaled Nassar
+ *
+ * Please note that this file was originally released under the
+ * GNU General Public License as published by the Free Software Foundation;
+ * either version 2 of the License, or (at your option) any later version.
+ *
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+use mlua::UserData;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Deserialize, Serialize)]
+pub struct CveReport {
+    pub name: Option<String>,
+    pub description: Option<String>,
+    pub url: Option<String>,
+    pub risk: Option<String>,
+    pub matchers: Option<Vec<String>>,
+}
+
+
+impl UserData for CveReport {
+    fn add_methods<'lua, M: mlua::UserDataMethods<'lua, Self>>(methods: &mut M) {
+        methods.add_method_mut("setRisk", |_, this, risk: String| {
+            this.risk = Some(risk);
+            Ok(())
+        });
+        methods.add_method_mut("setName", |_, this, name: String| {
+            this.name = Some(name);
+            Ok(())
+        });
+        methods.add_method_mut("setUrl", |_, this, url: String| {
+            this.url = Some(url);
+            Ok(())
+        });
+
+        methods.add_method_mut("setDescription", |_, this, description: String| {
+            this.description = Some(description);
+            Ok(())
+        });
+
+        methods.add_method_mut("setMatchers", |_, this, matchers: Vec<String>| {
+            this.matchers = Some(matchers);
+            Ok(())
+        });
+
+    }
+}
+
+impl CveReport {
+    pub fn init() -> CveReport {
+        CveReport { 
+            name: None,
+            description: None,
+            url: None,
+            risk: None,
+            matchers: None
+        }
+    }
+}
+
diff --git a/src/output/mod.rs b/src/output/mod.rs
@@ -1 +1,2 @@
-pub mod report;
+pub mod cve;
+pub mod vuln;
diff --git a/src/output/vuln.rs b/src/output/vuln.rs
@@ -18,6 +18,15 @@
 
 use mlua::UserData;
 use serde::{Deserialize, Serialize};
+use crate::output::cve::CveReport;
+
+
+#[derive(Clone, Deserialize, Serialize)]
+#[serde(tag = "report_type")]
+pub enum LotusReport {
+    CVE(CveReport),
+    VULN(OutReport),
+}
 
 #[derive(Clone, Deserialize, Serialize)]
 pub struct OutReport {
@@ -33,13 +42,17 @@ pub struct OutReport {
 
 #[derive(Clone)]
 pub struct AllReports {
-    pub reports: Vec<OutReport>,
+    pub reports: Vec<LotusReport>,
 }
 
 impl UserData for AllReports {
     fn add_methods<'lua, M: mlua::UserDataMethods<'lua, Self>>(methods: &mut M) {
-        methods.add_method_mut("addReport", |_, this, the_report: OutReport| {
-            this.reports.push(the_report);
+        methods.add_method_mut("addVulnReport", |_, this, the_report: OutReport| {
+            this.reports.push(LotusReport::VULN(the_report));
+            Ok(())
+        });
+        methods.add_method_mut("addCveReport", |_, this, the_report: CveReport| {
+            this.reports.push(LotusReport::CVE(the_report));
             Ok(())
         });
     }
diff --git a/src/scan/mod.rs b/src/scan/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -68,4 +68,6 @@ pub struct Opts {`
`68`	`68`
`69`	`69`	`#[structopt(long = "headers", parse(try_from_str = parse_headers), required = false, default_value = "{}")]`
`70`	`70`	`pub headers: HeaderMap,`
	`71`	`+ #[structopt(long = "exit-after-errors", default_value = "2000")]`
	`72`	`+ pub exit_after: i32,`
`71`	`73`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`		`-pub mod report;`
	`1`	`+pub mod cve;`
	`2`	`+pub mod vuln;`