assemblyline_models/messages/
task.rs

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

assemblyline_models/messages/
task.rs
