check: add --json formatter

Author: adrienthebo

check/src/main.rs

@@ -1,18 +1,41 @@
-use clap::Parser;
+use chrono::{DateTime, FixedOffset};
+use clap::{Parser, ValueEnum};
 use futures::TryStreamExt;
 use log::{debug, error, info, warn};
 use pcap_file_tokio::pcapng::{Block, PcapNgReader};
 use rayhunter::{
-    analysis::analyzer::{AnalysisRow, AnalyzerConfig, EventType, Harness},
+    analysis::analyzer::{AnalysisRow, AnalyzerConfig, Event, EventType, Harness},
     diag::DataType,
     gsmtap_parser,
     pcap::GsmtapPcapWriter,
     qmdl::QmdlReader,
 };
+use serde_json;
 use std::{collections::HashMap, future, path::PathBuf, pin::pin};
 use tokio::fs::File;
 use walkdir::WalkDir;
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
+enum OutputFormat {
+    Text,
+    Json,
+}
+
+impl Default for OutputFormat {
+    fn default() -> Self {
+        OutputFormat::Text
+    }
+}
+
+impl std::fmt::Display for OutputFormat {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            OutputFormat::Text => write!(f, "text"),
+            OutputFormat::Json => write!(f, "json"),
+        }
+    }
+}
+
 #[derive(Parser, Debug)]
 #[command(version, about)]
 struct Args {
@@ -22,6 +45,9 @@ struct Args {
     #[arg(short = 'P', long, help = "Convert qmdl files to pcap before analysis")]
     pcapify: bool,
 
+    #[arg(short, long, default_value_t = OutputFormat::Text)]
+    format: OutputFormat,
+
     #[arg(long, help = "Show why some packets were skipped during analysis")]
     show_skipped: bool,
 
@@ -32,23 +58,30 @@ struct Args {
     debug: bool,
 }
 
+trait Reporter {
+    fn process_row(&mut self, row: AnalysisRow);
+    fn finish(&self, show_skipped: bool) -> std::io::Result<()>;
+}
+
 #[derive(Default)]
-struct Report {
+struct TextReporter {
     skipped_reasons: HashMap<String, u32>,
     total_messages: u32,
     warnings: u32,
     skipped: u32,
     file_path: String,
 }
 
-impl Report {
+impl TextReporter {
     fn new(file_path: &str) -> Self {
-        Report {
+        TextReporter {
             file_path: file_path.to_string(),
             ..Default::default()
         }
     }
+}
 
+impl Reporter for TextReporter {
     fn process_row(&mut self, row: AnalysisRow) {
         self.total_messages += 1;
         if let Some(reason) = row.skipped_message_reason {
@@ -76,7 +109,7 @@ impl Report {
         }
     }
 
-    fn print_summary(&self, show_skipped: bool) {
+    fn finish(&self, show_skipped: bool) -> std::result::Result<(), std::io::Error> {
         if show_skipped && self.skipped > 0 {
             info!("{}: messages skipped:", self.file_path);
             for (reason, count) in self.skipped_reasons.iter() {
@@ -87,16 +120,90 @@ impl Report {
             "{}: {} messages analyzed, {} warnings, {} messages skipped",
             self.file_path, self.total_messages, self.warnings, self.skipped
         );
+        Ok(())
     }
 }
 
-async fn analyze_pcap(pcap_path: &str, show_skipped: bool) {
+#[derive(serde::Serialize)]
+struct JsonEvent {
+    timestamp: DateTime<FixedOffset>,
+    event: Event,
+}
+
+#[derive(Default)]
+struct JsonReporter {
+    skipped_reasons: HashMap<String, u32>,
+    total_messages: u32,
+    warnings: u32,
+    skipped: u32,
+    file_path: String,
+    events: Vec<JsonEvent>,
+}
+
+impl JsonReporter {
+    fn new(file_path: &str) -> Self {
+        JsonReporter {
+            file_path: file_path.to_string(),
+            ..Default::default()
+        }
+    }
+}
+
+impl Reporter for JsonReporter {
+    fn process_row(&mut self, row: AnalysisRow) {
+        self.total_messages += 1;
+        if let Some(reason) = row.skipped_message_reason {
+            *self.skipped_reasons.entry(reason).or_insert(0) += 1;
+            self.skipped += 1;
+            return;
+        }
+        for maybe_event in row.events {
+            let Some(event) = maybe_event else { continue };
+            let Some(timestamp) = row.packet_timestamp else {
+                continue;
+            };
+
+            match &event.event_type {
+                EventType::Low | EventType::Medium | EventType::High => {
+                    self.warnings += 1;
+                }
+                _ => (),
+            }
+
+            self.events.push(JsonEvent { timestamp, event });
+        }
+    }
+
+    fn finish(&self, show_skipped: bool) -> std::result::Result<(), std::io::Error> {
+        let mut report: serde_json::Value = serde_json::json!({
+            "file_path": &self.file_path,
+            "events": &self.events,
+            "summary": {
+                "total_messages": self.total_messages,
+                "warnings": self.warnings,
+            }
+        });
+
+        if show_skipped && self.skipped > 0 {
+            report.as_object_mut().unwrap().insert(
+                "skipped_messages".to_string(),
+                serde_json::json!({
+                    "total_skipped": self.skipped,
+                    "reasons": &self.skipped_reasons,
+                }),
+            );
+        }
+
+        serde_json::to_writer_pretty(std::io::stdout(), &report).map_err(|e| From::from(e))
+    }
+}
+
+async fn analyze_pcap(pcap_path: &str, show_skipped: bool, mut reporter: Box<dyn Reporter>) {
     let mut harness = Harness::new_with_config(&AnalyzerConfig::default());
     let pcap_file = &mut File::open(&pcap_path).await.expect("failed to open file");
     let mut pcap_reader = PcapNgReader::new(pcap_file)
         .await
         .expect("failed to read PCAP file");
-    let mut report = Report::new(pcap_path);
     while let Some(Ok(block)) = pcap_reader.next_block().await {
         let row = match block {
             Block::EnhancedPacket(packet) => harness.analyze_pcap_packet(packet),
@@ -105,12 +212,15 @@ async fn analyze_pcap(pcap_path: &str, show_skipped: bool) {
                 continue;
             }
         };
-        report.process_row(row);
+        reporter.process_row(row);
+    }
+
+    if let Err(e) = reporter.finish(show_skipped) {
+        error!("Couldn't generate report summary: {e}");
     }
-    report.print_summary(show_skipped);
 }
 
-async fn analyze_qmdl(qmdl_path: &str, show_skipped: bool) {
+async fn analyze_qmdl(qmdl_path: &str, show_skipped: bool, mut reporter: Box<dyn Reporter>) {
     let mut harness = Harness::new_with_config(&AnalyzerConfig::default());
     let qmdl_file = &mut File::open(&qmdl_path).await.expect("failed to open file");
     let file_size = qmdl_file
@@ -124,17 +234,19 @@ async fn analyze_qmdl(qmdl_path: &str, show_skipped: bool) {
             .as_stream()
             .try_filter(|container| future::ready(container.data_type == DataType::UserSpace))
     );
-    let mut report = Report::new(qmdl_path);
     while let Some(container) = qmdl_stream
         .try_next()
         .await
         .expect("failed getting QMDL container")
     {
         for row in harness.analyze_qmdl_messages(container) {
-            report.process_row(row);
+            reporter.process_row(row);
         }
     }
-    report.print_summary(show_skipped);
+
+    if let Err(e) = reporter.finish(show_skipped) {
+        error!("Couldn't generate report summary: {e}");
+    }
 }
 
 async fn pcapify(qmdl_path: &PathBuf) {
@@ -206,16 +318,21 @@ async fn main() {
         let path_str = path.to_str().unwrap();
         // instead of relying on the QMDL extension, can we check if a file is
         // QMDL by inspecting the contents?
+
+        let mut reporter: Box<dyn Reporter> = match args.format {
+            OutputFormat::Text => Box::new(TextReporter::new(path_str)),
+            OutputFormat::Json => Box::new(JsonReporter::new(path_str)),
+        };
         if name_str.ends_with(".qmdl") {
             info!("**** Beginning analysis of {name_str}");
-            analyze_qmdl(path_str, args.show_skipped).await;
+            analyze_qmdl(path_str, args.show_skipped, reporter).await;
             if args.pcapify {
                 pcapify(&path.to_path_buf()).await;
             }
         } else if name_str.ends_with(".pcap") || name_str.ends_with(".pcapng") {
             // TODO: if we've already analyzed a QMDL, skip its corresponding pcap
             info!("**** Beginning analysis of {name_str}");
-            analyze_pcap(path_str, args.show_skipped).await;
+            analyze_pcap(path_str, args.show_skipped, reporter).await;
         }
     }
 }