assemblyline_models/messages/
task.rs

1use std::collections::HashMap;
2
3use md5::Digest;
4use rand::Rng;
5use serde::{Serialize, Deserialize};
6
7use crate::datastore::tagging::TagValue;
8use crate::random_word;
9use crate::types::{JsonMap, SSDeepHash, ServiceName, Sha1, Sha256, Sid, Wildcard, MD5};
10use crate::{datastore::file::URIInfo, config::ServiceSafelist};
11
12
13/// File Information
14#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
15pub struct FileInfo {
16    /// The output from libmagic which was used to determine the tag
17    pub magic: String,
18    /// MD5 of the file
19    pub md5: MD5,
20    /// The libmagic mime type
21    pub mime: Option<String>,
22    /// SHA1 hash of the file
23    pub sha1: Sha1,
24    /// SHA256 hash of the file
25    pub sha256: Sha256,
26    /// Size of the file in bytes
27    pub size: u64,
28    /// SSDEEP hash of the file"
29    pub ssdeep: Option<SSDeepHash>,
30    /// TLSH hash of the file"
31    pub tlsh: Option<String>,
32    /// Type of file as identified by Assemblyline
33    #[serde(rename="type")]
34    pub file_type: String,
35    /// URI structure to speed up specialty file searching
36    pub uri_info: Option<URIInfo>,
37}
38
39/// Tag Item
40#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
41pub struct TagItem {
42    /// Type of tag item
43    #[serde(rename="type")]
44    pub tag_type: String,
45    ///Short version of tag type
46    pub short_type: String,
47    /// Value of tag item
48    pub value: TagValue,
49    /// Score of tag item
50    #[serde(skip_serializing_if="Option::is_none")]
51    pub score: Option<i32>,
52}
53
54
55/// Data Item
56#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
57pub struct DataItem {
58    pub name: String,
59    pub value: serde_json::Value,
60}
61
62/// Service Task Model
63#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
64pub struct Task {
65    /// A random ID to differentiate this task
66    pub task_id: u64,
67    /// Id of the dispatcher that issued this task
68    pub dispatcher: String,
69    pub dispatcher_address: String,
70    /// Submission ID
71    pub sid: Sid,
72    /// Metadata associated to the submission
73    pub metadata: HashMap<String, Wildcard>,
74    /// Minimum classification of the file being scanned
75    pub min_classification: String,
76    /// File info block
77    pub fileinfo: FileInfo,
78    /// File name
79    pub filename: String,
80    /// Service name
81    pub service_name: ServiceName,
82    /// Service specific parameters
83    pub service_config: JsonMap,
84    /// File depth relative to initital submitted file
85    pub depth: u32,
86    /// Maximum number of files that submission can have
87    pub max_files: i32,
88    /// Task TTL
89    pub ttl: i32,
90
91    /// List of tags
92    pub tags: Vec<TagItem>,
93    /// Temporary submission data
94    pub temporary_submission_data: Vec<DataItem>,
95
96    /// Perform deep scanning
97    pub deep_scan: bool,
98
99    /// Whether the service cache should be ignored during the processing of this task
100    pub ignore_cache: bool, 
101
102    /// Whether the service should ignore the dynamic recursion prevention or not
103    pub ignore_recursion_prevention: bool,
104
105    /// Should the service filter it's output?
106    pub ignore_filtering: bool,
107
108    /// Priority for processing order
109    pub priority: i32,
110
111    /// Safelisting configuration (as defined in global configuration)
112    #[serde(default="task_default_safelist_config")]
113    pub safelist_config: ServiceSafelist, // ", default={'enabled': False})
114}
115
116#[cfg(feature = "rand")]
117impl rand::distr::Distribution<Task> for rand::distr::StandardUniform {
118    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Task {
119        Task {
120            task_id: rng.random(),
121            dispatcher: random_word(rng),
122            dispatcher_address: "localhost:8080".to_string(),
123            sid: rng.random(),
124            metadata: Default::default(),
125            min_classification: Default::default(),
126            fileinfo: FileInfo { 
127                magic: "".to_string(), 
128                md5: rng.random(), 
129                mime: None, 
130                sha1: rng.random(), 
131                sha256: rng.random(), 
132                size: rng.random(), 
133                ssdeep: Some(rng.random()),
134                tlsh: None, 
135                file_type: "unknown".to_owned(), 
136                uri_info: None 
137            },
138            filename: random_word(rng),
139            service_name: ServiceName::from_string(random_word(rng)),
140            service_config: Default::default(),
141            depth: rng.random(),
142            max_files: rng.random(),
143            ttl: rng.random_range(0..100),
144            tags: Default::default(),
145            temporary_submission_data: Default::default(),
146            deep_scan: rng.random(),
147            ignore_cache: rng.random(),
148            ignore_recursion_prevention: rng.random(),
149            ignore_filtering: rng.random(),
150            priority: rng.random(),
151            safelist_config: Default::default(),
152        }
153    }
154}
155
156pub fn task_default_safelist_config() -> ServiceSafelist {
157    ServiceSafelist {
158        enabled: false,
159        ..Default::default()
160    }
161}
162
163impl Task {
164    pub fn make_key(sid: Sid, service_name: &str, sha: &Sha256) -> String {
165        format!("{sid}_{service_name}_{sha}")
166    }
167
168    pub fn key(&self) -> String {
169        Self::make_key(self.sid, &self.service_name, &self.fileinfo.sha256)
170    }
171
172    pub fn signature(&self) -> TaskSignature {
173        TaskSignature { 
174            task_id: self.task_id, 
175            sid: self.sid, 
176            service: self.service_name, 
177            hash: self.fileinfo.sha256.clone() 
178        }
179    } 
180}
181
182pub fn generate_conf_key(service_tool_version: Option<&str>, task: Option<&Task>, partial: Option<bool>) -> Result<String, serde_json::Error> {
183    if let Some(task) = task {
184        let service_config = serde_json::to_string(&{
185            let mut pairs: Vec<_> = task.service_config.iter().collect();
186            pairs.sort_unstable_by_key(|row| row.0);
187            pairs
188        })?;
189
190        let submission_params_str = serde_json::to_string(&[
191            ("deep_scan", serde_json::json!(task.deep_scan)),
192            ("ignore_filtering", serde_json::json!(task.ignore_filtering)),
193            ("max_files", serde_json::json!(task.max_files)),
194            ("min_classification", serde_json::json!(task.min_classification)),
195        ])?;
196
197        let ignore_salt = if task.ignore_cache || partial.unwrap_or_default() {
198             &rand::rng().random::<u128>().to_string()
199        } else {
200            "None"
201        };
202
203        let service_tool_version = service_tool_version.unwrap_or("None");
204        let total_str = format!("{service_tool_version}_{service_config}_{submission_params_str}_{ignore_salt}");
205
206        // get an md5 hash
207        let mut hasher = md5::Md5::new();
208        hasher.update(total_str);
209        let hash = hasher.finalize();
210        
211        // truncate it to 8 bytes and interpret it as a number
212        let number = u64::from_be_bytes(hash[0..8].try_into().unwrap());
213        
214        // encode it as a string
215        Ok(base62::encode(number))
216    } else {
217        Ok("0".to_string())
218    }
219}
220
221
222#[derive(Hash, PartialEq, Eq)]
223pub struct TaskSignature {
224    pub task_id: u64,
225    pub sid: Sid,
226    pub service: ServiceName,
227    pub hash: Sha256,
228}
229
230
231
232/// Service Task Model
233#[derive(Serialize, Deserialize)]
234pub struct TaskToken {
235    pub task_id: u64,
236    pub dispatcher: String,
237}
238
239// ============================================================================
240//MARK: Responses 
241
242#[derive(Serialize, Deserialize, Clone, Debug)]
243pub struct ResultSummary {
244    pub key: String,
245    pub drop: bool,
246    pub score: i32,
247    pub partial: bool,
248    pub children: Vec<(Sha256, String)>
249}
250
251#[derive(Serialize, Deserialize)]
252pub enum ServiceResponse {
253    Result(Box<ServiceResult>),
254    Error(Box<ServiceError>),
255}
256
257impl std::fmt::Debug for ServiceResponse {
258    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
259        f.debug_struct("ServiceResponse").field("result", &matches!(self, ServiceResponse::Result(..))).finish()
260    }
261}
262
263impl ServiceResponse {
264    pub fn sid(&self) -> Sid {
265        match self {
266            ServiceResponse::Result(item) => item.sid,
267            ServiceResponse::Error(item) => item.sid,
268        }
269    }
270
271    pub fn sha256(&self) -> Sha256 {
272        match self {
273            ServiceResponse::Result(item) => item.sha256.clone(),
274            ServiceResponse::Error(item) => item.service_task.fileinfo.sha256.clone(),
275        }
276    }
277
278    pub fn service_name(&self) -> ServiceName {
279        match self {
280            ServiceResponse::Result(item) => item.service_name,
281            ServiceResponse::Error(item) => item.service_task.service_name,
282        }
283    }
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
287pub struct TagEntry {
288    pub score: i32,    
289    #[serde(rename="type")]
290    pub tag_type: String,
291    pub value: TagValue,
292
293}
294
295#[derive(Serialize, Deserialize)]
296pub struct ServiceResult {
297    pub dynamic_recursion_bypass: Vec<Sha256>,
298    pub sid: Sid,
299    pub sha256: Sha256,
300    pub service_name: ServiceName,
301    pub service_version: String,
302    pub service_tool_version: Option<String>,
303    pub expiry_ts: Option<chrono::DateTime<chrono::Utc>>,
304    pub result_summary: ResultSummary,
305    pub tags: HashMap<String, TagEntry>,
306    pub extracted_names: HashMap<Sha256, String>,
307    pub temporary_data: JsonMap,
308    pub extra_errors: Vec<String>,
309}
310
311#[derive(Serialize, Deserialize)]
312pub struct ServiceError {
313    pub sid: Sid,
314    pub service_task: Task,
315    pub error: crate::datastore::Error,
316    pub error_key: String,
317}
assemblyline_models/messages/task.rs

assemblyline_models/messages/
task.rs
