From 9297ba4e0cf0a8f2d87493a4ff7749caf49aa44b Mon Sep 17 00:00:00 2001
From: NotAShelf <raf@notashelf.dev>
Date: Fri, 2 May 2025 11:22:52 +0300
Subject: [PATCH] eris: allow various log formats

Supports plain, pretty, json and pretty-json
---
 src/config.rs | 19 ++++++++++++
 src/main.rs   | 80 +++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 93 insertions(+), 6 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index ab85136..f4c40b2 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -85,6 +85,22 @@ pub struct Args {
         help = "Log level: trace, debug, info, warn, error"
     )]
     pub log_level: String,
+
+    #[clap(
+        long,
+        default_value = "pretty",
+        help = "Log format: plain, pretty, json, pretty-json"
+    )]
+    pub log_format: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
+pub enum LogFormat {
+    Plain,
+    #[default]
+    Pretty,
+    Json,
+    PrettyJson,
 }
 
 // Trap pattern structure. It can be either a plain string
@@ -145,6 +161,7 @@ pub struct Config {
     pub data_dir: String,
     pub config_dir: String,
     pub cache_dir: String,
+    pub log_format: LogFormat,
 }
 
 impl Default for Config {
@@ -206,6 +223,7 @@ impl Default for Config {
             data_dir: "./data".to_string(),
             config_dir: "./conf".to_string(),
             cache_dir: "./cache".to_string(),
+            log_format: LogFormat::Pretty,
         }
     }
 }
@@ -406,6 +424,7 @@ mod tests {
             base_dir: Some(PathBuf::from("/tmp/eris")),
             config_file: None,
             log_level: "debug".to_string(),
+            log_format: "pretty".to_string(),
         };
 
         let config = Config::from_args(&args);
diff --git a/src/main.rs b/src/main.rs
index 9c6c73f..c8d80a6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -29,15 +29,78 @@ async fn main() -> std::io::Result<()> {
     // Parse command line arguments
     let args = Args::parse();
 
-    // Initialize the logger
-    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(&args.log_level))
-        .format_timestamp_millis()
-        .init();
+    // Determine log format from args
+    let log_format = match args.log_format.to_lowercase().as_str() {
+        "json" => config::LogFormat::Json,
+        "pretty-json" => config::LogFormat::PrettyJson,
+        "plain" => config::LogFormat::Plain,
+        _ => config::LogFormat::Pretty,
+    };
 
-    log::info!("Starting eris tarpit system");
+    // Initialize the logger with proper formatting
+    let env = env_logger::Env::default().default_filter_or(&args.log_level);
+    let mut builder = env_logger::Builder::from_env(env);
+
+    match log_format {
+        config::LogFormat::Plain => {
+            builder.format_timestamp_millis().init();
+        }
+        config::LogFormat::Pretty => {
+            builder
+                .format(|buf, record| {
+                    use std::io::Write;
+                    let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:%S%.3f");
+                    writeln!(
+                        buf,
+                        "[{}] {} [{}] {}",
+                        timestamp,
+                        record.level(),
+                        record.target(),
+                        record.args()
+                    )
+                })
+                .init();
+        }
+        config::LogFormat::Json => {
+            builder
+                .format(|buf, record| {
+                    use std::io::Write;
+                    let json = serde_json::json!({
+                        "timestamp": chrono::Local::now().to_rfc3339(),
+                        "level": record.level().to_string(),
+                        "target": record.target(),
+                        "message": record.args().to_string(),
+                        "module_path": record.module_path(),
+                        "file": record.file(),
+                        "line": record.line(),
+                    });
+                    writeln!(buf, "{}", json)
+                })
+                .init();
+        }
+        config::LogFormat::PrettyJson => {
+            builder
+                .format(|buf, record| {
+                    use std::io::Write;
+                    let json = serde_json::json!({
+                        "timestamp": chrono::Local::now().to_rfc3339(),
+                        "level": record.level().to_string(),
+                        "target": record.target(),
+                        "message": record.args().to_string(),
+                        "module_path": record.module_path(),
+                        "file": record.file(),
+                        "line": record.line(),
+                    });
+                    writeln!(buf, "{}", serde_json::to_string_pretty(&json).unwrap())
+                })
+                .init();
+        }
+    }
+
+    log::info!("Starting Eris tarpit system");
 
     // Load configuration
-    let config = if let Some(config_path) = &args.config_file {
+    let mut config = if let Some(config_path) = &args.config_file {
         log::info!("Loading configuration from {config_path:?}");
         match Config::load_from_file(config_path) {
             Ok(cfg) => {
@@ -55,6 +118,11 @@ async fn main() -> std::io::Result<()> {
         Config::from_args(&args)
     };
 
+    // Log format from the command line needs to be preserved
+    if args.config_file.is_none() {
+        config.log_format = log_format;
+    }
+
     // Ensure required directories exist
     match config.ensure_dirs_exist() {
         Ok(()) => log::info!("Directory setup completed"),