From a4da9214bd888c8c13b879caaec71301c1e8e51f Mon Sep 17 00:00:00 2001
From: Mingwei Zhang <mingwei@mwzhang.com>
Date: Thu, 26 Dec 2024 16:51:02 -0800
Subject: [PATCH] Add MRT output file writing functionality

Introduce support for writing filtered MRT messages to a file. Added
a new `mrt_path` CLI argument and implemented handling to write MRT
output to the specified path when provided. Retained existing stdout
behavior as the default.
---
 src/bin/monocle.rs | 45 ++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/src/bin/monocle.rs b/src/bin/monocle.rs
index 5e7ba0f..3c805c7 100644
--- a/src/bin/monocle.rs
+++ b/src/bin/monocle.rs
@@ -204,6 +204,10 @@ enum Commands {
         #[clap(long)]
         pretty: bool,
 
+        /// MRT output file path
+        #[clap(long, short = 'M')]
+        mrt_path: Option<PathBuf>,
+
         /// Filter by AS path regex string
         #[clap(flatten)]
         filters: ParseFilters,
@@ -228,7 +232,7 @@ enum Commands {
         sqlite_path: Option<PathBuf>,
 
         /// MRT output file path
-        #[clap(long)]
+        #[clap(long, short = 'M')]
         mrt_path: Option<PathBuf>,
 
         /// SQLite reset database content if exists
@@ -421,6 +425,7 @@ fn main() {
             file_path,
             json,
             pretty,
+            mrt_path,
             filters,
         } => {
             if let Err(e) = filters.validate() {
@@ -445,13 +450,39 @@ fn main() {
             .unwrap();
 
             let mut stdout = std::io::stdout();
-            for elem in parser {
-                let output_str = elem_to_string(&elem, json, pretty, "");
-                if let Err(e) = writeln!(stdout, "{}", &output_str) {
-                    if e.kind() != std::io::ErrorKind::BrokenPipe {
-                        eprintln!("{e}");
+
+            match mrt_path {
+                None => {
+                    for elem in parser {
+                        // output to stdout
+                        let output_str = elem_to_string(&elem, json, pretty, "");
+                        if let Err(e) = writeln!(stdout, "{}", &output_str) {
+                            if e.kind() != std::io::ErrorKind::BrokenPipe {
+                                eprintln!("{e}");
+                            }
+                            std::process::exit(1);
+                        }
+                    }
+                }
+                Some(p) => {
+                    let path = p.to_str().unwrap().to_string();
+                    println!("processing. filtered messages output to {}...", &path);
+                    let mut encoder = MrtUpdatesEncoder::new();
+                    let mut writer = match oneio::get_writer(&path) {
+                        Ok(w) => w,
+                        Err(e) => {
+                            eprintln!("{e}");
+                            std::process::exit(1);
+                        }
+                    };
+                    let mut total_count = 0;
+                    for elem in parser {
+                        total_count += 1;
+                        encoder.process_elem(&elem);
                     }
-                    std::process::exit(1);
+                    writer.write_all(&encoder.export_bytes()).unwrap();
+                    drop(writer);
+                    println!("done. total of {} message wrote", total_count);
                 }
             }
         }