Skip to content

Commit

Permalink
Refactor color handling
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
joerivanruth committed Apr 25, 2024
1 parent 3b3ab35 commit 7d91da2
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 37 deletions.
52 changes: 52 additions & 0 deletions src/colors.rs
Original file line number Diff line number Diff line change
@@ -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"),
};
23 changes: 14 additions & 9 deletions src/main.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#![doc = include_str!("../README.md")]

mod addr;
mod colors;
mod event;
mod mapi;
mod pcap;
Expand All @@ -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;

Expand Down Expand Up @@ -53,7 +55,7 @@ fn mymain() -> AResult<()> {
let mut pcap_file: Option<PathBuf> = 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()? {
Expand All @@ -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'"),
}
}
Expand Down Expand Up @@ -102,19 +104,22 @@ fn mymain() -> AResult<()> {

args.no_more_stashed()?;

let colored_default;
let default_colors;
let out: Box<dyn io::Write + Send + 'static> = 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);

Expand Down
1 change: 1 addition & 0 deletions src/mapi/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down
67 changes: 39 additions & 28 deletions src/render.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Box<dyn io::Write + 'static + Send>>,
current_style: Style,
Expand All @@ -21,10 +25,11 @@ pub struct Renderer {
}

impl Renderer {
pub fn new(colored: bool, out: Box<dyn io::Write + 'static + Send>) -> Self {
pub fn new(colors: &'static Colors, out: Box<dyn io::Write + 'static + Send>) -> 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,
Expand Down Expand Up @@ -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()?;
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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;
Expand All @@ -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(())
}
}
Expand Down

0 comments on commit 7d91da2

Please sign in to comment.