assemblyline_models/datastore/
result.rs

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

assemblyline_models/datastore/
result.rs
