tivilsta/
lib.rs

1// Tivilsta - A different whitelisting mechanism
2//
3// Author:
4//      Nissar Chababy, @funilrys, contactTATAfunilrysTODTODcom
5//
6// License:
7//      Copyright (c) 2022, 2023, 2024, 2025 Nissar Chababy
8//
9//      Licensed under the Apache License, Version 2.0 (the "License");
10//      you may not use this file except in compliance with the License.
11//      You may obtain a copy of the License at
12//
13//          http://www.apache.org/licenses/LICENSE-2.0
14//
15//      Unless required by applicable law or agreed to in writing, software
16//      distributed under the License is distributed on an "AS IS" BASIS,
17//      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18//      See the License for the specific language governing permissions and
19//      limitations under the License.
20
21mod data;
22mod utils;
23
24use crate::data::iana;
25use crate::data::psl;
26use crate::utils::{download_file, extract_netloc};
27use fancy_regex::Regex;
28use std::collections::hash_map::Entry;
29use std::collections::HashMap;
30use std::collections::HashSet;
31use std::fs;
32use std::fs::File;
33use std::io::{BufRead, BufReader};
34
35#[derive(Debug, Clone)]
36struct RulerSettings {
37    handle_complement: bool,
38    extensions: Vec<String>,
39}
40
41#[derive(Debug, Clone)]
42struct RulerTmps {
43    downloaded_files: Vec<String>,
44}
45
46#[derive(Debug, Clone)]
47pub struct Ruler {
48    strict: HashMap<String, HashSet<String>>,
49    ends: HashMap<String, HashSet<String>>,
50    present: HashMap<String, HashSet<String>>,
51    regex: String,
52    compiled_regex: Regex,
53    settings: RulerSettings,
54    tmps: RulerTmps,
55}
56
57impl Ruler {
58    /// Creates a new empty Ruler object.
59    ///
60    /// # Arguments
61    ///
62    /// * `handle_complement` - Whether we should follow and cleanup complements.
63    /// A complement is `www.example.org` if `example.org` has been given - and vice-versa.
64    ///
65    /// # Returns
66    ///
67    /// A new Ruler object.
68    ///
69    /// # Example
70    ///
71    /// ### Parsing a vector
72    ///
73    /// ```rust
74    /// use tivilsta::Ruler;
75    ///
76    /// let mut ruler = Ruler::new(false);
77    ///
78    /// let my_subjects: Vec<String> = vec![
79    ///     String::from("example.com"),
80    ///     String::from("example.org"),
81    ///     String::from("api.example.org"),
82    ///     String::from("test.example.com"),
83    /// ];
84    ///
85    /// let whitelisting_rules: Vec<String> = vec![
86    ///     String::from("api.example.org"),
87    ///     String::from("ALL .com"),
88    /// ];
89    ///
90    /// // Check that no rule is loaded.
91    /// assert_eq!(ruler.is_whitelisted(&String::from("example.com")), false);
92    /// assert_eq!(ruler.is_whitelisted(&String::from("example.org")), false);
93    /// assert_eq!(ruler.is_whitelisted(&String::from("api.example.com")), false);
94    /// assert_eq!(ruler.is_whitelisted(&String::from("test.example.com")), false);
95    ///
96    /// // Let's parse our rules.
97    /// ruler.parse_vec(&whitelisting_rules);
98    ///
99    /// assert_eq!(ruler.is_whitelisted(&String::from("example.com")), true);
100    /// assert_eq!(ruler.is_whitelisted(&String::from("example.org")), false);
101    /// assert_eq!(ruler.is_whitelisted(&String::from("api.example.com")), true);
102    /// assert_eq!(ruler.is_whitelisted(&String::from("test.example.com")), true);
103    ///
104    /// // Let's unparse our rules.
105    /// ruler.unparse_vec(&whitelisting_rules);
106    ///
107    /// assert_eq!(ruler.is_whitelisted(&String::from("example.com")), false);
108    /// assert_eq!(ruler.is_whitelisted(&String::from("example.org")), false);
109    /// assert_eq!(ruler.is_whitelisted(&String::from("api.example.com")), false);
110    /// assert_eq!(ruler.is_whitelisted(&String::from("test.example.com")), false);
111    /// ```
112    pub fn new(handle_complement: bool) -> Ruler {
113        Ruler {
114            strict: HashMap::new(),
115            ends: HashMap::new(),
116            present: HashMap::new(),
117            regex: String::from(""),
118            compiled_regex: Regex::new("").unwrap(),
119            settings: RulerSettings {
120                handle_complement,
121                extensions: vec![],
122            },
123            tmps: RulerTmps {
124                downloaded_files: vec![],
125            },
126        }
127    }
128
129    fn reduce(&self, element: &String) -> String {
130        if let Some(stripped) = element.strip_prefix("www.") {
131            stripped.to_string()
132        } else {
133            element.to_string()
134        }
135    }
136
137    fn extensions() -> Vec<String> {
138        let mut extensions: Vec<String> = Vec::new();
139
140        let mut iana_extensions = iana::extensions().unwrap();
141        let mut psl_suffixes = psl::suffixes().unwrap();
142
143        extensions.append(&mut iana_extensions);
144        extensions.append(&mut psl_suffixes);
145
146        extensions
147    }
148
149    fn search_keys(&mut self, record: &str) -> (String, String) {
150        let common_search_key = record.chars().take(4).collect::<String>();
151        let ends_search_key = record
152            .chars()
153            .rev()
154            .take(3)
155            .collect::<Vec<_>>()
156            .into_iter()
157            .rev()
158            .collect::<String>();
159
160        (common_search_key, ends_search_key)
161    }
162
163    fn push_strict(&mut self, record: &String) {
164        let (search_key, _) = self.search_keys(&self.reduce(record));
165
166        match self.strict.entry(search_key) {
167            Entry::Occupied(mut entry) => {
168                entry.get_mut().insert(record.to_string());
169            }
170            Entry::Vacant(entry) => {
171                let mut dataset = HashSet::new();
172
173                dataset.insert(record.to_string());
174                entry.insert(dataset);
175            }
176        }
177    }
178
179    fn pull_strict(&mut self, record: &String) {
180        let (search_key, _) = self.search_keys(&self.reduce(record));
181
182        match self.strict.entry(search_key) {
183            Entry::Occupied(mut entry) => {
184                entry.get_mut().remove(record);
185            }
186            Entry::Vacant(entry) => {
187                let _ = entry;
188            }
189        }
190    }
191
192    fn push_present(&mut self, record: &String) {
193        let (search_key, _) = self.search_keys(&self.reduce(record));
194
195        match self.present.entry(search_key) {
196            Entry::Occupied(mut entry) => {
197                entry.get_mut().insert(record.to_string());
198            }
199            Entry::Vacant(entry) => {
200                let mut dataset = HashSet::new();
201
202                dataset.insert(record.to_string());
203                entry.insert(dataset);
204            }
205        }
206    }
207
208    fn pull_present(&mut self, record: &String) {
209        let (search_key, _) = self.search_keys(&self.reduce(record));
210
211        match self.present.entry(search_key) {
212            Entry::Occupied(mut entry) => {
213                entry.get_mut().remove(record);
214            }
215            Entry::Vacant(entry) => {
216                let _ = entry;
217            }
218        }
219    }
220
221    fn push_ends(&mut self, record: &String) {
222        let (_, search_key) = self.search_keys(&self.reduce(record));
223
224        match self.ends.entry(search_key) {
225            Entry::Occupied(mut entry) => {
226                entry.get_mut().insert(record.to_string());
227            }
228            Entry::Vacant(entry) => {
229                let mut dataset = HashSet::new();
230
231                dataset.insert(record.to_string());
232                entry.insert(dataset);
233            }
234        }
235    }
236
237    fn pull_ends(&mut self, record: &String) {
238        let (_, search_key) = self.search_keys(&self.reduce(record));
239
240        match self.ends.entry(search_key) {
241            Entry::Occupied(mut entry) => {
242                entry.get_mut().remove(record);
243            }
244            Entry::Vacant(entry) => {
245                let _ = entry;
246            }
247        }
248    }
249
250    fn push_regex(&mut self, record: &String) {
251        if self.regex.is_empty() {
252            self.regex.push_str(&record.to_string());
253        } else {
254            self.regex.push_str(&format!("|{}", record));
255        }
256
257        self.compiled_regex = Regex::new(&self.regex[..]).unwrap();
258    }
259
260    fn pull_regex(&mut self, record: &String) {
261        if self.regex.starts_with(record) && self.regex.ends_with(record) {
262            self.regex = String::from("");
263        } else if self.regex.starts_with(record) {
264            self.regex = self.regex.replace(&format!("{}|", record), "");
265        } else {
266            self.regex = self.regex.replace(&format!("|{}", record), "");
267        }
268
269        self.compiled_regex = Regex::new(&self.regex[..]).unwrap();
270    }
271
272    fn parse_all(&mut self, line: &str) -> bool {
273        let record: String;
274
275        if line.starts_with("ALL ") {
276            record = line.replacen("ALL ", "", 1).trim().to_string()
277        } else if line.starts_with("all ") {
278            record = line.replacen("all ", "", 1).trim().to_string()
279        } else {
280            return false;
281        }
282
283        if let Some(stripped) = record.strip_prefix('.') {
284            if record.matches('.').count() > 1 {
285                if self.settings.handle_complement {
286                    self.push_strict(&format!("www.{}", stripped));
287                }
288                self.push_strict(&stripped.to_string());
289            }
290            self.push_ends(&record);
291        } else {
292            self.parse(&format!("ALL .{}", record));
293        }
294
295        true
296    }
297
298    fn unparse_all(&mut self, line: &str) -> bool {
299        let record: String;
300
301        if line.starts_with("ALL ") {
302            record = line.replacen("ALL ", "", 1).trim().to_string()
303        } else if line.starts_with("all ") {
304            record = line.replacen("all ", "", 1).trim().to_string()
305        } else {
306            return false;
307        }
308
309        if let Some(stripped) = record.strip_prefix('.') {
310            if record.matches('.').count() > 1 {
311                if self.settings.handle_complement {
312                    self.pull_strict(&format!("www.{}", stripped));
313                }
314                self.pull_strict(&stripped.to_string());
315            }
316            self.pull_ends(&record);
317        } else {
318            self.unparse(&format!("ALL .{}", record));
319        }
320
321        true
322    }
323
324    fn parse_root_zone_db(&mut self, line: &str) -> bool {
325        let mut record: String;
326
327        if line.starts_with("RZD ") {
328            record = line.replacen("RZD ", "", 1).trim().to_string()
329        } else if line.starts_with("rzd ") {
330            record = line.replacen("rzd ", "", 1).trim().to_string()
331        } else {
332            return false;
333        }
334
335        if self.settings.handle_complement && record.starts_with("www.") {
336            record = record.replacen("www.", "", 1).trim().to_string();
337        }
338
339        if self.settings.extensions.is_empty() {
340            self.settings.extensions = Ruler::extensions()
341        }
342
343        for extension in &self.settings.extensions.clone() {
344            self.push_present(&format!("{}.{}", record, extension));
345
346            if self.settings.handle_complement {
347                self.push_present(&format!("www.{}.{}", record, extension));
348            }
349        }
350
351        true
352    }
353
354    fn unparse_root_zone_db(&mut self, line: &str) -> bool {
355        let mut record: String;
356
357        if line.starts_with("RZD ") {
358            record = line.replacen("RZD ", "", 1).trim().to_string()
359        } else if line.starts_with("rzd ") {
360            record = line.replacen("rzd ", "", 1).trim().to_string()
361        } else {
362            return false;
363        }
364
365        if self.settings.handle_complement && record.starts_with("www.") {
366            record = record.replacen("www.", "", 1).trim().to_string();
367        }
368
369        if self.settings.extensions.is_empty() {
370            self.settings.extensions = Ruler::extensions()
371        }
372
373        for extension in &self.settings.extensions.clone() {
374            self.pull_present(&format!("{}.{}", record, extension));
375
376            if self.settings.handle_complement {
377                self.pull_present(&format!("www.{}.{}", record, extension));
378            }
379        }
380
381        true
382    }
383
384    fn parse_regex(&mut self, line: &str) -> bool {
385        let record: String;
386
387        if line.starts_with("REG ") {
388            record = line.replacen("REG ", "", 1).trim().to_string()
389        } else if line.starts_with("reg ") {
390            record = line.replacen("reg ", "", 1).trim().to_string()
391        } else {
392            return false;
393        }
394
395        self.push_regex(&record);
396
397        true
398    }
399
400    fn unparse_regex(&mut self, line: &str) -> bool {
401        let record: String;
402
403        if line.starts_with("REG ") {
404            record = line.replacen("REG ", "", 1).trim().to_string()
405        } else if line.starts_with("reg ") {
406            record = line.replacen("reg ", "", 1).trim().to_string()
407        } else {
408            return false;
409        }
410
411        self.pull_regex(&record);
412
413        true
414    }
415
416    fn parse_plain(&mut self, line: &String) -> bool {
417        let record: String = if self.settings.handle_complement && line.starts_with("www.") {
418            line.replacen("www.", "", 1).trim().to_string()
419        } else {
420            line.to_string()
421        };
422
423        self.push_strict(&record);
424
425        if self.settings.handle_complement {
426            self.push_strict(&format!("www.{}", record));
427        }
428
429        true
430    }
431
432    fn unparse_plain(&mut self, line: &String) -> bool {
433        let record: &String = &self.reduce(line);
434        self.pull_strict(record);
435
436        if self.settings.handle_complement {
437            self.pull_strict(&format!("www.{}", record));
438        }
439
440        true
441    }
442
443    /// Parses the given String into the ruler.
444    ///
445    /// # Arguments
446    ///
447    /// * `line` - The line to parse.
448    ///
449    /// # Returns
450    ///
451    /// Nothing.
452    pub fn parse(&mut self, line: &String) {
453        if line.is_empty() || line.starts_with('#') {
454            return;
455        }
456
457        let idnazed_line = self.idnaze_line(line);
458
459        let _ = self.parse_all(&idnazed_line)
460            || self.parse_regex(&idnazed_line)
461            || self.parse_root_zone_db(&idnazed_line)
462            || self.parse_plain(&idnazed_line);
463    }
464
465    /// Parses the given Vector of Strings into the ruler.
466    ///
467    /// # Arguments
468    ///
469    /// * `lines` - The lines to parse.
470    ///
471    /// # Returns
472    ///
473    /// Nothing.
474    pub fn parse_vec(&mut self, lines: &[String]) {
475        for line in lines {
476            self.parse(line);
477        }
478    }
479
480    /// Parses the content of the given file into the ruler.
481    ///
482    /// # Arguments
483    ///
484    /// * `file` - The file to parse.
485    ///
486    /// # Returns
487    ///
488    /// Nothing.
489    pub fn parse_file(&mut self, path: &str) {
490        let file = File::open(path).unwrap();
491        let reader = BufReader::new(file);
492
493        for line in reader.lines() {
494            self.parse(&line.unwrap());
495        }
496    }
497
498    /// Parses the content of the given URL (after downloading it) into the ruler.
499    ///
500    /// # Arguments
501    ///
502    /// * `url` - The URL to download and parse.
503    ///
504    /// # Returns
505    ///
506    /// Nothing.
507    pub fn parse_link(&mut self, url: &str) {
508        let (real_path, downloaded) = download_file(&url.to_string());
509
510        if downloaded {
511            self.tmps.downloaded_files.push(real_path.clone());
512        }
513
514        self.parse_file(real_path.as_str());
515    }
516
517    /// Unparses the given String into the ruler.
518    ///
519    /// # Arguments
520    ///
521    /// * `line` - The line to parse.
522    ///
523    /// # Returns
524    ///
525    /// Nothing.
526    pub fn unparse(&mut self, line: &String) {
527        if line.is_empty() || line.starts_with('#') {
528            return;
529        }
530
531        let _ = self.unparse_all(line)
532            || self.unparse_regex(line)
533            || self.unparse_root_zone_db(line)
534            || self.unparse_plain(line);
535    }
536
537    /// Unparses the given Vector of Strings into the ruler.
538    ///
539    /// # Arguments
540    ///
541    /// * `lines` - The lines to parse.
542    ///
543    /// # Returns
544    ///
545    /// Nothing.
546    pub fn unparse_vec(&mut self, lines: &[String]) {
547        for line in lines {
548            self.unparse(line);
549        }
550    }
551
552    /// Unparses the content of the given file into the ruler.
553    ///
554    /// # Arguments
555    ///
556    /// * `file` - The file to parse.
557    ///
558    /// # Returns
559    ///
560    /// Nothing.
561    pub fn unparse_file(&mut self, path: &str) {
562        let file = File::open(path).unwrap();
563        let reader = BufReader::new(file);
564
565        for line in reader.lines() {
566            self.unparse(&line.unwrap());
567        }
568    }
569
570    /// Unparses the content of the given URL (after downloading it) into the ruler.
571    ///
572    /// # Arguments
573    ///
574    /// * `url` - The URL to download and parse.
575    ///
576    /// # Returns
577    ///
578    /// Nothing.
579    pub fn unparse_link(&mut self, url: &str) {
580        let (real_path, downloaded) = download_file(&url.to_string());
581
582        if downloaded {
583            self.tmps.downloaded_files.push(real_path.clone());
584        }
585
586        self.unparse_file(real_path.as_str());
587    }
588
589    /// IDNAze the given `subject`.
590    ///
591    /// # Arguments
592    ///
593    /// * `subject` - The subject to IDNAze.
594    ///
595    /// # Returns
596    ///
597    /// The IDNAzed subject.
598    ///
599    /// # Example
600    ///
601    /// ```rust
602    /// use tivilsta::Ruler;
603    ///
604    /// let mut ruler = Ruler::new(false);
605    ///
606    /// let subject = String::from("www.äxample.org");
607    ///
608    /// assert_eq!(ruler.idnaze_subject(&subject), "www.xn--xample-9ta.org");
609    /// ```
610    pub fn idnaze_subject(&mut self, subject: &String) -> String {
611        match idna::domain_to_ascii(subject.as_str()) {
612            Ok(result) => result,
613            Err(_) => subject.to_string(),
614        }
615    }
616
617    /// IDNAze the given `line`.
618    ///
619    /// # Arguments
620    ///
621    /// * `line` - The line to IDNAze.
622    ///
623    /// # Returns
624    ///
625    /// The IDNAzed line.
626    ///
627    /// # Example
628    ///
629    /// This example shows how to IDNAze a line, where the line contains a comment.
630    /// In such cases, the comment will be kept as is.
631    ///
632    /// ```rust
633    /// use tivilsta::Ruler;
634    ///
635    /// let mut ruler = Ruler::new(false);
636    ///
637    /// let line = String::from("www.äxample.org # äxample.org");
638    ///
639    /// assert_eq!(ruler.idnaze_line(&line), "www.xn--xample-9ta.org # äxample.org");
640    /// ```
641    pub fn idnaze_line(&mut self, line: &String) -> String {
642        let tab = "\t";
643        let space = " ";
644
645        let separator;
646
647        let regex_ignore = Regex::new(r"localhost$|localdomain$|local$|broadcasthost$|0\.0\.0\.0$|allhosts$|allnodes$|allrouters$|localnet$|loopback$|mcastprefix$").unwrap();
648
649        if line.is_empty() || line.starts_with('#') || regex_ignore.is_match(&line[..]).unwrap() {
650            return line.clone();
651        }
652
653        if line.contains(tab) {
654            separator = tab
655        } else if line.contains(space) {
656            separator = space
657        } else {
658            separator = ""
659        }
660
661        if !separator.is_empty() {
662            let mut idnazed_data: Vec<String> = Vec::new();
663
664            let subjects: &str;
665            let mut comment = "";
666
667            if line.contains('#') {
668                (subjects, comment) = line.split_once('#').unwrap();
669            } else {
670                subjects = line;
671            }
672
673            let mut splitted_subject: Vec<&str> = subjects.split(separator).collect();
674
675            for data in splitted_subject.iter_mut() {
676                if data.is_empty() || regex_ignore.is_match(data).unwrap() {
677                    idnazed_data.push(data.to_string());
678                    continue;
679                }
680
681                let idnazed = if data.contains('#') {
682                    let (element, comment) = data.split_once('#').unwrap();
683                    let idnazed_line =
684                        format!("{} #{}", self.idnaze_subject(&element.to_string()), comment);
685
686                    idnazed_line
687                } else {
688                    self.idnaze_subject(&data.to_string())
689                };
690
691                idnazed_data.push(idnazed);
692            }
693
694            if !comment.is_empty() {
695                return idnazed_data.join(separator) + "#" + comment;
696            }
697
698            return idnazed_data.join(separator);
699        }
700
701        self.idnaze_subject(line)
702    }
703
704    /// Checks the given `line` against the rules.
705    ///
706    /// # Arguments
707    ///
708    /// * `line` - The line to check. **WARNING:** We assume 1 rule per line.
709    ///
710    ///   **Note:** If a URL (e.g `https://example.org/`) is given, the sub-domain
711    ///   will be used to determine if the line has been whitelisted.
712    ///
713    /// # Returns
714    ///
715    /// A `bool` indicating whether the line matches the rules.
716    /// Any `true` value should be considered positive.
717    /// Meaning that the line matches one of the rule.
718    pub fn is_whitelisted(&mut self, line: &String) -> bool {
719        if line.is_empty() || line.starts_with('#') {
720            return false;
721        }
722
723        let mut flines: Vec<String> = vec![extract_netloc(line)];
724
725        if line.starts_with("http://") || line.starts_with("https://") {
726            flines.push(line.to_string());
727        }
728
729        for fline in flines.iter() {
730            let (common_skey, ends_skey) = self.search_keys(&self.reduce(&fline));
731
732            let mut matching_state;
733
734            match self.strict.entry(common_skey.to_string()) {
735                Entry::Occupied(entry) => matching_state = entry.get().contains(fline),
736                Entry::Vacant(_) => matching_state = false,
737            }
738
739            if matching_state {
740                return true;
741            }
742
743            match self.present.entry(common_skey) {
744                Entry::Occupied(entry) => matching_state = entry.get().contains(fline),
745                Entry::Vacant(_) => matching_state = false,
746            }
747
748            if matching_state {
749                return true;
750            }
751
752            match self.ends.entry(ends_skey) {
753                Entry::Occupied(entry) => {
754                    let mut matching = entry.get().iter().map(|x| fline.ends_with(x)).peekable();
755                    matching_state = *matching.peek().unwrap_or(&false);
756                }
757                Entry::Vacant(_) => matching_state = false,
758            }
759
760            if matching_state {
761                return true;
762            }
763
764            if !self.regex.is_empty() {
765                return self.compiled_regex.is_match(&fline[..]).unwrap();
766            }
767        }
768        false
769    }
770}
771
772impl Drop for Ruler {
773    fn drop(&mut self) {
774        for file in &self.tmps.downloaded_files {
775            let _ = fs::remove_file(file);
776        }
777    }
778}
779
780#[cfg(test)]
781mod tests {
782    use super::*;
783
784    #[test]
785    fn test_new_ruler_gen_complement_true() {
786        let ruler = Ruler::new(true);
787
788        assert_eq!(ruler.settings.handle_complement, true)
789    }
790
791    #[test]
792    fn test_new_ruler_gen_complement_false() {
793        let ruler = Ruler::new(false);
794
795        assert_eq!(ruler.settings.handle_complement, false)
796    }
797
798    #[test]
799    fn test_reduce() {
800        let ruler = Ruler::new(false);
801
802        assert_eq!(
803            ruler.reduce(&"www.example.org".to_string()),
804            "example.org".to_string()
805        )
806    }
807
808    #[test]
809    fn test_reduce_no_www() {
810        let ruler = Ruler::new(false);
811
812        assert_eq!(
813            ruler.reduce(&"example.org".to_string()),
814            "example.org".to_string()
815        )
816    }
817
818    #[test]
819    fn test_reduce_multiple_www() {
820        let ruler = Ruler::new(false);
821
822        assert_eq!(
823            ruler.reduce(&"www.www.example.org".to_string()),
824            "www.example.org".to_string()
825        )
826    }
827
828    #[test]
829    fn test_search_keys() {
830        let mut ruler = Ruler::new(false);
831
832        assert_eq!(
833            ruler.search_keys(&"example.org".to_string()),
834            ("exam".to_string(), "org".to_string())
835        )
836    }
837
838    #[test]
839    fn test_search_keys_long_extension() {
840        let mut ruler = Ruler::new(false);
841
842        assert_eq!(
843            ruler.search_keys(&"example.example".to_string()),
844            ("exam".to_string(), "ple".to_string())
845        )
846    }
847
848    #[test]
849    fn test_idnaze_subject() {
850        let mut ruler = Ruler::new(false);
851
852        assert_eq!(
853            ruler.idnaze_subject(&"www.äxample.org".to_string()),
854            "www.xn--xample-9ta.org".to_string()
855        );
856
857        assert_eq!(
858            ruler.idnaze_subject(&"www.example.org".to_string()),
859            "www.example.org".to_string()
860        );
861    }
862
863    #[test]
864    fn test_idnaze_line() {
865        let mut ruler = Ruler::new(false);
866
867        assert_eq!(
868            ruler.idnaze_line(&"www.äxample.org".to_string()),
869            "www.xn--xample-9ta.org".to_string()
870        );
871
872        assert_eq!(
873            ruler.idnaze_line(&"www.example.org".to_string()),
874            "www.example.org".to_string()
875        );
876
877        assert_eq!(
878            ruler.idnaze_line(&"www.example.org # example.org".to_string()),
879            "www.example.org # example.org".to_string()
880        );
881
882        assert_eq!(
883            ruler.idnaze_line(&"www.example.org # äxample.org".to_string()),
884            "www.example.org # äxample.org".to_string()
885        );
886
887        assert_eq!(
888            ruler.idnaze_line(&"www.example.org # äxample.org # example.org".to_string()),
889            "www.example.org # äxample.org # example.org".to_string()
890        );
891
892        assert_eq!(
893            ruler.idnaze_line(&"www.example.org       # äxample.org # example.org".to_string()),
894            "www.example.org       # äxample.org # example.org".to_string()
895        );
896
897        assert_eq!(
898            ruler.idnaze_line(&"www.äxample.org # äxample.org # example.org #".to_string()),
899            "www.xn--xample-9ta.org # äxample.org # example.org #".to_string()
900        );
901    }
902
903    #[test]
904    fn test_push_strict() {
905        let mut ruler = Ruler::new(false);
906
907        // Ensure that it's really empty :)
908        assert_eq!(ruler.strict.get_key_value("exam"), None);
909
910        ruler.push_strict(&"www.example.org".to_string());
911
912        let mut expected = HashSet::new();
913        expected.insert("www.example.org".to_string());
914
915        assert_eq!(
916            ruler.strict.get_key_value("exam"),
917            Some((&"exam".to_string(), &expected))
918        );
919
920        // Let's add another one.
921
922        ruler.push_strict(&"example.net".to_string());
923        expected.insert("example.net".to_string());
924
925        assert_eq!(
926            ruler.strict.get_key_value("exam"),
927            Some((&"exam".to_string(), &expected))
928        );
929    }
930
931    #[test]
932    fn test_pull_strict() {
933        let mut ruler = Ruler::new(false);
934
935        // Ensure that it's really empty :)
936        assert_eq!(ruler.strict.get_key_value("exam"), None);
937
938        // Add some data into it :)
939        ruler.push_strict(&"www.example.org".to_string());
940        ruler.push_strict(&"example.net".to_string());
941
942        ruler.pull_strict(&"www.example.org".to_string());
943
944        let mut expected = HashSet::new();
945        expected.insert("example.net".to_string());
946
947        assert_eq!(
948            ruler.strict.get_key_value("exam"),
949            Some((&"exam".to_string(), &expected))
950        );
951
952        // Let's remove another one.
953        ruler.pull_strict(&"example.net".to_string());
954        expected.remove("example.net");
955
956        assert_eq!(
957            ruler.strict.get_key_value("exam"),
958            Some((&"exam".to_string(), &expected))
959        );
960    }
961
962    #[test]
963    fn test_push_present() {
964        let mut ruler = Ruler::new(false);
965
966        // Ensure that it's really empty :)
967        assert_eq!(ruler.present.get_key_value("exam"), None);
968
969        ruler.push_present(&"www.example.net".to_string());
970
971        let mut expected = HashSet::new();
972        expected.insert("www.example.net".to_string());
973
974        assert_eq!(
975            ruler.present.get_key_value("exam"),
976            Some((&"exam".to_string(), &expected))
977        );
978
979        // Let's add another one.
980
981        ruler.push_present(&"example.com".to_string());
982        expected.insert("example.com".to_string());
983
984        assert_eq!(
985            ruler.present.get_key_value("exam"),
986            Some((&"exam".to_string(), &expected))
987        );
988    }
989
990    #[test]
991    fn test_pull_present() {
992        let mut ruler = Ruler::new(false);
993
994        // Ensure that it's really empty :)
995        assert_eq!(ruler.present.get_key_value("exam"), None);
996
997        // Add some data into it :)
998        ruler.push_present(&"www.example.net".to_string());
999        ruler.push_present(&"example.org".to_string());
1000
1001        ruler.pull_present(&"www.example.net".to_string());
1002
1003        let mut expected = HashSet::new();
1004        expected.insert("example.org".to_string());
1005
1006        assert_eq!(
1007            ruler.present.get_key_value("exam"),
1008            Some((&"exam".to_string(), &expected))
1009        );
1010
1011        // Let's remove another one.
1012        ruler.pull_present(&"example.org".to_string());
1013        expected.remove("example.org");
1014
1015        assert_eq!(
1016            ruler.present.get_key_value("exam"),
1017            Some((&"exam".to_string(), &expected))
1018        );
1019    }
1020
1021    #[test]
1022    fn test_push_ends() {
1023        let mut ruler = Ruler::new(false);
1024
1025        // Ensure that it's really empty :)
1026        assert_eq!(ruler.ends.get_key_value("ple"), None);
1027
1028        ruler.push_ends(&"www.example.example".to_string());
1029
1030        let mut expected = HashSet::new();
1031        expected.insert("www.example.example".to_string());
1032
1033        assert_eq!(
1034            ruler.ends.get_key_value("ple"),
1035            Some((&"ple".to_string(), &expected))
1036        );
1037
1038        // Let's add another one.
1039
1040        ruler.push_ends(&"example.com".to_string());
1041
1042        let mut expected = HashSet::new();
1043        expected.insert("example.com".to_string());
1044
1045        assert_eq!(
1046            ruler.ends.get_key_value("com"),
1047            Some((&"com".to_string(), &expected))
1048        );
1049
1050        // Let's add another one.
1051
1052        ruler.push_ends(&"example.co".to_string());
1053
1054        let mut expected = HashSet::new();
1055        expected.insert("example.co".to_string());
1056
1057        assert_eq!(
1058            ruler.ends.get_key_value(".co"),
1059            Some((&".co".to_string(), &expected))
1060        );
1061
1062        assert_eq!(ruler.ends.contains_key("com"), true);
1063        assert_eq!(ruler.ends.contains_key("ple"), true);
1064        assert_eq!(ruler.ends.contains_key(".co"), true);
1065    }
1066
1067    #[test]
1068    fn test_pull_ends() {
1069        let mut ruler = Ruler::new(false);
1070
1071        // Ensure that it's really empty :)
1072        assert_eq!(ruler.ends.get_key_value("ple"), None);
1073
1074        // Add some data into it :)
1075        ruler.push_ends(&"www.example.example".to_string());
1076        ruler.push_ends(&"example.com".to_string());
1077        ruler.push_ends(&"example.co".to_string());
1078
1079        assert_eq!(ruler.ends.contains_key("com"), true);
1080        assert_eq!(ruler.ends.contains_key("ple"), true);
1081        assert_eq!(ruler.ends.contains_key(".co"), true);
1082
1083        ruler.pull_ends(&"www.example.example".to_string());
1084
1085        let expected = HashSet::new();
1086
1087        assert_eq!(
1088            ruler.ends.get_key_value("ple"),
1089            Some((&"ple".to_string(), &expected))
1090        );
1091
1092        let mut expected = HashSet::new();
1093        expected.insert("example.com".to_string());
1094
1095        assert_eq!(
1096            ruler.ends.get_key_value("com"),
1097            Some((&"com".to_string(), &expected))
1098        );
1099
1100        let mut expected = HashSet::new();
1101        expected.insert("example.co".to_string());
1102
1103        assert_eq!(
1104            ruler.ends.get_key_value(".co"),
1105            Some((&".co".to_string(), &expected))
1106        );
1107
1108        // Let's remove another one.
1109        ruler.pull_ends(&"example.com".to_string());
1110
1111        let expected = HashSet::new();
1112
1113        assert_eq!(
1114            ruler.ends.get_key_value("com"),
1115            Some((&"com".to_string(), &expected))
1116        );
1117
1118        assert_eq!(ruler.ends.contains_key("com"), true);
1119        assert_eq!(ruler.ends.contains_key("ple"), true);
1120        assert_eq!(ruler.ends.contains_key(".co"), true);
1121    }
1122
1123    #[test]
1124    fn test_push_regex() {
1125        let mut ruler = Ruler::new(false);
1126
1127        // Ensure that it's really empty :)
1128        assert_eq!(ruler.regex, "");
1129        assert_eq!(ruler.compiled_regex.as_str(), "");
1130
1131        ruler.push_regex(&"^(www.)?example.com$".to_string());
1132
1133        let expected = "^(www.)?example.com$".to_string();
1134
1135        assert_eq!(ruler.regex, expected);
1136        assert_eq!(ruler.compiled_regex.as_str(), &expected[..]);
1137
1138        // Let's add another one.
1139        ruler.push_regex(&"^(api.)?example.org$".to_string());
1140
1141        let expected = "^(www.)?example.com$|^(api.)?example.org$".to_string();
1142
1143        assert_eq!(ruler.regex, expected);
1144        assert_eq!(ruler.compiled_regex.as_str(), &expected[..]);
1145    }
1146
1147    #[test]
1148    fn test_pull_regex() {
1149        let mut ruler = Ruler::new(false);
1150
1151        // Ensure that it's really empty :)
1152        assert_eq!(ruler.regex, "");
1153        assert_eq!(ruler.compiled_regex.as_str(), "");
1154
1155        // Add some data into it :)
1156        ruler.push_regex(&"^(www.)?example.com$".to_string());
1157        ruler.push_regex(&"^(api.)?example.org$".to_string());
1158
1159        ruler.pull_regex(&"^(www.)?example.com$".to_string());
1160
1161        let expected = "^(api.)?example.org$".to_string();
1162
1163        assert_eq!(ruler.regex, expected);
1164        assert_eq!(ruler.compiled_regex.as_str(), &expected[..]);
1165
1166        // Let's remove another one.
1167        ruler.pull_regex(&"^(api.)?example.org$".to_string());
1168
1169        let expected = "".to_string();
1170
1171        assert_eq!(ruler.regex, expected);
1172        assert_eq!(ruler.compiled_regex.as_str(), &expected[..]);
1173    }
1174
1175    #[test]
1176    fn test_parse_all() {
1177        let mut ruler = Ruler::new(false);
1178
1179        let given = &"example.org".to_string();
1180        let mut expected_res = false;
1181
1182        let mut expected_ends: HashMap<String, HashSet<String>> = HashMap::new();
1183        let mut expected_strict: HashMap<String, HashSet<String>> = HashMap::new();
1184        let expected_present: HashMap<String, HashSet<String>> = HashMap::new();
1185        let expected_regex = "".to_string();
1186
1187        assert_eq!(ruler.parse_all(given), expected_res);
1188        assert_eq!(ruler.ends, expected_ends);
1189        assert_eq!(ruler.strict, expected_strict);
1190        assert_eq!(ruler.present, expected_present);
1191        assert_eq!(ruler.regex, expected_regex);
1192
1193        // Let's add a new one.
1194        let given = &"ALL example.org".to_string();
1195        expected_res = true;
1196
1197        let mut ends_set = HashSet::new();
1198        ends_set.insert(".example.org".to_string());
1199        expected_ends.insert("org".to_string(), ends_set);
1200
1201        let mut strict_set = HashSet::new();
1202        strict_set.insert("example.org".to_string());
1203        expected_strict.insert("exam".to_string(), strict_set);
1204
1205        assert_eq!(ruler.parse_all(given), expected_res);
1206        assert_eq!(ruler.ends, expected_ends);
1207        assert_eq!(ruler.strict, expected_strict);
1208        assert_eq!(ruler.present, expected_present);
1209        assert_eq!(ruler.regex, expected_regex);
1210
1211        // Let's add another one but the marker is in lowercase.
1212        let given = &"all .example.net".to_string();
1213        expected_res = true;
1214
1215        let mut new_set = HashSet::new();
1216        new_set.insert(".example.net".to_string());
1217        expected_ends.insert("net".to_string(), new_set);
1218
1219        let mut new_set = HashSet::new();
1220        new_set.insert("example.org".to_string());
1221        new_set.insert("example.net".to_string());
1222        expected_strict.insert("exam".to_string(), new_set);
1223
1224        assert_eq!(ruler.parse_all(given), expected_res);
1225        assert_eq!(ruler.ends, expected_ends);
1226        assert_eq!(ruler.strict, expected_strict);
1227        assert_eq!(ruler.present, expected_present);
1228        assert_eq!(ruler.regex, expected_regex);
1229
1230        // Let's add another one but this time with the complement generation.
1231        ruler.settings.handle_complement = true;
1232
1233        let given = &"ALL .example.de".to_string();
1234        expected_res = true;
1235
1236        let mut new_set = HashSet::new();
1237        new_set.insert(".example.de".to_string());
1238        expected_ends.insert(".de".to_string(), new_set);
1239
1240        let mut new_set = HashSet::new();
1241        new_set.insert("example.org".to_string());
1242        new_set.insert("example.net".to_string());
1243        new_set.insert("example.de".to_string());
1244        new_set.insert("www.example.de".to_string());
1245
1246        expected_strict.insert("exam".to_string(), new_set);
1247
1248        assert_eq!(ruler.parse_all(given), expected_res);
1249        assert_eq!(ruler.ends, expected_ends);
1250        assert_eq!(ruler.strict, expected_strict);
1251        assert_eq!(ruler.present, expected_present);
1252        assert_eq!(ruler.regex, expected_regex);
1253    }
1254
1255    #[test]
1256    fn test_unparse_all() {
1257        let mut ruler = Ruler::new(false);
1258
1259        let given = &"ALL example.com".to_string();
1260        let mut expected_ends: HashMap<String, HashSet<String>> = HashMap::new();
1261        let mut expected_strict: HashMap<String, HashSet<String>> = HashMap::new();
1262        let expected_present: HashMap<String, HashSet<String>> = HashMap::new();
1263        let expected_regex = "".to_string();
1264
1265        // Fill ruler with some data
1266        ruler.parse_all(&"ALL .hello.com".to_string());
1267        ruler.parse_all(&"ALL .github.com".to_string());
1268        ruler.parse_all(&"ALL .example.com".to_string());
1269
1270        let mut ends_set = HashSet::new();
1271        ends_set.insert(".github.com".to_string());
1272        ends_set.insert(".hello.com".to_string());
1273        expected_ends.insert("com".to_string(), ends_set);
1274
1275        let mut strict_set1 = HashSet::new();
1276        strict_set1.insert("hello.com".to_string());
1277        expected_strict.insert("hell".to_string(), strict_set1);
1278
1279        let mut strict_set2 = HashSet::new();
1280        strict_set2.insert("github.com".to_string());
1281        expected_strict.insert("gith".to_string(), strict_set2);
1282        expected_strict.insert("exam".to_string(), HashSet::new());
1283
1284        assert_eq!(ruler.unparse_all(given), true);
1285        assert_eq!(ruler.ends, expected_ends);
1286        assert_eq!(ruler.strict, expected_strict);
1287        assert_eq!(ruler.present, expected_present);
1288        assert_eq!(ruler.regex, expected_regex);
1289
1290        // Let's remove another one but this time with the complement generation.
1291        ruler.settings.handle_complement = true;
1292
1293        ruler.parse_all(&"ALL .hello.com".to_string());
1294
1295        let mut strict_set1 = HashSet::new();
1296        strict_set1.insert("hello.com".to_string());
1297        strict_set1.insert("www.hello.com".to_string());
1298        expected_strict.insert("hell".to_string(), strict_set1);
1299
1300        let mut strict_set2 = HashSet::new();
1301        strict_set2.insert("github.com".to_string());
1302        expected_strict.insert("gith".to_string(), strict_set2);
1303        expected_strict.insert("exam".to_string(), HashSet::new());
1304
1305        let given = &"ALL .hello.world".to_string();
1306
1307        assert_eq!(ruler.strict, expected_strict);
1308
1309        assert_eq!(ruler.unparse_all(given), true);
1310        assert_eq!(ruler.ends, expected_ends);
1311        assert_eq!(ruler.strict, expected_strict);
1312        assert_eq!(ruler.present, expected_present);
1313        assert_eq!(ruler.regex, expected_regex);
1314    }
1315}
tivilsta/lib.rs

tivilsta/
lib.rs
