Skip to main content

assemblyline_models/
lib.rs

1#![allow(clippy::needless_return)]
2use std::fmt::Display;
3
4use serde::Deserialize;
5
6#[cfg(feature = "rand")]
7use rand::RngExt;
8
9
10pub mod datastore;
11pub mod config;
12pub mod messages;
13pub mod serialize;
14pub mod meta;
15pub mod types;
16
17pub use meta::ElasticMeta;
18pub use types::classification::{disable_global_classification, set_global_classification};
19
20pub const HEXCHARS: [char; 16] = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'];
21
22pub trait Readable: for <'de> Deserialize<'de> {
23    fn set_from_archive(&mut self, from_archive: bool);
24}
25
26#[derive(Debug)]
27pub enum ModelError {
28    InvalidSha256(String),
29    InvalidMd5(String),
30    InvalidSha1(String),
31    InvalidSid(String),
32    InvalidSSDeep(String),
33    ClassificationNotInitialized,
34    InvalidClassification(Option<assemblyline_markings::errors::Errors>),
35}
36
37impl Display for ModelError {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            ModelError::InvalidSha256(content) => f.write_fmt(format_args!("Invalid value provided for a sha256: {content}")),
41            ModelError::InvalidMd5(content) => f.write_fmt(format_args!("Invalid value provided for a md5: {content}")),
42            ModelError::InvalidSha1(content) => f.write_fmt(format_args!("Invalid value provided for a sha1: {content}")),
43            ModelError::InvalidSid(content) => f.write_fmt(format_args!("Invalid value provided for a sid: {content}")),
44            ModelError::ClassificationNotInitialized => f.write_str("The classification engine has not been initialized."),
45            ModelError::InvalidClassification(_) => f.write_str("An invalid classification string was provided."),
46            ModelError::InvalidSSDeep(content) =>  f.write_fmt(format_args!("Invalid value provided for a ssdeep hash: {content}")),
47        }
48    }
49}
50
51impl From<base62::DecodeError> for ModelError {
52    fn from(value: base62::DecodeError) -> Self {
53        Self::InvalidSid(value.to_string())
54    }
55}
56
57impl From<assemblyline_markings::errors::Errors> for ModelError {
58    fn from(value: assemblyline_markings::errors::Errors) -> Self {
59        Self::InvalidClassification(Some(value))
60    }
61}
62
63impl std::error::Error for ModelError {}
64
65
66const WORDS: [&str; 187] = ["The", "Cyber", "Centre", "stays", "on", "the", "cutting", "edge", "of", "technology", "by",
67    "working", "with", "commercial", "vendors", "of", "cyber", "security", "technology", "to", "support", "their",
68    "development", "of", "enhanced", "cyber", "defence", "tools", "To", "do", "this", "our", "experts", "survey",
69    "the", "cyber", "security", "market", "evaluate", "emerging", "technologies", "in", "order", "to", "determine",
70    "their", "potential", "to", "improve", "cyber", "security", "across", "the", "country", "The", "Cyber", "Centre",
71    "supports", "innovation", "by", "collaborating", "with", "all", "levels", "of", "government", "private", "industry",
72    "academia", "to", "examine", "complex", "problems", "in", "cyber", "security", "We", "are", "constantly",
73    "engaging", "partners", "to", "promote", "an", "open", "innovative", "environment", "We", "invite", "partners",
74    "to", "work", "with", "us", "but", "also", "promote", "other", "Government", "of", "Canada", "innovation",
75    "programs", "One", "of", "our", "key", "partnerships", "is", "with", "the", "Government", "of", "Canada", "Build",
76    "in", "Canada", "Innovation", "Program", "BCIP", "The", "BCIP", "helps", "Canadian", "companies", "of", "all",
77    "sizes", "transition", "their", "state", "of", "the", "art", "goods", "services", "from", "the", "laboratory",
78    "to", "the", "marketplace", "For", "certain", "cyber", "security", "innovations", "the", "Cyber", "Centre",
79    "performs", "the", "role", "of", "technical", "authority", "We", "evaluate", "participating", "companies",
80    "new", "technology", "provide", "feedback", "in", "order", "to", "assist", "them", "in", "bringing", "their",
81    "product", "to", "market", "To", "learn", "more", "about", "selling", "testing", "an", "innovation", "visit",
82    "the", "BCIP", "website"];
83
84#[cfg(feature = "rand")]
85pub fn random_word<R: rand::Rng + ?Sized>(prng: &mut R) -> String {
86    WORDS[prng.random_range(0..WORDS.len())].to_string()
87}
88
89#[cfg(feature = "rand")]
90pub fn random_words<R: rand::Rng + ?Sized>(prng: &mut R, count: usize) -> Vec<String> {
91    let mut output = vec![];
92    while output.len() < count {
93        output.push(WORDS[prng.random_range(0..WORDS.len())].to_string())
94    }
95    output
96}
97
98#[cfg(feature = "rand")]
99pub fn random_hex<R: rand::prelude::Rng + ?Sized>(rng: &mut R, size: usize) -> String {
100    let mut buffer = String::with_capacity(size);
101    for _ in 0..size {
102        let index = rng.random_range(0..HEXCHARS.len());
103        buffer.push(HEXCHARS[index]);
104    }
105    buffer
106}
107
108#[cfg(test)]
109mod test {
110    use rand::RngExt;
111
112    use crate::types::{SSDeepHash, Sha1, Sha256, MD5};
113
114    #[test]
115    fn random_ssdeep() {
116        let mut prng = rand::rng();
117        for _ in 0..100 {
118            let hash: SSDeepHash = prng.random();
119            assert_eq!(hash, hash.to_string().parse().unwrap());
120        }
121    }
122
123    #[test]
124    fn random_sha256() {
125        let mut prng = rand::rng();
126        for _ in 0..100 {
127            let hash: Sha256 = prng.random();
128            assert_eq!(hash, hash.to_string().parse().unwrap());
129        }
130    }
131
132    #[test]
133    fn random_sha1() {
134        let mut prng = rand::rng();
135        for _ in 0..100 {
136            let hash: Sha1 = prng.random();
137            assert_eq!(hash, hash.to_string().parse().unwrap());
138        }
139    }
140
141    #[test]
142    fn random_md5() {
143        let mut prng = rand::rng();
144        for _ in 0..100 {
145            let hash: MD5 = prng.random();
146            assert_eq!(hash, hash.to_string().parse().unwrap());
147        }
148    }
149}
150