diff --git a/src/output/console.rs b/src/output/console.rs index 1d12d12..97d63b5 100644 --- a/src/output/console.rs +++ b/src/output/console.rs @@ -6,7 +6,7 @@ use std::io::{self, BufWriter, Write}; use std::path::Path; use std::sync::Mutex; -use super::Output; +use super::{escape_csv_field, Output}; use crate::derive::DerivedKey; use crate::matcher::MatchInfo; @@ -76,8 +76,13 @@ impl Output for ConsoleOutput { // Compact format: source,transform,privkey,address_compressed writeln!( w, - "{},{},{},{}", - source, transform, derived.private_key_hex, derived.p2pkh_compressed + "{}", + format_compact_csv_row( + source, + transform, + &derived.private_key_hex, + &derived.p2pkh_compressed, + ) )?; } @@ -122,6 +127,19 @@ impl Output for ConsoleOutput { } } +fn format_compact_csv_row( + source: &str, + transform: &str, + private_key: &str, + address: &str, +) -> String { + [source, transform, private_key, address] + .into_iter() + .map(escape_csv_field) + .collect::>() + .join(",") +} + #[cfg(test)] mod tests { use super::*; @@ -164,6 +182,48 @@ mod tests { assert!(content.contains("test_source,sha256,abc123,1Address")); } + #[test] + fn test_to_file_writes_compact_format_with_commas_in_source() { + let temp = NamedTempFile::new().unwrap(); + let output = ConsoleOutput::to_file(temp.path()).unwrap(); + + output + .key("test,source", "sha256", &make_test_key()) + .unwrap(); + output.flush().unwrap(); + + let content = std::fs::read_to_string(temp.path()).unwrap(); + assert!(content.contains("\"test,source\",sha256,abc123,1Address")); + } + + #[test] + fn test_to_file_writes_compact_format_with_quotes_in_source() { + let temp = NamedTempFile::new().unwrap(); + let output = ConsoleOutput::to_file(temp.path()).unwrap(); + + output + .key("say \"hello\"", "sha256", &make_test_key()) + .unwrap(); + output.flush().unwrap(); + + let content = std::fs::read_to_string(temp.path()).unwrap(); + assert!(content.contains("\"say \"\"hello\"\"\",sha256,abc123,1Address")); + } + + #[test] + fn test_to_file_writes_compact_format_with_newlines_in_source() { + let temp = NamedTempFile::new().unwrap(); + let output = ConsoleOutput::to_file(temp.path()).unwrap(); + let source = "line1\nline2"; + + output.key(source, "sha256", &make_test_key()).unwrap(); + output.flush().unwrap(); + + let content = std::fs::read_to_string(temp.path()).unwrap(); + let expected_row = format_compact_csv_row(source, "sha256", "abc123", "1Address"); + assert_eq!(content, format!("{}\n", expected_row)); + } + #[test] fn test_to_file_verbose_writes_yaml_format() { let temp = NamedTempFile::new().unwrap(); diff --git a/src/output/mod.rs b/src/output/mod.rs index d91d342..0f30e54 100644 --- a/src/output/mod.rs +++ b/src/output/mod.rs @@ -18,6 +18,22 @@ use crate::derive::DerivedKey; use crate::matcher::MatchInfo; use anyhow::Result; +pub(crate) fn escape_csv_field(field: &str) -> String { + let has_boundary_whitespace = field.chars().next().is_some_and(char::is_whitespace) + || field.chars().last().is_some_and(char::is_whitespace); + + if has_boundary_whitespace + || field.contains(',') + || field.contains('"') + || field.contains('\n') + || field.contains('\r') + { + format!("\"{}\"", field.replace('"', "\"\"")) + } else { + field.to_string() + } +} + /// Output trait for handling generated keys. pub trait Output: Send + Sync { /// Output a key (no matcher, output all keys). diff --git a/src/output/query_format.rs b/src/output/query_format.rs index 9d7b075..b8cdec7 100644 --- a/src/output/query_format.rs +++ b/src/output/query_format.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use arrow::datatypes::Schema; use comfy_table::{presets::UTF8_FULL, Cell, Color, ContentArrangement, Table}; +use super::escape_csv_field; use crate::storage::{Row, Value}; #[derive(Clone, Copy, Debug, Default, PartialEq)] @@ -216,14 +217,6 @@ fn format_value_csv(value: Option<&Value>) -> String { } } -fn escape_csv_field(field: &str) -> String { - if field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r') { - format!("\"{}\"", field.replace('"', "\"\"")) - } else { - field.to_string() - } -} - pub fn format_schema(schema: &Schema) -> String { let mut table = Table::new(); table @@ -270,8 +263,14 @@ mod tests { #[test] fn output_format_from_str() { - assert_eq!("table".parse::().unwrap(), OutputFormat::Table); - assert_eq!("TABLE".parse::().unwrap(), OutputFormat::Table); + assert_eq!( + "table".parse::().unwrap(), + OutputFormat::Table + ); + assert_eq!( + "TABLE".parse::().unwrap(), + OutputFormat::Table + ); assert_eq!("json".parse::().unwrap(), OutputFormat::Json); assert_eq!("JSON".parse::().unwrap(), OutputFormat::Json); assert_eq!("csv".parse::().unwrap(), OutputFormat::Csv); @@ -375,6 +374,22 @@ mod tests { assert_eq!(escape_csv_field("say \"hi\""), "\"say \"\"hi\"\"\""); } + #[test] + fn escape_csv_field_with_boundary_whitespace() { + assert_eq!(escape_csv_field(" leading"), "\" leading\""); + assert_eq!(escape_csv_field("trailing "), "\"trailing \""); + assert_eq!(escape_csv_field("\ttab"), "\"\ttab\""); + assert_eq!(escape_csv_field("tab\t"), "\"tab\t\""); + assert_eq!( + escape_csv_field("\u{00A0}nbsp-leading"), + "\"\u{00A0}nbsp-leading\"" + ); + assert_eq!( + escape_csv_field("nbsp-trailing\u{00A0}"), + "\"nbsp-trailing\u{00A0}\"" + ); + } + #[test] fn format_value_json_types() { assert_eq!(format_value_json(&Value::Null), "null");