use crate::patterns::MaskRule;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use memchr::memchr2;
use rayon::prelude::*;
use regex::RegexSet;
use std::{
fs::{self, File},
io::{self, BufRead, BufReader, BufWriter, Write},
path::{Path, PathBuf},
};
pub struct Sanitizer {
rules: Vec<MaskRule>,
rule_set: RegexSet,
}
impl Sanitizer {
pub fn new() -> Self {
let rules = MaskRule::rules();
let pattern_strings: Vec<&str> = rules.iter().map(|r| r.pattern.as_str()).collect();
let rule_set = RegexSet::new(&pattern_strings).expect("Failed to build RegexSet");
Self { rules, rule_set }
}
pub fn sanitize_auto<P: AsRef<Path>>(
&self,
input_path: P,
output_dir: P,
) -> io::Result<(Vec<PathBuf>, u64)> {
let input = input_path.as_ref();
let mut log_files: Vec<PathBuf> = Vec::new();
if input.is_file() {
log_files.push(input.to_path_buf());
} else if input.is_dir() {
for entry in fs::read_dir(input)? {
let entry = entry?;
let path = entry.path();
if path.is_file() && path.extension().is_some_and(|ext| ext == "log") {
log_files.push(path);
}
}
}
self.sanitize_batch_parallel(&log_files, output_dir.as_ref())
}
fn sanitize_batch_parallel(
&self,
input_paths: &[PathBuf],
output_dir: &Path,
) -> io::Result<(Vec<PathBuf>, u64)> {
if !output_dir.exists() {
fs::create_dir_all(output_dir)?;
}
let mp = MultiProgress::new();
let results: io::Result<Vec<(PathBuf, u64)>> = input_paths
.par_iter()
.map(|input_path| {
let file_name = input_path
.file_name()
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "Invalid path"))?;
let destination_path = output_dir.join(format!("{}.sanitized", file_name.to_string_lossy()));
let total_bytes = fs::metadata(input_path)?.len();
let pb = mp.add(ProgressBar::new(total_bytes));
pb.set_style(
ProgressStyle::with_template(
"{prefix:.bold} {spinner:.green} [{wide_bar:.cyan/blue}] {bytes}/{total_bytes}",
)
.unwrap()
.progress_chars("#>-"),
);
pb.set_prefix(file_name.to_string_lossy().into_owned());
self.sanitize_file_ultra_fast(input_path, &destination_path, &pb)?;
pb.finish_with_message("Done");
Ok((destination_path, total_bytes))
})
.collect();
let results_vec = results?;
let total_bytes: u64 = results_vec.iter().map(|(_, bytes)| bytes).sum();
let paths: Vec<PathBuf> = results_vec.into_iter().map(|(path, _)| path).collect();
Ok((paths, total_bytes))
}
fn sanitize_file_ultra_fast(
&self,
input_path: &Path,
output_path: &Path,
pb: &ProgressBar,
) -> io::Result<()> {
let input_file = File::open(input_path)?;
let reader = BufReader::with_capacity(512 * 1024, input_file);
let output_file = File::create(output_path)?;
let mut writer = BufWriter::with_capacity(512 * 1024, output_file);
let mut bytes_processed: u64 = 0;
let mut line_counter: u64 = 0;
for line_result in reader.lines() {
let mut line = line_result?;
let line_bytes = line.as_bytes();
bytes_processed += line_bytes.len() as u64 + 1;
line_counter += 1;
let has_trigger_char = memchr2(b'@', b'=', line_bytes).is_some()
|| memchr2(b':', b'-', line_bytes).is_some();
if has_trigger_char && self.rule_set.is_match(&line) {
let matches = self.rule_set.matches(&line);
for index in matches.iter() {
let rule = &self.rules[index];
line = rule
.pattern
.replace_all(&line, rule.replacement)
.into_owned();
}
}
writeln!(writer, "{}", line)?;
if line_counter.is_multiple_of(512) {
pb.set_position(bytes_processed);
}
}
pb.set_position(bytes_processed);
writer.flush()?;
Ok(())
}
}