dtxt_detect/
lib.rs

1use aho_corasick::AhoCorasick;
2use deunicode::deunicode;
3
4pub struct Engine {
5    pub tier1_keys: Vec<String>,
6    pub tier2_keys: Vec<String>,
7    pub tier3_keys: Vec<String>,
8    t1_corasick: Option<AhoCorasick>,
9    t2_corasick: Option<AhoCorasick>,
10    t3_corasick: Option<AhoCorasick>,
11    pub keep_tier2_entries: bool,
12    pub keep_tier3_entries: bool,
13    pub fail_on_tier2: bool,
14    pub fail_on_tier3: bool,
15    pub fail_on_tier1_and_2: bool,
16    pub normalize_unicode: bool,
17    pub reset_string_on_fail: bool,
18}
19pub struct DtxtOutput {
20    pub string: Option<String>,
21    pub warnings: u32,
22    pub fails: u32,
23    /// Needs `keep_tier2_entries` to be true in the Engine
24    pub tier2_entries: Vec<String>,
25    /// Needs `keep_tier3_entires` to be true in the Engine
26    pub tier3_entries: Vec<String>,
27}
28
29impl Default for DtxtOutput {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl DtxtOutput {
36    pub fn new() -> DtxtOutput {
37        DtxtOutput {
38            string: None,
39            warnings: 0,
40            fails: 0,
41            tier2_entries: vec![],
42            tier3_entries: vec![],
43        }
44    }
45}
46impl Default for Engine {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52impl Engine {
53    pub fn new() -> Engine {
54        Engine {
55            tier1_keys: vec![],
56            tier2_keys: vec![],
57            tier3_keys: vec![],
58            t1_corasick: None,
59            t2_corasick: None,
60            t3_corasick: None,
61            keep_tier2_entries: false,
62            keep_tier3_entries: false,
63            fail_on_tier2: false,
64            fail_on_tier3: true,
65            fail_on_tier1_and_2: true,
66            normalize_unicode: true,
67            reset_string_on_fail: true,
68        }
69    }
70    /// Rebuild Engine, used when updating
71    /// Normally it is done by the `Engine::load` function
72    pub fn rebuild(&mut self) {
73        self.t1_corasick = Some(AhoCorasick::builder()
74            .ascii_case_insensitive(true)
75            .build(&self.tier1_keys)
76            .unwrap());
77        self.t2_corasick = Some(AhoCorasick::builder()
78            .ascii_case_insensitive(true)
79            .build(&self.tier2_keys)
80            .unwrap());
81        self.t3_corasick = Some(AhoCorasick::builder()
82            .ascii_case_insensitive(true)
83            .build(&self.tier3_keys)
84            .unwrap());
85    }
86    /// Load key datasets inside of the Engine
87    pub fn load(
88        &mut self,
89        tier1_keys: Vec<String>,
90        tier2_keys: Vec<String>,
91        tier3_keys: Vec<String>,
92        concat: bool,
93        rebuild: bool,
94    ) {
95        if concat {
96            self.tier1_keys.extend(tier1_keys);
97            self.tier2_keys.extend(tier2_keys);
98            self.tier3_keys.extend(tier3_keys);
99        } else {
100            self.tier1_keys = tier1_keys;
101            self.tier2_keys = tier2_keys;
102            self.tier3_keys = tier3_keys;
103        }
104
105        if rebuild {
106            // Needed as the Corasick's objects need to be rebuilt  with the new keys
107            self.rebuild();
108        }
109    }
110    /// Process a String using the Engine
111    pub fn process(&self, input: String) -> DtxtOutput {
112        let mut work_string = input;
113        let mut dtxt_output = DtxtOutput::new();
114
115        // Normalize unicode
116        if self.normalize_unicode {
117            work_string = deunicode(&work_string);
118        }
119
120        // Detect tier-1 words
121        let ac1 = self.t1_corasick
122            .as_ref()
123            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
124        let mut tier1_entries = vec![];
125        for mat in ac1.find_iter(&work_string) {
126            let pattern_id = mat.pattern().as_usize();
127            tier1_entries.push(self.tier1_keys[pattern_id].clone());
128        }
129
130        // Detect tier-2 words
131        let ac2 = self.t2_corasick
132            .as_ref()
133            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
134        let mut tier2_entries = vec![];
135        for mat in ac2.find_iter(&work_string) {
136            let pattern_id = mat.pattern().as_usize();
137            tier2_entries.push(self.tier2_keys[pattern_id].clone());
138        }
139
140        // Detect tier-3 words
141        let ac3 = self.t3_corasick
142            .as_ref()
143            .expect("[dtxt-detect] Engine not built. Call `.rebuild()` after loading keys.");
144        let mut tier3_entries = vec![];
145        for mat in ac3.find_iter(&work_string) {
146            let pattern_id = mat.pattern().as_usize();
147            tier3_entries.push(self.tier3_keys[pattern_id].clone());
148        }
149
150        // Check for warnings and fails
151        if self.fail_on_tier2 {
152            if !tier2_entries.is_empty() {
153                dtxt_output.fails += tier2_entries.len() as u32;
154            }
155        } else {
156            dtxt_output.warnings += tier2_entries.len() as u32;
157        }
158        if self.fail_on_tier3 {
159            if !tier3_entries.is_empty() {
160                dtxt_output.fails += tier3_entries.len() as u32;
161            }
162        } else {
163            dtxt_output.warnings += tier3_entries.len() as u32;
164        }
165        if self.fail_on_tier1_and_2
166            && !tier1_entries.is_empty()
167            && !tier2_entries.is_empty()
168            && !self.fail_on_tier2
169        {
170            // Do not consider if already failed on tier2 in order to have recounting the same failure
171            dtxt_output.fails += 0
172        }
173
174        // Register needed entries
175        if self.keep_tier2_entries {
176            dtxt_output.tier2_entries = tier2_entries;
177        }
178        if self.keep_tier3_entries {
179            dtxt_output.tier3_entries = tier3_entries;
180        }
181
182        // Return empty message if configured this way
183        if self.reset_string_on_fail && dtxt_output.fails > 0 {
184            dtxt_output.string = None;
185        } else {
186            dtxt_output.string = Some(work_string);
187        }
188        dtxt_output
189    }
190}
dtxt_detect/lib.rs

dtxt_detect/
lib.rs
