assemblyline_models/messages/
submission.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::types::Wildcard;
6use crate::{ClassificationString, Sid};
7pub use crate::datastore::submission::{File, SubmissionParams};
8
9#[derive(Serialize, Deserialize)]
10pub enum MessageType {
11    SubmissionIngested, 
12    SubmissionReceived, 
13    SubmissionStarted, 
14    SubmissionCompleted
15}
16
17/// Notification Model
18#[derive(Serialize, Deserialize, Default, Debug, Clone)]
19#[serde(default)]
20pub struct Notification {
21    /// Queue to publish the completion message
22    pub queue: Option<String>,
23    /// Notify only if this score threshold is met
24    pub threshold: Option<i32>,
25}
26
27/// Submission Model
28#[derive(Serialize, Deserialize, Debug, Clone)]
29pub struct Submission {
30    /// Submission ID to use
31    pub sid: Sid,
32    /// Message time
33    #[serde(default="chrono::Utc::now")]
34    pub time: chrono::DateTime<chrono::Utc>,
35    /// File block
36    pub files: Vec<File>, 
37    /// Metadata submitted with the file
38    pub metadata: HashMap<String, Wildcard>, 
39    /// Notification queue parameters
40    #[serde(default)]
41    pub notification: Notification,
42    /// Parameters of the submission
43    pub params: SubmissionParams,
44    /// Key used to track groups of submissions ingester will see as duplicates
45    pub scan_key: Option<String>,
46}
47
48impl Submission {
49    pub fn new(classification: ClassificationString) -> Self {
50        Self { 
51            sid: Sid(0), 
52            time: chrono::Utc::now(), 
53            files: Default::default(), 
54            metadata: Default::default(), 
55            notification: Default::default(), 
56            params: SubmissionParams::new(classification), 
57            scan_key: Default::default() 
58        }
59    }
60}
61
62impl From<&crate::datastore::submission::Submission> for Submission {
63    fn from(value: &crate::datastore::submission::Submission) -> Self {
64        Self {
65            sid: value.sid,
66            files: value.files.clone(),
67            metadata: value.metadata.clone(),
68            params: value.params.clone(),
69            scan_key: value.scan_key.clone(),
70            time: chrono::Utc::now(),
71            notification: Default::default(),
72        }
73    }
74}
75
76/// Model of Submission Message
77#[derive(Serialize, Deserialize)]
78pub struct SubmissionMessage {
79    /// Body of the message
80    pub msg: Submission,
81    /// Class to use to load the message as an object
82    #[serde(default="default_message_loader")]
83    pub msg_loader: String,
84        /// Type of message
85    pub msg_type: MessageType,
86    /// Sender of the message
87    pub sender: String,
88}
89
90pub fn default_message_loader() -> String {"assemblyline.odm.messages.submission.SubmissionMessage".to_string()}
91
92impl SubmissionMessage {
93    pub fn ingested(sub: Submission) -> Self {
94        Self {
95            msg: sub,
96            msg_loader: default_message_loader(),
97            msg_type: MessageType::SubmissionIngested,
98            sender: "ingester".to_owned()
99        }
100    }
101    pub fn started(sub: Submission) -> Self {
102        Self {
103            msg: sub,
104            msg_loader: default_message_loader(),
105            msg_type: MessageType::SubmissionStarted,
106            sender: "dispatcher".to_owned()
107        }
108    }
109    pub fn completed(sub: Submission, sender: String) -> Self {
110        Self {
111            msg: sub,
112            msg_loader: default_message_loader(),
113            msg_type: MessageType::SubmissionCompleted,
114            sender
115        }
116    }
117    pub fn received(sub: Submission, sender: String) -> Self {
118        Self {
119            msg: sub,
120            msg_loader: default_message_loader(),
121            msg_type: MessageType::SubmissionReceived,
122            sender
123        }
124    }
125}