use super::*;
#[derive(Debug, Default)]
pub(crate) struct StreamingControlSanitizer {
pub(super) state: ControlState,
}
#[derive(Debug, Default)]
pub(super) enum ControlState {
#[default]
Normal,
Escape,
Csi,
String {
escape: bool,
},
}
impl StreamingControlSanitizer {
pub(crate) fn push(&mut self, text: &str) -> String {
let mut out = String::new();
for ch in text.chars() {
match &mut self.state {
ControlState::Normal => {
if ch == '\u{1b}' {
self.state = ControlState::Escape;
} else if !ch.is_control() || matches!(ch, '\n' | '\t' | '\r') {
out.push(ch);
}
}
ControlState::Escape => {
self.state = match ch {
'[' => ControlState::Csi,
']' | 'P' | '^' | '_' | 'X' => ControlState::String { escape: false },
_ => ControlState::Normal,
};
}
ControlState::Csi => {
if ('@'..='~').contains(&ch) {
self.state = ControlState::Normal;
}
}
ControlState::String { escape } => {
if ch == '\u{7}' {
self.state = ControlState::Normal;
} else if *escape {
self.state = if ch == '\\' {
ControlState::Normal
} else if ch == '\u{1b}' {
ControlState::String { escape: true }
} else {
ControlState::String { escape: false }
};
} else if ch == '\u{1b}' {
*escape = true;
}
}
}
}
out
}
pub(crate) fn reset(&mut self) {
self.state = ControlState::Normal;
}
}
pub(crate) fn sanitize_preview(text: &str) -> String {
redact_sensitive_text(strip_control_sequences(text).as_str())
}
fn strip_control_sequences(text: &str) -> String {
let mut out = String::new();
let mut chars = text.chars().peekable();
while let Some(ch) = chars.next() {
if ch != '\u{1b}' {
if ch.is_control() && !matches!(ch, '\n' | '\t' | '\r') {
continue;
}
out.push(ch);
continue;
}
match chars.peek().copied() {
Some('[') => {
chars.next();
consume_csi(&mut chars);
}
Some(']') | Some('P') | Some('^') | Some('_') | Some('X') => {
chars.next();
consume_string_control(&mut chars);
}
Some(ch) if ('@'..='_').contains(&ch) => {
chars.next();
}
Some(_) => {
chars.next();
}
None => {}
}
}
out
}
fn consume_csi(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
for ch in chars.by_ref() {
if ('@'..='~').contains(&ch) {
break;
}
}
}
fn consume_string_control(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
while let Some(ch) = chars.next() {
if ch == '\u{7}' {
break;
}
if ch == '\u{1b}' && chars.peek().copied() == Some('\\') {
chars.next();
break;
}
}
}