threatflux_string_analysis/
categorizer.rs

1//! String categorization functionality.
2
3use crate::types::{
4    AnalysisError, AnalysisResult, MAX_DESCRIPTION_BYTES, compact_string, validate_identifier,
5};
6use regex::Regex;
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeSet;
9use std::net::IpAddr;
10use std::sync::LazyLock;
11
12const MAX_CATEGORY_RULES: usize = 4_096;
13
14static EMAIL_REGEX: LazyLock<Regex> = LazyLock::new(|| {
15    Regex::new(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
16        .expect("the built-in email regex must compile")
17});
18
19/// Thread-safe predicate used by a [`CategoryRule`].
20pub type CategoryMatcher = Box<dyn Fn(&str) -> bool + Send + Sync>;
21
22/// A named category that can be assigned to a string.
23#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
24#[serde(deny_unknown_fields)]
25pub struct StringCategory {
26    /// Stable category name.
27    pub name: String,
28    /// Optional parent category name.
29    pub parent: Option<String>,
30    /// Human-readable category description.
31    pub description: String,
32}
33
34impl StringCategory {
35    pub(crate) fn compact(mut self) -> Self {
36        self.name = compact_string(self.name);
37        self.parent = self.parent.map(compact_string);
38        self.description = compact_string(self.description);
39        self
40    }
41}
42
43/// Rule for categorizing strings.
44pub struct CategoryRule {
45    /// Unique rule name.
46    pub name: String,
47    /// Predicate that determines whether a string matches.
48    pub matcher: CategoryMatcher,
49    /// Category assigned on a match.
50    pub category: StringCategory,
51    /// Priority; larger values are evaluated first.
52    pub priority: i32,
53}
54
55impl CategoryRule {
56    fn validate(&self) -> AnalysisResult<()> {
57        validate_identifier("category rule", &self.name)?;
58        validate_category(&self.category)
59    }
60}
61
62/// Interface for deterministic, thread-safe string categorizers.
63pub trait Categorizer: Send + Sync {
64    /// Categorize a string in deterministic rule order.
65    fn categorize(&self, value: &str) -> Vec<StringCategory>;
66
67    /// Validate and add a uniquely named rule.
68    fn add_rule(&mut self, rule: CategoryRule) -> AnalysisResult<()>;
69
70    /// Remove an existing rule by name.
71    fn remove_rule(&mut self, name: &str) -> AnalysisResult<()>;
72
73    /// Return distinct categories in deterministic rule order.
74    fn get_categories(&self) -> Vec<StringCategory>;
75}
76
77/// Default heuristic categorizer.
78pub struct DefaultCategorizer {
79    rules: Vec<CategoryRule>,
80}
81
82impl DefaultCategorizer {
83    /// Create a categorizer with the built-in informational rules.
84    pub fn new() -> Self {
85        let mut categorizer = Self::empty();
86        categorizer.add_default_rules();
87        categorizer
88    }
89
90    /// Create a categorizer without built-in rules.
91    pub fn empty() -> Self {
92        Self { rules: Vec::new() }
93    }
94
95    fn add_default_rules(&mut self) {
96        self.rules = vec![
97            rule(
98                "url",
99                "url",
100                "network",
101                "URL or web address",
102                100,
103                |value| {
104                    [
105                        "http://",
106                        "https://",
107                        "ftp://",
108                        "ssh://",
109                        "telnet://",
110                        "rdp://",
111                    ]
112                    .iter()
113                    .any(|scheme| starts_with_ascii_case(value, scheme))
114                },
115            ),
116            rule(
117                "registry",
118                "registry",
119                "windows",
120                "Windows registry key",
121                95,
122                |value| {
123                    starts_with_ascii_case(value, "HKEY_")
124                        || contains_ascii_case(value, "\\SOFTWARE\\")
125                },
126            ),
127            rule(
128                "ip_address",
129                "ip_address",
130                "network",
131                "Syntactically valid IPv4 or IPv6 address",
132                95,
133                |value| value.parse::<IpAddr>().is_ok(),
134            ),
135            rule(
136                "path",
137                "path",
138                "filesystem",
139                "Absolute or drive-qualified file-system path",
140                90,
141                |value| {
142                    value.starts_with('/')
143                        || value.starts_with('\\')
144                        || (value.len() >= 3
145                            && value.as_bytes()[1] == b':'
146                            && matches!(value.as_bytes()[2], b'\\' | b'/'))
147                },
148            ),
149            rule(
150                "api_call",
151                "api_call",
152                "system",
153                "Known system API function name",
154                90,
155                is_known_api_call,
156            ),
157            rule(
158                "library",
159                "library",
160                "binary",
161                "Shared library or DLL name",
162                85,
163                |value| {
164                    ends_with_ascii_case(value, ".dll")
165                        || ends_with_ascii_case(value, ".so")
166                        || ends_with_ascii_case(value, ".dylib")
167                        || contains_ascii_case(value, ".so.")
168                },
169            ),
170            rule(
171                "email",
172                "email",
173                "contact",
174                "Email-address-shaped string",
175                85,
176                |value| EMAIL_REGEX.is_match(value),
177            ),
178            rule(
179                "command",
180                "command",
181                "execution",
182                "Command or shell interpreter reference",
183                80,
184                is_command_reference,
185            ),
186        ];
187        sort_rules(&mut self.rules);
188    }
189}
190
191impl Categorizer for DefaultCategorizer {
192    fn categorize(&self, value: &str) -> Vec<StringCategory> {
193        let mut seen = BTreeSet::new();
194        let mut categories = Vec::new();
195        for rule in &self.rules {
196            if (rule.matcher)(value) && seen.insert(rule.category.name.clone()) {
197                categories.push(rule.category.clone());
198            }
199        }
200
201        if categories.is_empty() {
202            categories.push(StringCategory {
203                name: "generic".to_string(),
204                parent: None,
205                description: "Generic string".to_string(),
206            });
207        }
208        categories
209    }
210
211    fn add_rule(&mut self, mut rule: CategoryRule) -> AnalysisResult<()> {
212        rule.validate()?;
213        if self.rules.len() >= MAX_CATEGORY_RULES {
214            return Err(AnalysisError::CapacityExceeded {
215                resource: "category rules",
216                limit: MAX_CATEGORY_RULES,
217            });
218        }
219        if self.rules.iter().any(|existing| existing.name == rule.name) {
220            return Err(AnalysisError::DuplicateName {
221                kind: "category rule",
222                name: rule.name.to_string(),
223            });
224        }
225        rule.name = compact_string(rule.name);
226        rule.category = rule.category.compact();
227        self.rules.push(rule);
228        sort_rules(&mut self.rules);
229        Ok(())
230    }
231
232    fn remove_rule(&mut self, name: &str) -> AnalysisResult<()> {
233        validate_identifier("category rule", name)?;
234        let Some(index) = self.rules.iter().position(|rule| rule.name == name) else {
235            return Err(AnalysisError::NotFound {
236                kind: "category rule",
237                name: name.to_string(),
238            });
239        };
240        self.rules.remove(index);
241        Ok(())
242    }
243
244    fn get_categories(&self) -> Vec<StringCategory> {
245        let mut seen = BTreeSet::new();
246        self.rules
247            .iter()
248            .filter(|rule| seen.insert(rule.category.name.clone()))
249            .map(|rule| rule.category.clone())
250            .collect()
251    }
252}
253
254impl Default for DefaultCategorizer {
255    fn default() -> Self {
256        Self::new()
257    }
258}
259
260pub(crate) fn validate_category(category: &StringCategory) -> AnalysisResult<()> {
261    validate_identifier("category", &category.name)?;
262    if let Some(parent) = &category.parent {
263        validate_identifier("parent category", parent)?;
264    }
265    if category.description.len() > MAX_DESCRIPTION_BYTES {
266        return Err(AnalysisError::InputTooLarge {
267            field: "category.description",
268            actual: category.description.len(),
269            limit: MAX_DESCRIPTION_BYTES,
270        });
271    }
272    if category.description.trim().is_empty() {
273        return Err(AnalysisError::InvalidIdentifier {
274            kind: "category description",
275            name: category.description.clone(),
276            reason: "must not be empty or whitespace-only",
277        });
278    }
279    if category.description.chars().any(char::is_control) {
280        return Err(AnalysisError::InvalidIdentifier {
281            kind: "category description",
282            name: category.description.to_string(),
283            reason: "must not contain control characters",
284        });
285    }
286    Ok(())
287}
288
289fn sort_rules(rules: &mut [CategoryRule]) {
290    rules.sort_by(|left, right| {
291        right
292            .priority
293            .cmp(&left.priority)
294            .then_with(|| left.name.cmp(&right.name))
295    });
296}
297
298fn rule(
299    rule_name: &str,
300    category_name: &str,
301    parent: &str,
302    description: &str,
303    priority: i32,
304    matcher: impl Fn(&str) -> bool + Send + Sync + 'static,
305) -> CategoryRule {
306    CategoryRule {
307        name: rule_name.to_string(),
308        matcher: Box::new(matcher),
309        category: StringCategory {
310            name: category_name.to_string(),
311            parent: Some(parent.to_string()),
312            description: description.to_string(),
313        },
314        priority,
315    }
316}
317
318fn is_command_reference(value: &str) -> bool {
319    value
320        .split(|character: char| !character.is_ascii_alphanumeric() && character != '.')
321        .any(|token| {
322            [
323                "cmd",
324                "cmd.exe",
325                "powershell",
326                "powershell.exe",
327                "pwsh",
328                "pwsh.exe",
329                "bash",
330                "dash",
331                "zsh",
332                "ksh",
333                "sh",
334            ]
335            .iter()
336            .any(|command| token.eq_ignore_ascii_case(command))
337        })
338}
339
340fn starts_with_ascii_case(value: &str, prefix: &str) -> bool {
341    value
342        .as_bytes()
343        .get(..prefix.len())
344        .is_some_and(|candidate| candidate.eq_ignore_ascii_case(prefix.as_bytes()))
345}
346
347fn ends_with_ascii_case(value: &str, suffix: &str) -> bool {
348    value
349        .as_bytes()
350        .get(value.len().saturating_sub(suffix.len())..)
351        .is_some_and(|candidate| candidate.eq_ignore_ascii_case(suffix.as_bytes()))
352}
353
354fn contains_ascii_case(value: &str, needle: &str) -> bool {
355    value
356        .as_bytes()
357        .windows(needle.len())
358        .any(|candidate| candidate.eq_ignore_ascii_case(needle.as_bytes()))
359}
360
361fn is_known_api_call(value: &str) -> bool {
362    matches!(
363        value,
364        "CreateProcess"
365            | "CreateProcessA"
366            | "CreateProcessW"
367            | "VirtualAlloc"
368            | "VirtualAllocEx"
369            | "WriteProcessMemory"
370            | "GetProcAddress"
371            | "LoadLibrary"
372            | "LoadLibraryA"
373            | "LoadLibraryW"
374            | "OpenProcess"
375            | "CreateRemoteThread"
376            | "malloc"
377            | "calloc"
378            | "realloc"
379            | "free"
380            | "fork"
381            | "exec"
382            | "open"
383            | "read"
384            | "write"
385    )
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[test]
393    fn added_rules_compact_owned_metadata() {
394        let mut name = String::with_capacity(1_000_000);
395        name.push_str("compact_rule");
396        let mut category_name = String::with_capacity(1_000_000);
397        category_name.push_str("compact_category");
398        let mut parent = String::with_capacity(1_000_000);
399        parent.push_str("parent");
400        let mut description = String::with_capacity(1_000_000);
401        description.push_str("Category compaction regression");
402        let mut categorizer = DefaultCategorizer::empty();
403        categorizer
404            .add_rule(CategoryRule {
405                name,
406                matcher: Box::new(|_| true),
407                category: StringCategory {
408                    name: category_name,
409                    parent: Some(parent),
410                    description,
411                },
412                priority: 0,
413            })
414            .unwrap();
415
416        let rule = &categorizer.rules[0];
417        assert_eq!(rule.name.capacity(), rule.name.len());
418        assert_eq!(rule.category.name.capacity(), rule.category.name.len());
419        assert_eq!(
420            rule.category.parent.as_ref().unwrap().capacity(),
421            rule.category.parent.as_ref().unwrap().len()
422        );
423        assert_eq!(
424            rule.category.description.capacity(),
425            rule.category.description.len()
426        );
427    }
428}
threatflux_string_analysis/categorizer.rs

threatflux_string_analysis/
categorizer.rs
