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

assemblyline_models/messages/
task.rs
