Skip to main content

assemblyline_models/types/
classification.rs

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