From e517bede72ab9d1a187e529254b958a0673ddcfb Mon Sep 17 00:00:00 2001
From: Kremilly <contact@kremilly.com>
Date: Tue, 10 Dec 2024 13:51:00 -0300
Subject: [PATCH] feat: refactor reports module to reports_xss, add JSON output
 support and enhance error handling

---
 .gitignore                 |  1 +
 src/plugins/mod.rs         |  2 +-
 src/plugins/reports.rs     | 29 -------------------
 src/plugins/reports_xss.rs | 58 ++++++++++++++++++++++++++++++++++++++
 src/plugins/scan_xss.rs    | 18 +++++++++---
 src/ui/errors_alerts.rs    |  9 ++++++
 6 files changed, 83 insertions(+), 34 deletions(-)
 delete mode 100644 src/plugins/reports.rs
 create mode 100644 src/plugins/reports_xss.rs

diff --git a/.gitignore b/.gitignore
index 751a4c8..6c56f56 100644
--- a/.gitignore
+++ b/.gitignore
@@ -37,5 +37,6 @@ dumpsync1.yml
 # Reports files
 *.csv
 *.pdf
+*.json
 *.html
 *.xlsx
\ No newline at end of file
diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs
index c844eda..6205643 100644
--- a/src/plugins/mod.rs
+++ b/src/plugins/mod.rs
@@ -1,2 +1,2 @@
-pub mod reports;
+pub mod reports_xss;
 pub mod scan_xss;
\ No newline at end of file
diff --git a/src/plugins/reports.rs b/src/plugins/reports.rs
deleted file mode 100644
index d16aad8..0000000
--- a/src/plugins/reports.rs
+++ /dev/null
@@ -1,29 +0,0 @@
-use csv::Writer;
-use std::error::Error;
-
-use crate::ui::scan_alerts::ScanAlerts;
-
-pub struct Reports;
-
-impl Reports {
-
-    pub fn xss(&self, detections: Vec<(String, usize, String, String)>, output_path: &str) -> Result<(), Box<dyn Error>> {
-        let mut writer = Writer::from_path(output_path)?;
-        writer.write_record(&["Table", "Row Index", "Column", "Value"])?;
-
-        for (table, row_index, column, value) in detections {
-            writer.write_record(&[
-                table,
-                row_index.to_string(),
-                column,
-                value,
-            ])?;
-        }
-
-        writer.flush()?;
-        ScanAlerts::reports_generated(output_path);
-        
-        Ok(())
-    }
-    
-}
diff --git a/src/plugins/reports_xss.rs b/src/plugins/reports_xss.rs
new file mode 100644
index 0000000..f6e0af2
--- /dev/null
+++ b/src/plugins/reports_xss.rs
@@ -0,0 +1,58 @@
+use csv::Writer;
+use std::fs::File;
+use serde::Serialize;
+use std::error::Error;
+use serde_json::to_writer_pretty;
+
+use crate::ui::scan_alerts::ScanAlerts;
+
+#[derive(Serialize)]
+struct Detection {
+    table: String,
+    row_index: usize,
+    column: String,
+    value: String,
+}
+
+pub struct ReportsXSS;
+
+impl ReportsXSS {
+
+    pub fn csv(&self, detections: Vec<(String, usize, String, String)>, output_path: &str) -> Result<(), Box<dyn Error>> {
+        let mut writer = Writer::from_path(output_path)?;
+        writer.write_record(&["Table", "Row Index", "Column", "Value"])?;
+
+        for (table, row_index, column, value) in detections {
+            writer.write_record(&[
+                table,
+                row_index.to_string(),
+                column,
+                value,
+            ])?;
+        }
+
+        writer.flush()?;
+        
+        ScanAlerts::reports_generated(output_path);
+        Ok(())
+    }
+
+    pub fn json(&self, detections: Vec<(String, usize, String, String)>, output_path: &str) -> Result<(), Box<dyn Error>> {
+        let detections: Vec<Detection> = detections
+            .into_iter()
+            .map(|(table, row_index, column, value)| Detection {
+                table,
+                row_index,
+                column,
+                value,
+            })
+            .collect();
+
+        let file = File::create(output_path)?;
+        to_writer_pretty(file, &detections)?;
+        
+        ScanAlerts::reports_generated(output_path);
+        Ok(())
+    }
+    
+}
diff --git a/src/plugins/scan_xss.rs b/src/plugins/scan_xss.rs
index 20d13d7..5365494 100644
--- a/src/plugins/scan_xss.rs
+++ b/src/plugins/scan_xss.rs
@@ -10,9 +10,13 @@ use mysql::{
 
 use crate::{
     consts::global::Global, 
-    plugins::reports::Reports, 
-    ui::scan_alerts::ScanAlerts, 
-    engine::connection::Connection, 
+    plugins::reports_xss::ReportsXSS,  
+    engine::connection::Connection,
+    
+    ui::{
+        scan_alerts::ScanAlerts,
+        errors_alerts::ErrorsAlerts,
+    },
     
     helpers::{
         scan_handlers::ScanHandlers,
@@ -118,7 +122,13 @@ impl ScanXSS {
         }
 
         if let Some(file) = &self.file {
-            Reports.xss(detections, file)?;
+            if file.ends_with(".csv") {
+                ReportsXSS.csv(detections, file)?;
+            } else if file.ends_with(".json") {
+                ReportsXSS.json(detections, file)?;
+            } else {
+                ErrorsAlerts::report_format();
+            }
         }
         
         Ok(())
diff --git a/src/ui/errors_alerts.rs b/src/ui/errors_alerts.rs
index f182a68..0aa29e1 100644
--- a/src/ui/errors_alerts.rs
+++ b/src/ui/errors_alerts.rs
@@ -49,4 +49,13 @@ impl ErrorsAlerts {
         );
     }
 
+    pub fn report_format() {
+        let message = "Invalid file format, only CSV and JSON are supported.";
+
+        println!("{}", "-".repeat(50));
+        println!(
+            "{}", message.red().bold(), 
+        );
+    }
+
 }