Add tests (#66)

* test: Tests for `create_inspector_url` * Add some unit tests for scanner.rs * Lint --------- Co-authored-by: Robin <[email protected]> Co-authored-by: Robin5605 <[email protected]> Co-authored-by: Bradley Reynolds <[email protected]>
vipyrsec · Feb 4, 2024 · 9998d25 · 9998d25
1 parent 55be966
commit 9998d25
Show file tree

Hide file tree

Showing 3 changed files with 343 additions and 2 deletions.
diff --git a/src/client/models.rs b/src/client/models.rs
@@ -6,7 +6,7 @@ use yara::{Compiler, Rules};
 
 use crate::error::DragonflyError;
 
-#[derive(Debug, Serialize)]
+#[derive(Debug, Serialize, PartialEq)]
 pub struct SubmitJobResultsSuccess {
     pub name: String,
     pub version: String,

diff --git a/src/scanner.rs b/src/scanner.rs
@@ -14,7 +14,7 @@ use crate::{
     utils::create_inspector_url,
 };
 
-#[derive(Debug, Hash, Eq, PartialEq)]
+#[derive(Debug, Hash, Eq, PartialEq, Clone)]
 pub struct RuleScore {
     pub name: String,
     pub score: i64,
@@ -273,3 +273,308 @@ fn scan_file(
 
     Ok(FileScanResult::new(path.to_path_buf(), rules))
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{collections::HashSet, path::PathBuf};
+    use yara::Compiler;
+
+    use super::{scan_file, DistributionScanResults, PackageScanResults};
+    use crate::scanner::{FileScanResult, RuleScore};
+
+    #[test]
+    fn test_file_score() {
+        let rules = vec![
+            RuleScore {
+                name: String::from("rule1"),
+                score: 5,
+            },
+            RuleScore {
+                name: String::from("rule2"),
+                score: 7,
+            },
+        ];
+
+        let file_scan_result = FileScanResult {
+            path: PathBuf::default(),
+            rules,
+        };
+        assert_eq!(file_scan_result.calculate_score(), 12);
+    }
+
+    #[test]
+    fn test_get_most_malicious_file() {
+        let file_scan_results = vec![
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule1"),
+                    score: 5,
+                }],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule2"),
+                    score: 7,
+                }],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule3"),
+                    score: 4,
+                }],
+            },
+        ];
+
+        let distribution_scan_results = DistributionScanResults {
+            file_scan_results,
+            inspector_url: reqwest::Url::parse("https://example.net").unwrap(),
+        };
+
+        assert_eq!(
+            distribution_scan_results
+                .get_most_malicious_file()
+                .unwrap()
+                .rules[0]
+                .name,
+            "rule2"
+        )
+    }
+
+    #[test]
+    fn test_get_matched_rules() {
+        let file_scan_results = vec![
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule1"),
+                        score: 5,
+                    },
+                    RuleScore {
+                        name: String::from("rule2"),
+                        score: 7,
+                    },
+                ],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule2"),
+                        score: 7,
+                    },
+                    RuleScore {
+                        name: String::from("rule3"),
+                        score: 9,
+                    },
+                ],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule3"),
+                        score: 9,
+                    },
+                    RuleScore {
+                        name: String::from("rule4"),
+                        score: 6,
+                    },
+                ],
+            },
+        ];
+
+        let distribution_scan_results = DistributionScanResults {
+            file_scan_results,
+            inspector_url: reqwest::Url::parse("https://example.net").unwrap(),
+        };
+
+        let matched_rules: HashSet<RuleScore> = distribution_scan_results
+            .get_matched_rules()
+            .into_iter()
+            .cloned()
+            .collect();
+
+        let expected_rules = HashSet::from([
+            RuleScore {
+                name: String::from("rule1"),
+                score: 5,
+            },
+            RuleScore {
+                name: String::from("rule2"),
+                score: 7,
+            },
+            RuleScore {
+                name: String::from("rule3"),
+                score: 9,
+            },
+            RuleScore {
+                name: String::from("rule4"),
+                score: 6,
+            },
+        ]);
+
+        assert_eq!(matched_rules, expected_rules)
+    }
+
+    #[test]
+    fn test_get_matched_rule_identifiers() {
+        let file_scan_results = vec![
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule1"),
+                        score: 5,
+                    },
+                    RuleScore {
+                        name: String::from("rule2"),
+                        score: 7,
+                    },
+                ],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule2"),
+                        score: 7,
+                    },
+                    RuleScore {
+                        name: String::from("rule3"),
+                        score: 9,
+                    },
+                ],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![
+                    RuleScore {
+                        name: String::from("rule3"),
+                        score: 9,
+                    },
+                    RuleScore {
+                        name: String::from("rule4"),
+                        score: 6,
+                    },
+                ],
+            },
+        ];
+
+        let distribution_scan_results = DistributionScanResults {
+            file_scan_results,
+            inspector_url: reqwest::Url::parse("https://example.net").unwrap(),
+        };
+
+        let matched_rule_identifiers = distribution_scan_results.get_matched_rule_identifiers();
+
+        let expected_rule_identifiers = vec!["rule1", "rule2", "rule3", "rule4"];
+
+        assert_eq!(
+            HashSet::<_>::from_iter(matched_rule_identifiers),
+            HashSet::<_>::from_iter(expected_rule_identifiers)
+        )
+    }
+
+    #[test]
+    fn test_build_package_scan_results_body() {
+        let file_scan_results1 = vec![
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule1"),
+                    score: 5,
+                }],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule2"),
+                    score: 7,
+                }],
+            },
+        ];
+        let distribution_scan_results1 = DistributionScanResults {
+            file_scan_results: file_scan_results1,
+            inspector_url: reqwest::Url::parse("https://example.net/distrib1.tar.gz").unwrap(),
+        };
+
+        let file_scan_results2 = vec![
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule3"),
+                    score: 2,
+                }],
+            },
+            FileScanResult {
+                path: PathBuf::default(),
+                rules: vec![RuleScore {
+                    name: String::from("rule4"),
+                    score: 9,
+                }],
+            },
+        ];
+        let distribution_scan_results2 = DistributionScanResults {
+            file_scan_results: file_scan_results2,
+            inspector_url: reqwest::Url::parse("https://example.net/distrib2.whl").unwrap(),
+        };
+
+        let package_scan_results = PackageScanResults {
+            name: String::from("remmy"),
+            version: String::from("4.20.69"),
+            distribution_scan_results: vec![distribution_scan_results1, distribution_scan_results2],
+            commit_hash: String::from("abc"),
+        };
+
+        let body = package_scan_results.build_body();
+
+        assert_eq!(
+            body.inspector_url,
+            Some(String::from("https://example.net/distrib1.tar.gz"))
+        );
+        assert_eq!(body.score, 12);
+        assert_eq!(
+            HashSet::from([
+                "rule1".into(),
+                "rule2".into(),
+                "rule3".into(),
+                "rule4".into()
+            ]),
+            HashSet::from_iter(body.rules_matched)
+        );
+    }
+
+    #[test]
+    fn test_scan_file() {
+        let rules = r#"
+            rule contains_rust {
+                meta:
+                    weight = 5
+                strings:
+                    $rust = "rust" nocase
+                condition:
+                    $rust
+            }
+        "#;
+
+        let compiler = Compiler::new().unwrap().add_rules_str(rules).unwrap();
+
+        let rules = compiler.compile_rules().unwrap();
+        let result =
+            scan_file(&mut "I love Rust!".as_bytes(), &PathBuf::default(), &rules).unwrap();
+
+        assert_eq!(result.path, PathBuf::default());
+        assert_eq!(
+            result.rules[0],
+            RuleScore {
+                name: "contains_rust".into(),
+                score: 5
+            }
+        );
+        assert_eq!(result.calculate_score(), 5);
+    }
+}
diff --git a/src/utils.rs b/src/utils.rs
@@ -16,3 +16,39 @@ pub fn create_inspector_url(name: &str, version: &str, download_url: &Url) -> Ur
 
     download_url
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! create_inspector_url_tests {
+        ($($name:ident: $value:expr,)*) => {
+            $(
+                #[test]
+                fn $name() {
+                    let ((n, version, download_url), exp) = $value;
+                    assert_eq!(exp, create_inspector_url(n, version, &download_url));
+                }
+            )*
+        }
+    }
+
+    create_inspector_url_tests! {
+        create_inspector_url_0: (
+            ("numpy", "1.24.3", Url::parse("https://files.pythonhosted.org/packages/f3/23/7cc851bae09cf4db90d42a701dfe525780883ada86bece45e3da7a07e76b/numpy-1.24.3-cp310-cp310-macosx_10_9_x86_64.whl/numpy/__init__.pyi").unwrap()),
+            Url::parse("https://inspector.pypi.io/project/numpy/1.24.3/packages/f3/23/7cc851bae09cf4db90d42a701dfe525780883ada86bece45e3da7a07e76b/numpy-1.24.3-cp310-cp310-macosx_10_9_x86_64.whl/numpy/__init__.pyi/").unwrap(),
+        ),
+        create_inspector_url_1: (
+            ("numpy", "1.24.3", Url::parse("https://files.pythonhosted.org/packages/f3/23/7cc851bae09cf4db90d42a701dfe525780883ada86bece45e3da7a07e76b/numpy-1.24.3-cp310-cp310-macosx_10_9_x86_64.whl/numpy/typing/tests/data/fail/twodim_base.pyi").unwrap()),
+            Url::parse("https://inspector.pypi.io/project/numpy/1.24.3/packages/f3/23/7cc851bae09cf4db90d42a701dfe525780883ada86bece45e3da7a07e76b/numpy-1.24.3-cp310-cp310-macosx_10_9_x86_64.whl/numpy/typing/tests/data/fail/twodim_base.pyi/").unwrap()
+        ),
+        create_inspector_url_2: (
+            ("discord-py","2.2.3", Url::parse("https://files.pythonhosted.org/packages/36/ce/3ad5a63240b504722dada49d880f9f6250ab861baaba5d27df4f4cb3e34a/discord.py-2.2.3.tar.gz/discord.py-2.2.3/discord/app_commands/checks.py").unwrap()),
+            Url::parse("https://inspector.pypi.io/project/discord-py/2.2.3/packages/36/ce/3ad5a63240b504722dada49d880f9f6250ab861baaba5d27df4f4cb3e34a/discord.py-2.2.3.tar.gz/discord.py-2.2.3/discord/app_commands/checks.py/").unwrap()
+        ),
+        create_inspector_url_3: (
+            ("requests", "2.19.1", Url::parse("https://files.pythonhosted.org/packages/54/1f/782a5734931ddf2e1494e4cd615a51ff98e1879cbe9eecbdfeaf09aa75e9/requests-2.19.1.tar.gz/requests-2.19.1/LICENSE").unwrap()),
+            Url::parse("https://inspector.pypi.io/project/requests/2.19.1/packages/54/1f/782a5734931ddf2e1494e4cd615a51ff98e1879cbe9eecbdfeaf09aa75e9/requests-2.19.1.tar.gz/requests-2.19.1/LICENSE/").unwrap()
+        ),
+    }
+}