Skip to main content

assemblyline_models/types/
classification.rs

1use std::ops::Deref;
2use std::sync::{Arc, Mutex};
3
4use serde::{Serialize, Deserialize};
5use assemblyline_markings::classification::{ClassificationParser, NormalizeOptions};
6use struct_metadata::Described;
7
8use crate::types::JsonMap;
9use crate::{ElasticMeta, ModelError};
10
11enum ClassificationMode {
12    Uninitialized,
13    Disabled,
14    Configured(Arc<ClassificationParser>)
15}
16
17static GLOBAL_CLASSIFICATION: Mutex<ClassificationMode> = Mutex::new(ClassificationMode::Uninitialized);
18
19// #[cfg(feature = "local_classification")]
20// tokio::task_local! {
21//     static TASK_CLASSIFICATION: Option<Arc<ClassificationParser>>;
22// }
23
24pub fn disable_global_classification() {
25    *GLOBAL_CLASSIFICATION.lock().unwrap() = ClassificationMode::Disabled;
26}
27
28pub fn set_global_classification(config: Arc<ClassificationParser>) {
29    *GLOBAL_CLASSIFICATION.lock().unwrap() = ClassificationMode::Configured(config);
30}
31
32pub fn unrestricted_classification() -> String {
33    match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
34        ClassificationMode::Uninitialized => panic!("classification handling without defining parser"),
35        ClassificationMode::Disabled => "".to_string(),
36        ClassificationMode::Configured(parser) => parser.unrestricted().to_owned(),
37    }
38}
39
40pub fn unrestricted_classification_string() -> ClassificationString {
41    match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
42        ClassificationMode::Uninitialized => panic!("classification handling without defining parser"),
43        ClassificationMode::Disabled => ClassificationString(Default::default()),
44        ClassificationMode::Configured(parser) => ClassificationString::unrestricted(parser),
45    }
46}
47
48pub fn unrestricted_expanding_classification() -> ExpandingClassification {
49    match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
50        ClassificationMode::Uninitialized => panic!("classification handling without defining parser"),
51        ClassificationMode::Disabled => ExpandingClassification {
52            classification: Default::default(),
53            __access_lvl__: Default::default(),
54            __access_req__: Default::default(),
55            __access_grp1__: Default::default(),
56            __access_grp2__: Default::default(),
57        },
58        ClassificationMode::Configured(parser) => ExpandingClassification::unrestricted(parser),
59    }
60}
61
62
63// #[cfg(feature = "local_classification")]
64// pub async fn with_local_classification<F: std::future::Future>(config: Arc<ClassificationParser>, f: F)  {
65//     TASK_CLASSIFICATION.scope(Some(config), f).await
66// }
67
68// MARK: ExpandingClassification
69
70/// Expanding classification type
71#[derive(Debug, Clone, Eq)]
72pub struct ExpandingClassification<const USER: bool=false> { 
73    pub classification: String,
74    pub __access_lvl__: i32,
75    pub __access_req__: Vec<String>,
76    pub __access_grp1__: Vec<String>,
77    pub __access_grp2__: Vec<String>,
78}
79
80impl<const USER: bool> ExpandingClassification<USER> {
81    pub fn new(classification: String, parser: &ClassificationParser) -> Result<Self, ModelError> {
82        if parser.original_definition.enforce {
83
84            let parts = parser.get_classification_parts(&classification, false, true, !USER)?;
85            let classification = parser.get_normalized_classification_text(parts.clone(), false, false)?;
86
87            Ok(Self {
88                classification,
89                __access_lvl__: parts.level,
90                __access_req__: parts.required,
91                __access_grp1__: if parts.groups.is_empty() { vec!["__EMPTY__".to_owned()] } else { parts.groups },
92                __access_grp2__: if parts.subgroups.is_empty() { vec!["__EMPTY__".to_owned()] } else { parts.subgroups },
93            })
94        } else {
95            Ok(Self {
96                classification,
97                __access_lvl__: Default::default(),
98                __access_req__: Default::default(),
99                __access_grp1__: vec!["__EMPTY__".to_owned()],
100                __access_grp2__: vec!["__EMPTY__".to_owned()],
101            })
102        }
103    }
104
105    pub fn try_unrestricted() -> Option<Self> {
106        if let ClassificationMode::Configured(parser) = &*GLOBAL_CLASSIFICATION.lock().unwrap() {
107            Some(Self::unrestricted(parser))
108        } else {
109            None
110        }   
111    }
112
113    pub fn unrestricted(parser: &ClassificationParser) -> Self {
114        Self::new(parser.unrestricted().to_string(), parser).unwrap()
115    }
116
117    pub fn as_str(&self) -> &str {
118        &self.classification
119    }
120
121    pub fn insert(parser: &ClassificationParser, output: &mut JsonMap, classification: &str) -> Result<(), ModelError> {
122        use serde_json::json;
123        if parser.original_definition.enforce {
124            let parts = parser.get_classification_parts(classification, true, true, !USER)?;
125            let classification = parser.get_normalized_classification_text(parts.clone(), true, false)?;
126
127            output.insert("classification".to_string(), json!(classification));
128            output.insert("__access_lvl__".to_string(), json!(parts.level));
129            output.insert("__access_req__".to_string(), json!(parts.required));
130            output.insert("__access_grp1__".to_string(), json!(if parts.groups.is_empty() { vec!["__EMPTY__".to_string()] } else { parts.groups }));
131            output.insert("__access_grp2__".to_string(), json!(if parts.subgroups.is_empty() { vec!["__EMPTY__".to_string()] } else { parts.subgroups }));
132            Ok(())
133        } else {
134            output.insert("classification".to_string(), json!(classification));
135            output.insert("__access_lvl__".to_string(), json!(0));
136            output.insert("__access_req__".to_string(), serde_json::Value::Array(Default::default()));
137            output.insert("__access_grp1__".to_string(), json!(&["__EMPTY__"]));
138            output.insert("__access_grp2__".to_string(), json!(&["__EMPTY__"]));
139            Ok(())
140        }
141    }
142}
143
144#[derive(Serialize, Deserialize)]
145struct RawClassification { 
146    pub classification: String,
147    #[serde(default)]
148    pub __access_lvl__: i32,
149    #[serde(default)]
150    pub __access_req__: Vec<String>,
151    #[serde(default)]
152    pub __access_grp1__: Vec<String>,
153    #[serde(default)]
154    pub __access_grp2__: Vec<String>,
155}
156
157impl<const U: bool> From<&ExpandingClassification<U>> for RawClassification {
158    fn from(value: &ExpandingClassification<U>) -> Self {
159        Self {
160            classification: value.classification.clone(),
161            __access_lvl__: value.__access_lvl__,
162            __access_req__: value.__access_req__.clone(),
163            __access_grp1__: value.__access_grp1__.clone(),
164            __access_grp2__: value.__access_grp2__.clone(),
165        }
166    }
167}
168
169impl<const U: bool> From<RawClassification> for ExpandingClassification<U> {
170    fn from(value: RawClassification) -> Self {
171        Self {
172            classification: value.classification.clone(),
173            __access_lvl__: value.__access_lvl__,
174            __access_req__: value.__access_req__.clone(),
175            __access_grp1__: value.__access_grp1__.clone(),
176            __access_grp2__: value.__access_grp2__.clone(),
177        }
178    }
179}
180
181
182impl<const U: bool> Serialize for ExpandingClassification<U> {
183    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
184    where S: serde::Serializer 
185    {
186        match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
187            ClassificationMode::Uninitialized => return Err(serde::ser::Error::custom("classification engine not initalized")),
188            // if it is disabled just serialize as raw
189            ClassificationMode::Disabled |
190            // assume the content of the struct is already normalized on load or construction
191            ClassificationMode::Configured(_) => {},
192        }
193        RawClassification::from(self).serialize(serializer)
194    }
195}
196
197impl<'de, const U: bool> Deserialize<'de> for ExpandingClassification<U> {
198    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
199    where D: serde::Deserializer<'de> 
200    {
201        let parser = match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
202            ClassificationMode::Uninitialized => return Err(serde::de::Error::custom("classification engine not initalized")),
203            // if it is disabled just handle as raw
204            ClassificationMode::Disabled => None,
205            // normalize on load
206            ClassificationMode::Configured(parser) => Some(parser.clone()),
207        };
208
209        if let Some(parser) = parser {
210            let raw = RawClassification::deserialize(deserializer)?;
211            Self::new(raw.classification, &parser).map_err(serde::de::Error::custom)
212        } else {
213            Ok(RawClassification::deserialize(deserializer)?.into())
214        }        
215    }
216}
217
218impl<const USER: bool> Described<ElasticMeta> for ExpandingClassification<USER> {
219    fn metadata() -> struct_metadata::Descriptor<ElasticMeta> {
220
221        // let group_meta = ElasticMeta {
222        //     index: Some(true),
223        //     store: None,
224        //     ..Default::default()
225        // };
226
227        struct_metadata::Descriptor { 
228            docs: None, 
229            metadata: ElasticMeta{mapping: Some("classification"), ..Default::default()}, 
230            kind: struct_metadata::Kind::new_struct("ExpandingClassification", vec![
231                struct_metadata::Entry { label: "classification", docs: None, metadata: ElasticMeta{mapping: Some("classification"), ..Default::default()}, type_info: String::metadata(), has_default: false, aliases: &["classification"] },
232                // struct_metadata::Entry { label: "__access_lvl__", docs: None, metadata: Default::default(), type_info: i32::metadata(), has_default: false, aliases: &["__access_lvl__"] },
233                // struct_metadata::Entry { label: "__access_req__", docs: None, metadata: Default::default(), type_info: Vec::<String>::metadata(), has_default: false, aliases: &["__access_req__"] },
234                // struct_metadata::Entry { label: "__access_grp1__", docs: None, metadata: Default::default(), type_info: Vec::<String>::metadata(), has_default: false, aliases: &["__access_grp1__"] },
235                // struct_metadata::Entry { label: "__access_grp2__", docs: None, metadata: Default::default(), type_info: Vec::<String>::metadata(), has_default: false, aliases: &["__access_grp2__"] },
236            ], &mut [], &mut[]),
237            // kind: struct_metadata::Kind::Aliased { 
238            //     name: "ExpandingClassification", 
239            //     kind: Box::new(String::metadata())
240            // },
241        }
242    }
243}
244
245impl<const U: bool> PartialEq<&str> for ExpandingClassification<U> {
246    fn eq(&self, other: &&str) -> bool {
247        let parser = match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
248            ClassificationMode::Uninitialized | ClassificationMode::Disabled => None,
249            ClassificationMode::Configured(parser) => Some(parser.clone()),
250        };
251
252        if let Some(parser) = parser {
253            match (parser.normalize_classification(&self.classification), parser.normalize_classification(other)) {
254                (Ok(a), Ok(b)) => a == b,
255                _ => false
256            }
257        } else {
258            self.classification.as_str() == *other
259        }
260    }
261}
262
263impl<const U: bool> PartialEq for ExpandingClassification<U> {
264    fn eq(&self, other: &Self) -> bool {
265        self.eq(&other.classification.as_str())
266    }
267}
268
269
270// MARK: ClassificationString
271
272/// A classification value stored as a string
273#[derive(Described, Debug, Clone, Eq)]
274#[metadata_type(ElasticMeta)]
275#[metadata(mapping="classification_string")]
276pub struct ClassificationString(pub (crate) String);
277
278impl From<ClassificationString> for String {
279    fn from(value: ClassificationString) -> Self {
280        value.0
281    }
282}
283
284impl From<ExpandingClassification> for ClassificationString {
285    fn from(value: ExpandingClassification) -> Self {
286        ClassificationString(value.classification)
287    }
288}
289
290impl Deref for ClassificationString {
291    type Target = str;
292
293    fn deref(&self) -> &Self::Target {
294        &self.0
295    }
296}
297
298impl ClassificationString {
299    pub fn new(classification: String, parser: &Arc<ClassificationParser>) -> Result<Self, ModelError> {
300        Ok(Self(parser.normalize_classification_options(&classification, NormalizeOptions::short())?))
301    }
302
303    pub fn new_unchecked(classification: String) -> Self {
304        Self(classification)
305    }
306
307    pub fn try_unrestricted() -> Option<Self> {
308        if let ClassificationMode::Configured(parser) = &*GLOBAL_CLASSIFICATION.lock().unwrap() {
309            Some(Self::unrestricted(parser))
310        } else {
311            None
312        }   
313    }
314    
315    pub fn unrestricted(parser: &ClassificationParser) -> Self {
316        Self(parser.normalize_classification_options(parser.unrestricted(), NormalizeOptions::short()).unwrap())
317    }
318
319    pub fn as_str(&self) -> &str {
320        &self.0
321    }
322
323    pub fn default_unrestricted() -> ClassificationString {
324        match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
325            ClassificationMode::Uninitialized => panic!("classification handling without defining parser"),
326            ClassificationMode::Disabled => ClassificationString(Default::default()),
327            ClassificationMode::Configured(parser) => ClassificationString::unrestricted(parser),
328        }
329    }
330}
331
332#[derive(Serialize, Deserialize)]
333struct RawClassificationString(String);
334
335
336impl From<&ClassificationString> for RawClassificationString {
337    fn from(value: &ClassificationString) -> Self {
338        Self(value.0.clone())
339    }
340}
341
342impl From<RawClassificationString> for ClassificationString {
343    fn from(value: RawClassificationString) -> Self {
344        Self(value.0.clone())
345    }
346}
347
348
349impl Serialize for ClassificationString {
350    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
351    where S: serde::Serializer 
352    {
353        match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
354            ClassificationMode::Uninitialized => return Err(serde::ser::Error::custom("classification engine not initalized")),
355            // if it is disabled just serialize as raw
356            ClassificationMode::Disabled |
357            // assume the content of the struct is already normalized on load or construction
358            ClassificationMode::Configured(_) => {},
359        }
360        RawClassificationString::from(self).serialize(serializer)
361    }
362}
363
364// fn response_or_empty_string<'de, D>(d: D) -> Result<Option<Response>, D::Error>
365// where
366//     D: Deserializer<'de>,
367// {
368//     #[derive(Deserialize)]
369//     #[serde(untagged)]
370//     enum MyHelper<'a> {
371//         S(&'a str),
372//         R(Response),
373//     }
374
375//     match MyHelper::deserialize(d) {
376//         Ok(MyHelper::R(r)) => Ok(Some(r)),
377//         Ok(MyHelper::S(s)) if s.is_empty() => Ok(None),
378//         Ok(MyHelper::S(_)) => Err(D::Error::custom("only empty strings may be provided")),
379//         Err(err) => Err(err),
380//     }
381// }
382
383impl<'de> Deserialize<'de> for ClassificationString {
384    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
385    where D: serde::Deserializer<'de> 
386    {
387        let parser = match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
388            ClassificationMode::Uninitialized => return Err(serde::de::Error::custom("classification engine not initalized")),
389            // if it is disabled just handle as raw
390            ClassificationMode::Disabled => None,
391            // normalize on load
392            ClassificationMode::Configured(parser) => Some(parser.clone()),
393        };
394
395        if let Some(parser) = parser {
396            let raw = RawClassificationString::deserialize(deserializer)?;
397            Self::new(raw.0, &parser).map_err(serde::de::Error::custom)
398        } else {
399            Ok(RawClassificationString::deserialize(deserializer)?.into())
400        }
401    }
402}
403
404impl PartialEq for ClassificationString {
405    fn eq(&self, other: &Self) -> bool {
406        let parser = match &*GLOBAL_CLASSIFICATION.lock().unwrap() {
407            ClassificationMode::Uninitialized | ClassificationMode::Disabled => None,
408            ClassificationMode::Configured(parser) => Some(parser.clone()),
409        };
410
411        if let Some(parser) = parser {
412            match (parser.normalize_classification(&self.0), parser.normalize_classification(&other.0)) {
413                (Ok(a), Ok(b)) => a == b,
414                _ => false
415            }
416        } else {
417            self.0 == other.0
418        }
419    }
420}