assemblyline_models/messages/
task.rs

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

assemblyline_models/messages/
task.rs
