assemblyline_models/datastore/
result.rs

1// from collections import defaultdict
2// from typing import Any, Dict
3
4// from assemblyline import odm
5// from assemblyline.common import forge
6// from assemblyline.common.caching import generate_conf_key
7// from assemblyline.common.dict_utils import flatten
8// from assemblyline.common.tagging import tag_dict_to_list
9// from assemblyline.odm.models.tagging import Tagging
10
11use std::collections::HashMap;
12
13use chrono::{DateTime, Utc};
14use serde::{Serialize, Deserialize};
15use serde_with::{SerializeDisplay, DeserializeFromStr};
16use struct_metadata::Described;
17
18use crate::datastore::tagging::LayoutError;
19use crate::messages::task::{generate_conf_key, TagEntry, Task};
20use crate::types::strings::Keyword;
21use crate::{random_word, ElasticMeta, Readable};
22use crate::types::{ClassificationString, ExpandingClassification, ServiceName, Sha256, Text};
23
24use super::tagging::Tagging;
25
26#[derive(SerializeDisplay, DeserializeFromStr, strum::Display, strum::EnumString, Debug, Described, Clone)]
27#[metadata_type(ElasticMeta)]
28#[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
29pub enum BodyFormat {
30    Text,
31    MemoryDump,
32    GraphData,
33    Url,
34    Json,
35    KeyValue,
36    ProcessTree,
37    Table,
38    Image,
39    Multi,
40    OrderedKeyValue,
41    Timeline,
42}
43
44// This needs to match the PROMOTE_TO StringTable in
45// assemblyline-v4-service/assemblyline_v4_service/common/result.py.
46// Any updates here need to go in that StringTable also.
47#[derive(Serialize, Deserialize, Debug, Described, Clone)]
48#[metadata_type(ElasticMeta)]
49#[serde(rename_all="SCREAMING_SNAKE_CASE")]
50pub enum PromoteTo {
51    Screenshot, 
52    Entropy, 
53    UriParams
54}
55
56// constants = forge.get_constants()
57
58#[derive(Serialize, Deserialize, Debug, Described, Clone)]
59#[metadata_type(ElasticMeta)]
60#[metadata(index=true, store=false)]
61pub struct Attack {
62    /// ID
63    #[metadata(copyto="__text__")]
64    pub attack_id: String,
65    /// Pattern Name
66    #[metadata(copyto="__text__")]
67    pub pattern: String,
68    /// Categories
69    pub categories: Vec<String>,
70}
71
72/// Heuristic Signatures
73#[derive(Serialize, Deserialize, Debug, Described, Clone)]
74#[metadata_type(ElasticMeta)]
75#[metadata(index=true, store=false)]
76pub struct Signature {
77    /// Name of the signature that triggered the heuristic
78    #[metadata(copyto="__text__")]
79    pub name: String,
80    /// Number of times this signature triggered the heuristic
81    #[serde(default = "default_signature_frequency")]
82    pub frequency: i32,
83    /// Is the signature safelisted or not
84    #[serde(default)]
85    pub safe: bool,
86}
87
88fn default_signature_frequency() -> i32 { 1 }
89
90/// Heuristic associated to the Section
91#[derive(Serialize, Deserialize, Debug, Described, Clone)]
92#[metadata_type(ElasticMeta)]
93#[metadata(index=true, store=false)]
94pub struct Heuristic {
95    /// ID of the heuristic triggered
96    #[metadata(copyto="__text__")]
97    pub heur_id: String,
98    /// Name of the heuristic
99    #[metadata(copyto="__text__")]
100    pub name: String,
101    /// List of Att&ck IDs related to this heuristic
102    #[serde(default)]
103    pub attack: Vec<Attack>,
104    /// List of signatures that triggered the heuristic
105    #[serde(default)]
106    pub signature: Vec<Signature>,
107    /// Calculated Heuristic score
108    pub score: i32,
109}
110
111/// Result Section
112#[derive(Serialize, Deserialize, Debug, Described, Clone)]
113#[metadata_type(ElasticMeta)]
114#[metadata(index=true, store=false)]
115pub struct Section {
116    /// Should the section be collapsed when displayed?
117    #[serde(default)]
118    pub auto_collapse: bool,
119    /// Text body of the result section
120    #[metadata(copyto="__text__")]
121    pub body: Option<Text>,
122    /// Classification of the section
123    pub classification: ClassificationString,
124    /// Type of body in this section
125    #[metadata(index=false)]
126    pub body_format: BodyFormat,
127    /// Configurations for the body of this section
128    #[metadata(index=false)]
129    pub body_config: Option<HashMap<String, serde_json::Value>>,
130    /// Depth of the section
131    #[metadata(index=false)]
132    pub depth: i32,
133    /// Heuristic used to score result section
134    pub heuristic: Option<Heuristic>,
135    /// List of tags associated to this section
136    #[serde(default)]
137    pub tags: Tagging,
138    /// List of safelisted tags
139    #[serde(default)]
140    #[metadata(store=false, mapping="flattenedobject")]
141    pub safelisted_tags: HashMap<String, Vec<Keyword>>,
142    /// Title of the section
143    #[metadata(copyto="__text__")]
144    pub title_text: Text,
145    /// This is the type of data that the current section should be promoted to.
146    pub promote_to: Option<PromoteTo>,
147}
148
149/// Result Body
150#[derive(Serialize, Deserialize, Debug, Default, Described, Clone)]
151#[metadata_type(ElasticMeta)]
152#[metadata(index=true, store=true)]
153pub struct ResultBody {
154    /// Aggregate of the score for all heuristics
155    #[serde(default)]
156    pub score: i32,
157    /// List of sections
158    #[serde(default)]
159    pub sections: Vec<Section>,
160}
161
162/// Service Milestones
163#[derive(Serialize, Deserialize, Debug, Default, Described, Clone)]
164#[metadata_type(ElasticMeta)]
165#[metadata(index=false, store=false)]
166pub struct Milestone {
167    /// Date the service started scanning
168    pub service_started: DateTime<Utc>,
169    /// Date the service finished scanning
170    pub service_completed: DateTime<Utc>,
171}
172
173#[cfg(feature = "rand")]
174impl rand::distr::Distribution<Milestone> for rand::distr::StandardUniform {
175    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Milestone {
176        let service_started = chrono::Utc::now() - chrono::TimeDelta::hours(rng.random_range(1..200));
177        let duration = chrono::TimeDelta::seconds(rng.random_range(1..900));
178        Milestone {
179            service_started,
180            service_completed: service_started + duration
181        }
182    }
183}
184
185
186/// File related to the Response
187#[derive(Serialize, Deserialize, Debug, Described, Clone)]
188#[metadata_type(ElasticMeta)]
189#[metadata(index=true, store=false)]
190pub struct File {
191    /// Name of the file
192    #[metadata(copyto="__text__")]
193    pub name: String,
194    /// SHA256 of the file
195    #[metadata(copyto="__text__")]
196    pub sha256: Sha256,
197    /// Description of the file
198    #[metadata(copyto="__text__")]
199    pub description: Text,
200    /// Classification of the file
201    pub classification: ClassificationString,
202    /// Is this an image used in an Image Result Section?
203    #[serde(default)]
204    pub is_section_image: bool,
205    /// File relation to parent, if any.
206    #[serde(default = "default_file_parent_relation")]
207    pub parent_relation: Text,
208    /// Allow file to be analysed during Dynamic Analysis even if Dynamic Recursion Prevention is enabled.
209    #[serde(default)]
210    pub allow_dynamic_recursion: bool,
211}
212
213fn default_file_parent_relation() -> Text { Text("EXTRACTED".to_owned()) }
214
215/// Response Body of Result
216#[derive(Serialize, Deserialize, Debug, Described, Clone)]
217#[metadata_type(ElasticMeta)]
218#[metadata(index=true, store=true)]
219pub struct ResponseBody {
220    /// Milestone block
221    #[serde(default)]
222    pub milestones: Milestone,
223    /// Version of the service
224    #[metadata(store=false)]
225    pub service_version: String,
226    /// Name of the service that scanned the file
227    #[metadata(copyto="__text__")]
228    pub service_name: ServiceName,
229    /// Tool version of the service
230    #[serde(default)]
231    #[metadata(copyto="__text__")]
232    pub service_tool_version: Option<String>,
233    /// List of supplementary files
234    #[serde(default)]
235    pub supplementary: Vec<File>,
236    /// List of extracted files
237    #[serde(default)]
238    pub extracted: Vec<File>,
239    /// Context about the service
240    #[serde(default)]
241    #[metadata(index=false, store=false)]
242    pub service_context: Option<String>,
243    /// Debug info about the service
244    #[serde(default)]
245    #[metadata(index=false, store=false)]
246    pub service_debug_info: Option<String>,
247}
248
249
250#[cfg(feature = "rand")]
251impl rand::distr::Distribution<ResponseBody> for rand::distr::StandardUniform {
252    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> ResponseBody {
253        ResponseBody {
254            milestones: rng.random(),
255            service_version: random_word(rng),
256            service_name: ServiceName::from_string(random_word(rng)),
257            service_tool_version: None,
258            supplementary: Default::default(),
259            extracted: Default::default(),
260            service_context: Default::default(),
261            service_debug_info: Default::default(),
262        }
263    }
264}
265
266
267/// Result Model
268#[derive(Serialize, Deserialize, Debug, Described, Clone)]
269#[metadata_type(ElasticMeta)]
270#[metadata(index=true, store=true)]
271pub struct Result {
272    /// Timestamp indicating when the result was archived.
273    pub archive_ts: Option<DateTime<Utc>>,
274    /// Aggregate classification for the result
275    #[serde(flatten)]
276    pub classification: ExpandingClassification,
277    /// Date at which the result object got created
278    #[serde(default="chrono::Utc::now")]
279    pub created: DateTime<Utc>,
280    /// Expiry timestamp
281    #[metadata(store=false)]
282    pub expiry_ts: Option<DateTime<Utc>>,
283    /// The body of the response from the service
284    pub response: ResponseBody,
285    /// The result body
286    #[serde(default)]
287    pub result: ResultBody,
288    /// SHA256 of the file the result object relates to
289    #[metadata(store=false)]
290    pub sha256: Sha256,
291    /// What type information is given along with this result
292    #[serde(rename = "type")]
293    pub result_type: Option<String>,
294    /// ???
295    pub size: Option<i32>,
296    /// Use to not pass to other stages after this run
297    #[serde(default)]
298    pub drop_file: bool,
299    /// Invalidate the current result cache creation
300    #[serde(default)]
301    pub partial: bool,
302    /// Was loaded from the archive
303    #[serde(default)]
304    #[metadata(index=false)]
305    pub from_archive: bool,
306}
307
308#[cfg(feature = "rand")]
309impl rand::distr::Distribution<Result> for rand::distr::StandardUniform {
310    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Result {
311        Result {
312            archive_ts: None,
313            classification: ExpandingClassification::try_unrestricted().unwrap(),
314            created: chrono::Utc::now(),
315            expiry_ts: None,
316            response: rng.random(),
317            result: Default::default(),
318            sha256: rng.random(),
319            result_type: None,
320            size: None,
321            partial: Default::default(),
322            drop_file: Default::default(),
323            from_archive: Default::default(),
324        }
325    }
326}
327
328impl Readable for Result {
329    fn set_from_archive(&mut self, from_archive: bool) {
330        self.from_archive = from_archive;
331    }
332}
333
334impl Result {
335    pub fn build_key(&self, task: Option<&Task>) -> std::result::Result<String, serde_json::Error> {
336        Self::help_build_key(
337            &self.sha256,
338            &self.response.service_name,
339            &self.response.service_version,
340            self.is_empty(),
341            self.partial,
342            self.response.service_tool_version.as_deref(),
343            task
344        )
345    }
346
347    pub fn help_build_key(sha256: &Sha256, service_name: &str, service_version: &str, is_empty: bool, partial: bool, service_tool_version: Option<&str>, task: Option<&Task>) -> std::result::Result<String, serde_json::Error> {
348        let mut key_list = vec![
349            sha256.to_string(),
350            service_name.replace('.', "_"),
351            format!("v{}", service_version.replace('.', "_")),
352            format!("c{}", generate_conf_key(service_tool_version, task, Some(partial))?),
353        ];
354
355        if is_empty {
356            key_list.push("e".to_owned())
357        }
358
359        Ok(key_list.join("."))
360    }
361
362    pub fn scored_tag_dict(&self) -> std::result::Result<HashMap<String, TagEntry>, LayoutError> {
363        let mut tags: HashMap<String, TagEntry> = Default::default();
364        // Save the tags and their score
365        for section in &self.result.sections {
366            for tag in section.tags.to_list(None)? {
367                let key = format!("{}:{}", tag.tag_type, tag.value);
368                let entry = tags.entry(key).or_insert(tag);
369                if let Some(heuristic) = &section.heuristic {
370                    entry.score += heuristic.score;
371                }
372            }
373        }
374        Ok(tags)
375    }
376
377    pub fn is_empty(&self) -> bool {
378        self.response.extracted.is_empty() && self.response.supplementary.is_empty() && self.result.sections.is_empty() && self.result.score == 0 && !self.partial
379    }
380}
assemblyline_models/datastore/result.rs

assemblyline_models/datastore/
result.rs
