assemblyline_models/messages/
submission.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::types::{Sid, Wildcard, Sha256};
6
7use crate::messages::dispatching::FileTreeData;
8
9pub use crate::datastore::submission::{File, SubmissionParams};
10
11#[derive(Serialize, Deserialize)]
12pub enum MessageType {
13    SubmissionIngested,
14    SubmissionReceived,
15    SubmissionStarted,
16    SubmissionCompleted
17}
18
19/// Notification Model
20#[derive(Serialize, Deserialize, Default, Debug, Clone)]
21#[serde(default)]
22pub struct Notification {
23    /// Queue to publish the completion message
24    pub queue: Option<String>,
25    /// Notify only if this score threshold is met
26    pub threshold: Option<i32>,
27}
28
29/// Submission Model
30#[derive(Serialize, Deserialize, Debug, Clone)]
31pub struct Submission {
32    /// Submission ID to use
33    pub sid: Sid,
34    /// Message time
35    #[serde(default="chrono::Utc::now")]
36    pub time: chrono::DateTime<chrono::Utc>,
37    /// File block
38    pub files: Vec<File>,
39    /// Metadata submitted with the file
40    pub metadata: HashMap<String, Wildcard>,
41    /// Notification queue parameters
42    #[serde(default)]
43    pub notification: Notification,
44    /// Parameters of the submission
45    pub params: SubmissionParams,
46    /// Key used to track groups of submissions ingester will see as duplicates
47    pub scan_key: Option<String>,
48    ///List of error keys
49    #[serde(default)]
50    pub errors: Vec<String>,
51    /// Result key value mapping
52    #[serde(default)]
53    pub results: HashMap<String, crate::datastore::result::Result>,
54    /// File sha256 map to File tree of the submission
55    #[serde(default)]
56    pub file_tree: HashMap<Sha256, FileTreeData>,
57    /// File sha256 map to file info
58    #[serde(default)]
59    pub file_infos: HashMap<Sha256, super::task::FileInfo>,
60
61}
62
63impl Submission {
64
65    pub fn new(sid: Sid,time: chrono::DateTime<chrono::Utc>, params: SubmissionParams) -> Self {
66        Self {
67            sid,
68            time,
69            params,
70
71            files: Default::default(),
72            metadata: Default::default(),
73            notification: Default::default(),
74            scan_key: Default::default(),
75            errors: Default::default(),
76            results: Default::default(),
77            file_infos: Default::default(),
78            file_tree: Default::default()
79        }
80    }
81
82    pub fn set_notification(mut self, notification: Notification) -> Self {
83        self.notification = notification;
84        self
85    }
86
87    pub fn set_files(mut self, files: Vec<File>) -> Self {
88        self.files = files;
89        self
90    }
91
92    pub fn set_errors(mut self, errors: Vec<String>) -> Self {
93        self.errors = errors;
94        self
95    }
96
97    pub fn set_metadata(mut self, metadata: HashMap<String, Wildcard>) -> Self {
98        self.metadata = metadata;
99        self
100    }
101
102    pub fn set_scan_key(mut self, scan_key: Option<String>) -> Self {
103        self.scan_key = scan_key;
104        self
105    }
106
107    pub fn set_results(mut self, results: HashMap<String, crate::datastore::result::Result>) -> Self {
108        self.results = results;
109        self
110    }
111
112    pub fn set_file_infos(mut self, file_infos: HashMap<Sha256, super::task::FileInfo>) -> Self {
113        self.file_infos = file_infos;
114        self
115    }
116
117    pub fn set_file_tree(mut self, file_tree: HashMap<Sha256, FileTreeData>) -> Self{
118        self.file_tree = file_tree;
119        self
120    }
121}
122
123impl From<&crate::datastore::submission::Submission> for Submission {
124    fn from(value: &crate::datastore::submission::Submission) -> Self {
125        Self {
126            sid: value.sid,
127            files: value.files.clone(),
128            metadata: value.metadata.clone(),
129            params: value.params.clone(),
130            scan_key: value.scan_key.clone(),
131            time: chrono::Utc::now(),
132            notification: Default::default(),
133            errors: Default::default(),
134            results: Default::default(),
135            file_infos: Default::default(),
136            file_tree: Default::default()
137        }
138    }
139}
140
141/// Model of Submission Message
142#[derive(Serialize, Deserialize)]
143pub struct SubmissionMessage {
144    /// Body of the message
145    pub msg: Submission,
146    /// Class to use to load the message as an object
147    #[serde(default="default_message_loader")]
148    pub msg_loader: String,
149        /// Type of message
150    pub msg_type: MessageType,
151    /// Sender of the message
152    pub sender: String,
153}
154
155pub fn default_message_loader() -> String {"assemblyline.odm.messages.submission.SubmissionMessage".to_string()}
156
157impl SubmissionMessage {
158    pub fn ingested(sub: Submission) -> Self {
159        Self {
160            msg: sub,
161            msg_loader: default_message_loader(),
162            msg_type: MessageType::SubmissionIngested,
163            sender: "ingester".to_owned()
164        }
165    }
166    pub fn started(sub: Submission) -> Self {
167        Self {
168            msg: sub,
169            msg_loader: default_message_loader(),
170            msg_type: MessageType::SubmissionStarted,
171            sender: "dispatcher".to_owned()
172        }
173    }
174    pub fn completed(sub: Submission, sender: String) -> Self {
175        Self {
176            msg: sub,
177            msg_loader: default_message_loader(),
178            msg_type: MessageType::SubmissionCompleted,
179            sender
180        }
181    }
182    pub fn received(sub: Submission, sender: String) -> Self {
183        Self {
184            msg: sub,
185            msg_loader: default_message_loader(),
186            msg_type: MessageType::SubmissionReceived,
187            sender
188        }
189    }
190}
assemblyline_models/messages/submission.rs

assemblyline_models/messages/
submission.rs
