string_analyze/

lib.rs

pub mod rules;
mod entropy;
mod tool;

use rayon::prelude::*;
use std::borrow::Cow;
use std::fmt;
use std::fmt::Formatter;

pub type UsizeRange = (usize, usize);

#[derive(Debug)]
pub struct Hit<'a, D> {
    pub describe: Cow<'a, str>,
    pub importance: u8,
    pub data: D,
}

impl<'a> fmt::Display for Hit<'a, RuleResult<'a>> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (reset, yellow, grey) = if f.alternate() {
            ("\x1b[0m", "\x1b[33m", "\x1b[90m")
        } else {
            ("", "", "")
        };

        // 打印 Hit 头信息
        writeln!(
            f,
            "    {}↳{} [{}{}{}]",
            grey, reset, yellow, self.describe, reset
        )?;

        // 打印 Hit 捕获的具体内容
        for (i, extracted) in self.data.discovered.iter().enumerate() {
            let cleaned_extracted = extracted.replace('\n', "\n          ");
            write!(f, "        {}↳{} {}", grey, reset, cleaned_extracted)?;
            // 如果不是最后一条数据，换行
            if i < self.data.discovered.len() - 1 {
                writeln!(f)?;
            }
        }

        Ok(())
    }
}

#[derive(Debug)]
pub struct State<'a> {
    pub input: &'a str,
    pub ranges: Vec<UsizeRange>,
}

impl<'a> State<'a> {
    pub fn new(input: &'a str) -> Self {
        Self {
            input,
            ranges: Vec::new(),
        }
    }

    pub fn retain<F>(&mut self, mut f: F)
    where
        F: FnMut(&'a str, UsizeRange) -> bool,
    {
        let input = self.input;
        self.ranges.retain(|&r| f(input, r));
    }
}
impl<'a> From<&'a str> for State<'a> {
    fn from(s: &'a str) -> State<'a> {
        State::new(s)
    }
}

#[derive(Debug)]
pub struct RuleResult<'a> {
    pub complete_input: &'a str,
    pub ranges: Vec<UsizeRange>,
    pub discovered: Vec<&'a str>,
}

impl<'a> From<State<'a>> for RuleResult<'a> {
    fn from(state: State<'a>) -> Self {
        let discovered = state
            .ranges
            .iter()
            .map(|&(s, e)| &state.input[s..e])
            .collect();
        Self {
            complete_input: state.input,
            ranges: state.ranges,
            discovered,
        }
    }
}

pub struct AnyRule {
    pub flow: Vec<Box<dyn for<'a> Fn(&mut State<'a>) -> bool + Send + Sync>>,
    pub out: Box<dyn for<'a> Fn(State<'a>) -> Hit<'a, RuleResult<'a>> + Send + Sync>,
}

impl AnyRule {
    pub fn new<O>(out: O) -> Self
    where
        O: for<'a> Fn(State<'a>) -> Hit<'a, RuleResult<'a>> + Send + Sync + 'static,
    {
        Self {
            flow: Vec::new(),
            out: Box::new(out),
        }
    }

    pub fn add_flow<F>(mut self, flow: F) -> Self
    where
        F: for<'a> Fn(&mut State<'a>) -> bool + Send + Sync + 'static,
    {
        self.flow.push(Box::new(flow));
        self
    }

    pub fn detect<'a>(&self, input: &'a str) -> Option<Hit<'a, RuleResult<'a>>> {
        let mut state = State::new(input);

        for flow in &self.flow {
            if !flow(&mut state) {
                return None;
            }
        }
        Some((self.out)(state))
    }
}

pub fn analyze_with<'a>(input: &'a str, rules: &[AnyRule]) -> AnalyzeResult<'a> {
    AnalyzeResult {
        input,
        results: rules
            .par_iter()
            .filter_map(|rule| rule.detect(input))
            .collect(),
    }
}

pub struct AnalyzeResult<'a> {
    pub input: &'a str,
    pub results: Vec<Hit<'a, RuleResult<'a>>>,
}

impl<'a> AnalyzeResult<'a> {
    pub fn new(input: &'a str, results: Vec<Hit<'a, RuleResult<'a>>>) -> Self {
        Self { input, results }
    }

    /// 计算概率融合后的风险得分
    pub fn score(&self) -> u8 {
        if self.results.is_empty() {
            return 0;
        }
        let mut fail_prob = 1.0f64;
        for hit in &self.results {
            let prob = (hit.importance as f64) / 100.0;
            fail_prob *= 1.0 - prob;
        }
        ((1.0 - fail_prob) * 100.0).round() as u8
    }
}

impl fmt::Display for AnalyzeResult<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        if self.results.is_empty() {
            return write!(f, "{}", self.input);
        }

        // 借助 f.alternate() (即使用 {:#}) 来判断是否应用 ANSI 颜色
        let ansi = f.alternate();
        let (reset, red) = if ansi {
            ("\x1b[0m", "\x1b[31m")
        } else {
            ("", "")
        };

        // 提取并区间合并
        let mut all_ranges: Vec<UsizeRange> = self
            .results
            .iter()
            .flat_map(|hit| hit.data.ranges.iter().copied())
            .collect();

        all_ranges.sort_unstable_by_key(|r| r.0);

        let mut merged_ranges: Vec<UsizeRange> = Vec::with_capacity(all_ranges.len());
        for r in all_ranges {
            if let Some(last) = merged_ranges.last_mut() {
                if r.0 <= last.1 {
                    last.1 = last.1.max(r.1);
                    continue;
                }
            }
            merged_ranges.push(r);
        }

        // 直接向 formatter 写入，免去巨大的 String 内存分配开销
        let mut cursor = 0;
        for &(start, end) in &merged_ranges {
            if !self.input.is_char_boundary(start) || !self.input.is_char_boundary(end) {
                continue;
            }

            if start > cursor {
                write!(f, "{}", &self.input[cursor..start])?;
            }
            write!(f, "{}{}{}", red, &self.input[start..end], reset)?;

            cursor = end;
        }
        if cursor < self.input.len() {
            write!(f, "{}", &self.input[cursor..])?;
        }

        writeln!(f)?; // 在原文本和分析报告之间加一个换行

        // 借助 Hit 的 fmt::Display 实现生成 Hits 列表报告
        for (i, hit) in self.results.iter().enumerate() {
            if ansi {
                write!(f, "{:#}", hit)?;
            } else {
                write!(f, "{}", hit)?;
            }

            if i < self.results.len() - 1 {
                writeln!(f)?;
            }
        }

        Ok(())
    }
}


#[cfg(test)]
mod tests {
    use crate::rules::ALL_RULES;
    use super::*;

    #[test]
    fn test_comprehensive() {
        let input = r#"
Base64: mTyqm7wjODkrNLcWl0eqO8K8gc1BPk1GNLgUpI== 444 m7wjODkrNLcWl0eqO8K8gc1BPk1GNLgUpI==
api_key =
"12345678abcdefgh";
        "#;

        let result = analyze_with(input, ALL_RULES.as_slice());
        let score = result.score();

        println!("[综合测试评分: {score}]\n{:#}", result);
    }
}