assemblyline_models/messages/
submission.rs

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