dragonfly_client_rs/client/
models.rs

1use color_eyre::Result;
2use serde::Serialize;
3use serde::{self, Deserialize};
4use std::collections::HashMap;
5use std::fmt::Display;
6use yara::{Compiler, Rules, ScanFlags};
7
8pub type ScanResult = Result<SubmitJobResultsSuccess, SubmitJobResultsError>;
9
10#[derive(Serialize, Debug)]
11#[serde(untagged)]
12#[serde(remote = "ScanResult")]
13enum ScanResultDef {
14    Ok(SubmitJobResultsSuccess),
15    Err(SubmitJobResultsError),
16}
17
18#[derive(Serialize)]
19pub struct ScanResultSerializer(#[serde(with = "ScanResultDef")] ScanResult);
20
21impl From<ScanResult> for ScanResultSerializer {
22    fn from(value: ScanResult) -> Self {
23        Self(value)
24    }
25}
26
27#[derive(Debug, Serialize, PartialEq)]
28pub struct SubmitJobResultsSuccess {
29    pub name: String,
30    pub version: String,
31    pub attempt: u64,
32    pub assignment_id: String,
33    pub score: i64,
34    pub inspector_url: Option<String>,
35
36    /// Contains all rule identifiers matched for the entire release.
37    pub rules_matched: Vec<String>,
38
39    /// The commit hash of the ruleset used to produce these results.
40    pub commit: String,
41}
42
43#[derive(Debug, Serialize, PartialEq)]
44pub struct SubmitJobResultsError {
45    pub name: String,
46    pub version: String,
47    pub attempt: u64,
48    pub assignment_id: String,
49    pub reason: String,
50}
51
52impl Display for SubmitJobResultsError {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        writeln!(f, "Name: {}", self.name)?;
55        writeln!(f, "Version: {}", self.version)?;
56        writeln!(f, "Reason: {}", self.reason)?;
57
58        Ok(())
59    }
60}
61
62#[derive(Debug, Deserialize, PartialEq)]
63pub struct Job {
64    pub hash: String,
65    pub name: String,
66    pub version: String,
67    pub distributions: Vec<String>,
68    pub attempt: u64,
69    pub assignment_id: String,
70}
71
72#[derive(Debug, Deserialize)]
73pub struct RulesResponse {
74    pub hash: String,
75    pub rules: HashMap<String, String>,
76}
77
78impl RulesResponse {
79    /// Compile the rules from the response
80    pub fn compile(&self) -> Result<Rules> {
81        let mut sources = self.rules.iter().collect::<Vec<_>>();
82        sources.sort_unstable_by_key(|(name, _)| *name);
83        let rules_str = sources
84            .into_iter()
85            .map(|(_, source)| source.as_str())
86            .collect::<Vec<_>>()
87            .join("\n");
88
89        let mut compiled_rules = Compiler::new()?
90            .add_rules_str(&rules_str)?
91            .compile_rules()?;
92        compiled_rules.set_flags(ScanFlags::FAST_MODE);
93
94        Ok(compiled_rules)
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::Job;
101
102    #[test]
103    fn independently_loaded_rules_compile_in_the_same_order() {
104        for _ in 0..16 {
105            let response: super::RulesResponse = serde_json::from_str(
106                r#"{
107                "hash":"same", "rules":{
108                    "z":"rule last { condition: true }",
109                    "a":"rule first { condition: true }"
110                }
111            }"#,
112            )
113            .unwrap();
114            let rules = response.compile().unwrap();
115            let matches = rules.scan_mem(b"same bytes", 10).unwrap();
116            assert_eq!(
117                matches.iter().map(|r| r.identifier).collect::<Vec<_>>(),
118                vec!["first", "last"]
119            );
120        }
121    }
122
123    #[test]
124    fn job_deserializes_assignment_lease() {
125        let job: Job = serde_json::from_str(
126            r#"{
127                "hash": "rules-commit",
128                "name": "example",
129                "version": "1.2.3",
130                "distributions": ["https://example.com/example.whl"],
131                "attempt": 2,
132                "assignment_id": "4e3702e8-27a3-46e6-b51c-4779a94fa4ab"
133            }"#,
134        )
135        .unwrap();
136
137        assert_eq!(job.attempt, 2);
138        assert_eq!(job.assignment_id, "4e3702e8-27a3-46e6-b51c-4779a94fa4ab");
139    }
140
141    #[test]
142    fn job_requires_assignment_lease() {
143        let result = serde_json::from_str::<Job>(
144            r#"{
145                "hash": "rules-commit",
146                "name": "example",
147                "version": "1.2.3",
148                "distributions": []
149            }"#,
150        );
151
152        assert!(result.is_err());
153    }
154}
dragonfly_client_rs/client/models.rs

dragonfly_client_rs/client/
models.rs
