From 7d91da223ee2c833fb8247114a853133f5ae8132 Mon Sep 17 00:00:00 2001 From: Joeri van Ruth Date: Thu, 25 Apr 2024 11:15:01 +0200 Subject: [PATCH] Refactor color handling Choose between multiple versions of a Colors object containing escape sequences rather than having a 'colored' boolean. Later on, html will just be another Colors object. --- src/colors.rs | 52 ++++++++++++++++++++++++++++++++++++++ src/main.rs | 23 ++++++++++------- src/mapi/mod.rs | 1 + src/render.rs | 67 ++++++++++++++++++++++++++++--------------------- 4 files changed, 106 insertions(+), 37 deletions(-) create mode 100644 src/colors.rs diff --git a/src/colors.rs b/src/colors.rs new file mode 100644 index 0000000..1e25275 --- /dev/null +++ b/src/colors.rs @@ -0,0 +1,52 @@ +pub struct Colors { + pub normal: EscapeSequence, + pub red: EscapeSequence, + pub green: EscapeSequence, + pub cyan: EscapeSequence, + pub blue: EscapeSequence, + pub bold: EscapeSequence, +} + +#[derive(Debug, Clone, Copy)] +pub struct EscapeSequence { + pub enable: &'static str, + pub disable: &'static str, +} + +impl EscapeSequence { + pub const fn new(enable: &'static str, disable: &'static str) -> Self { + EscapeSequence { enable, disable } + } +} + +const EMPTY: EscapeSequence = EscapeSequence::new("", ""); + +pub static NO_COLORS: &Colors = &Colors { + normal: EMPTY, + bold: EMPTY, + red: EMPTY, + green: EMPTY, + cyan: EMPTY, + blue: EMPTY, +}; + +#[allow(dead_code)] +pub static DEBUG_COLORS: &Colors = &Colors { + normal: EscapeSequence::new("«-»", ""), + red: EscapeSequence::new("«red»", ""), + green: EscapeSequence::new("«green»", ""), + cyan: EscapeSequence::new("«cyan»", ""), + blue: EscapeSequence::new("«blue»", ""), + bold: EscapeSequence::new("«bold»", "«/bold»"), +}; + +// Black=30 Red=31 Green=32 Yellow=33 Blue=34 Magenta=35 Cyan=36 White=37 + +pub static VT100_COLORS: &Colors = &Colors { + normal: EscapeSequence::new("\u{1b}[39m", ""), + red: EscapeSequence::new("\u{1b}[31m", ""), + green: EscapeSequence::new("\u{1b}[32m", ""), + cyan: EscapeSequence::new("\u{1b}[36m", ""), + blue: EscapeSequence::new("\u{1b}[34m", ""), + bold: EscapeSequence::new("\u{1b}[1m", "\u{1b}[0m"), +}; diff --git a/src/main.rs b/src/main.rs index 2940009..9180f64 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ #![doc = include_str!("../README.md")] mod addr; +mod colors; mod event; mod mapi; mod pcap; @@ -17,6 +18,7 @@ use std::{io, panic, process, thread}; use addr::MonetAddr; use anyhow::{bail, Context, Result as AResult}; use argsplitter::{ArgError, ArgSplitter}; +use colors::{NO_COLORS, VT100_COLORS}; use event::{MapiEvent, Timestamp}; use pcap::Tracker; @@ -53,7 +55,7 @@ fn mymain() -> AResult<()> { let mut pcap_file: Option = None; let mut level = None; let mut force_binary = false; - let mut colored = None; + let mut colors = None; let mut args = ArgSplitter::from_env(); while let Some(flag) = args.flag()? { @@ -65,10 +67,10 @@ fn mymain() -> AResult<()> { "-r" | "--raw" => level = Some(Level::Raw), "-B" | "--binary" => force_binary = true, "--color" => { - colored = match args.param()?.to_lowercase().as_str() { - "always" => Some(true), + colors = match args.param()?.to_lowercase().as_str() { + "always" => Some(VT100_COLORS), "auto" => None, - "never" => Some(false), + "never" => Some(NO_COLORS), other => bail!("--color={other}: must be 'always', 'auto' or 'never'"), } } @@ -102,19 +104,22 @@ fn mymain() -> AResult<()> { args.no_more_stashed()?; - let colored_default; + let default_colors; let out: Box = if let Some(p) = output_file { let f = File::create(&p) .with_context(|| format!("could not open output file {}", p.display()))?; - colored_default = false; + default_colors = NO_COLORS; Box::new(f) } else { let out = io::stdout(); - colored_default = is_terminal::is_terminal(&out); + default_colors = if is_terminal::is_terminal(&out) { + VT100_COLORS + } else { + NO_COLORS + }; Box::new(out) }; - let colored = colored.unwrap_or(colored_default); - let mut renderer = Renderer::new(colored, out); + let mut renderer = Renderer::new(colors.unwrap_or(default_colors), out); let mapi_state = mapi::State::new(level, force_binary); diff --git a/src/mapi/mod.rs b/src/mapi/mod.rs index ef2ef54..daf4139 100644 --- a/src/mapi/mod.rs +++ b/src/mapi/mod.rs @@ -424,6 +424,7 @@ impl Binary { renderer.style(style); renderer.put([hi, lo])?; } + renderer.style(Style::Normal); for i in self.col..16 { self.put_sep(i, &mut cur_head, Style::Frame, renderer)?; diff --git a/src/render.rs b/src/render.rs index 9297283..a2b93f6 100644 --- a/src/render.rs +++ b/src/render.rs @@ -9,10 +9,14 @@ use std::{ use chrono::{DateTime, Local}; -use crate::event::{ConnectionId, Direction, Timestamp}; +use crate::{ + colors::{Colors, EscapeSequence}, + event::{ConnectionId, Direction, Timestamp}, +}; pub struct Renderer { - colored: bool, + colors: &'static Colors, + color_stack: Vec<&'static EscapeSequence>, timing: TrackTime, out: BufWriter>, current_style: Style, @@ -21,10 +25,11 @@ pub struct Renderer { } impl Renderer { - pub fn new(colored: bool, out: Box) -> Self { + pub fn new(colors: &'static Colors, out: Box) -> Self { let buffered = BufWriter::with_capacity(4 * 8192, out); Renderer { - colored, + colors, + color_stack: vec![], out: buffered, current_style: Style::Normal, desired_style: Style::Normal, @@ -64,7 +69,7 @@ impl Renderer { message: &dyn Display, ) -> Result<(), io::Error> { self.style(Style::Frame); - self.fix_style()?; + self.switch_style()?; write!(self.out, "‣{} {message}", context)?; self.nl()?; self.out.flush()?; @@ -78,7 +83,7 @@ impl Renderer { ) -> io::Result<()> { self.render_timing()?; self.style(Style::Frame); - self.fix_style()?; + self.switch_style()?; write!(self.out, "┌{}", context.into())?; let mut sep = " "; for item in items { @@ -95,7 +100,7 @@ impl Renderer { self.nl()?; } self.style(Style::Frame); - self.fix_style()?; + self.switch_style()?; write!(self.out, "└")?; let mut sep = " "; for item in items { @@ -110,19 +115,19 @@ impl Renderer { pub fn put(&mut self, data: impl AsRef<[u8]>) -> io::Result<()> { if self.at_start { let old_style = self.style(Style::Frame); - self.fix_style()?; + self.switch_style()?; self.out.write_all("│".as_bytes())?; self.style(old_style); self.at_start = false; } - self.fix_style()?; + self.switch_style()?; self.out.write_all(data.as_ref())?; Ok(()) } pub fn nl(&mut self) -> io::Result<()> { let old_style = self.style(Style::Normal); - self.fix_style()?; + self.switch_style()?; writeln!(self.out)?; self.style(old_style); self.at_start = true; @@ -138,31 +143,37 @@ impl Renderer { style } - fn fix_style(&mut self) -> io::Result<()> { - if self.current_style == self.desired_style { + fn switch_style(&mut self) -> io::Result<()> { + let style = self.desired_style; + if style == self.current_style { return Ok(()); } - if self.colored { - self.write_escape_sequence(self.desired_style)?; + + while let Some(sequence) = self.color_stack.pop() { + self.out.write_all(sequence.disable.as_bytes())?; + } + + let colors = self.colors; + match style { + Style::Normal => self.push_style(&colors.normal)?, + Style::Error => { + self.push_style(&colors.red)?; + self.push_style(&colors.bold)?; + } + Style::Frame => self.push_style(&colors.cyan)?, + Style::Header => self.push_style(&colors.bold)?, + Style::Whitespace => self.push_style(&colors.red)?, + Style::Digit => self.push_style(&colors.green)?, + Style::Letter => self.push_style(&colors.blue)?, } + self.current_style = self.desired_style; Ok(()) } - fn write_escape_sequence(&mut self, style: Style) -> io::Result<()> { - // Black=30 Red=31 Green=32 Yellow=33 Blue=34 Magenta=35 Cyan=36 White=37 - - let escape_sequence = match style { - Style::Normal => "", - Style::Header => "\u{1b}[1m", // bold - Style::Frame => "\u{1b}[36m", // cyan - Style::Error => "\u{1b}[1m\u{1b}[31m", // bold red - Style::Whitespace => "\u{1b}[31m", // red - Style::Digit => "\u{1b}[32m", // green - Style::Letter => "\u{1b}[34m", // blue - }; - self.out.write_all(b"\x1b[m")?; // NORMAL - self.out.write_all(escape_sequence.as_bytes())?; + fn push_style(&mut self, seq: &'static EscapeSequence) -> io::Result<()> { + self.out.write_all(seq.enable.as_bytes())?; + self.color_stack.push(seq); Ok(()) } }