cwe_xml/cwe/
mod.rs

1use std::collections::{HashMap, HashSet};
2use std::collections::hash_map::Entry::{Occupied, Vacant};
3use std::fmt::{Display, Formatter};
4use std::fs::File;
5use std::io::{BufRead, BufReader, Read};
6use std::process::id;
7use std::rc::Rc;
8
9use crate::cwe::categories::Category;
10use crate::cwe::weakness_catalog::WeaknessCatalog;
11use crate::cwe::weaknesses::{RelatedNature, Weakness};
12use crate::errors::Error;
13
14pub mod views;
15pub mod external_references;
16pub mod weaknesses;
17pub mod categories;
18pub mod relationships;
19pub mod notes;
20pub mod content_history;
21pub mod structured_text;
22pub mod weakness_catalog;
23
24/// A CWE weakness database.
25#[derive(Debug, Default)]
26pub struct CweDatabase {
27    catalogs: HashMap<String, WeaknessCatalog>,
28    weakness_index: HashMap<i64, Rc<Weakness>>,
29    cwe_categories_index: HashMap<i64 /*cwe-id*/, HashSet<Rc<Category>>>,
30    weakness_children_index: HashMap<i64, HashSet<Rc<Weakness>>>,
31    weakness_roots_index: HashMap<i64, Rc<Weakness>>,
32}
33
34pub trait WeaknessVisitor {
35    fn visit(&mut self, db: &CweDatabase, level: usize, weakness: Rc<Weakness>);
36}
37
38impl CweDatabase {
39    /// Create a new empty CWE database.
40    pub fn new() -> CweDatabase {
41        CweDatabase::default()
42    }
43
44    /// Import a CWE catalog from a string containing the XML.
45    pub fn import_weakness_catalog_from_str(&mut self, xml: &str) -> Result<(), Error> {
46        let weakness_catalog: WeaknessCatalog = quick_xml::de::from_str(xml).map_err(|e| Error::InvalidCweFile {
47            file: "".to_string(),
48            error: e.to_string(),
49        })?;
50        let catalog_name = weakness_catalog.name.clone();
51        self.update_indexes(&weakness_catalog);
52        self.catalogs.insert(catalog_name, weakness_catalog);
53        Ok(())
54    }
55
56    /// Import a CWE catalog from an URL string containing the XML.
57    pub fn import_weakness_catalog_from_url(&mut self, url: &str) -> Result<(), Error> {
58        let xml_content = download_xml(url).map_err(|e| Error::InvalidCweFile {
59            file: url.to_string(),
60            error: e.to_string(),
61        })?;
62        self.import_weakness_catalog_from_str(&xml_content)
63    }
64
65    /// Import a CWE catalog from a file containing the XML.
66    pub fn import_weakness_catalog_from_file(&mut self, xml_file: &str) -> Result<(), Error> {
67        let file = File::open(xml_file).map_err(|e| Error::InvalidCweFile {
68            file: xml_file.to_string(),
69            error: e.to_string(),
70        })?;
71        let reader = BufReader::new(file);
72        let weakness_catalog: WeaknessCatalog = quick_xml::de::from_reader(reader).map_err(|e| Error::InvalidCweFile {
73            file: xml_file.to_string(),
74            error: e.to_string(),
75        })?;
76        let catalog_name = weakness_catalog.name.clone();
77        self.update_indexes(&weakness_catalog);
78        self.catalogs.insert(catalog_name, weakness_catalog);
79        Ok(())
80    }
81
82    /// Import a CWE catalog from a reader containing the XML.
83    pub fn import_weakness_catalog_from_reader<R>(&mut self, reader: R) -> Result<(), Error> where R: BufRead {
84        let weakness_catalog: WeaknessCatalog = quick_xml::de::from_reader(reader).map_err(|e| Error::InvalidCweFile {
85            file: "".to_string(),
86            error: e.to_string(),
87        })?;
88        let catalog_name = weakness_catalog.name.clone();
89        self.update_indexes(&weakness_catalog);
90        self.catalogs.insert(catalog_name, weakness_catalog);
91        Ok(())
92    }
93
94    /// Returns a reference to a Weakness struct if the CWE-ID exists in the catalog.
95    pub fn weakness_by_cwe_id(&self, cwe_id: i64) -> Option<Rc<Weakness>> {
96        self.weakness_index.get(&cwe_id)
97            .map(|weakness| weakness.clone())
98    }
99
100    /// Returns a list of categories for a given CWE-ID.
101    pub fn categories_by_cwe_id(&self, cwe_id: i64) -> HashSet<Rc<Category>> {
102        self.cwe_categories_index.get(&cwe_id)
103            .map(|categories| categories.clone())
104            .unwrap_or_default()
105    }
106
107    /// Returns a list of weaknesses that are children of a given CWE-ID.
108    pub fn weakness_children_by_cwe_id(&self, cwe_id: i64) -> HashSet<Rc<Weakness>> {
109        self.weakness_children_index.get(&cwe_id)
110            .map(|weaknesses| weaknesses.clone())
111            .unwrap_or_default()
112    }
113
114    /// Returns a list of weaknesses that are children of a given CWE-ID.
115    /// The list does not contain the weakness for the given CWE-ID.
116    pub fn weakness_subtree_by_cwe_id(&self, cwe_id: i64) -> Vec<Rc<Weakness>> {
117        let mut visitor = CweIdSubTreeVisitor::default();
118
119        if let Some(weakness) = self.weakness_by_cwe_id(cwe_id) {
120            self.visit_weakness(&mut visitor, 0, &weakness);
121        }
122
123        if visitor.cwe_ids.is_empty() {
124            vec![]
125        } else {
126            visitor.cwe_ids.iter().map(|cwe_id| self.weakness_by_cwe_id(*cwe_id).expect("should never happen")).collect()
127        }
128    }
129
130    /// Returns a list of weaknesses that are roots, i.e. they have no parents.
131    pub fn weakness_roots(&self) -> HashSet<Rc<Weakness>> {
132        self.weakness_roots_index.values().cloned().collect()
133    }
134
135    /// Visit all root weaknesses in the database and their children.
136    pub fn visit_weaknesses(&self, visitor: &mut impl WeaknessVisitor) {
137        for weakness in self.weakness_roots().iter() {
138            self.visit_weakness(visitor, 0, weakness);
139        }
140    }
141
142    /// Returns the direct weakness ancestors of a given CWE-ID.
143    pub fn direct_ancestors_by_cwe_id(&self, cwe_id: i64) -> HashSet<Rc<Weakness>> {
144        let mut ancestors = HashSet::new();
145        if let Some(weakness) = self.weakness_by_cwe_id(cwe_id) {
146            for ancestor_id in weakness.direct_ancestors() {
147                if let Some(ancestor) = self.weakness_by_cwe_id(ancestor_id) {
148                    ancestors.insert(ancestor);
149                }
150            }
151        }
152        ancestors
153    }
154
155    /// Merge the given categories into the category index for the given CWE-ID.
156    pub fn merge_categories_by_cwe_id(&mut self, cwe_id: i64, categories: HashSet<Rc<Category>>) {
157        self.cwe_categories_index.entry(cwe_id).or_insert_with(HashSet::new).extend(categories.iter().cloned());
158    }
159
160    /// Returns a list of all categories in the database (across all catalogs).
161    pub fn all_categories(&self) -> HashSet<Rc<Category>> {
162        let mut categories = HashSet::new();
163        for catalog in self.catalogs.values() {
164            if let Some(cats) = catalog.categories.as_ref() {
165                categories.extend(cats.categories.iter().cloned());
166            }
167        }
168        categories
169    }
170
171    /// Sub-weaknesses inherit the categories of their parent weaknesses.
172    pub fn infer_categories_from_ancestors(&mut self) {
173        let mut inferred_categories = HashMap::new();
174        let category_index = self.cwe_categories_index.clone();
175        for (cwe_id, categories) in category_index.iter() {
176            self.propagate_categories_to_subtree(*cwe_id, categories.clone(), &mut inferred_categories);
177        }
178
179        for (cwe_id, categories) in inferred_categories {
180            self.cwe_categories_index.entry(cwe_id).or_insert_with(HashSet::new).extend(categories.iter().cloned());
181        }
182    }
183
184    /// Propagate categories to ancestors that don't have any categories yet and have only one child
185    /// that has no category defined.
186    /// This is a heuristic to infer categories for weaknesses that have no categories defined.
187    /// This process is repeated until no more categories can be propagated.
188    pub fn infer_categories_from_descendants(&mut self) {
189        struct PropagatedCategories {
190            categories: HashSet<Rc<Category>>,
191            no_category_count: usize,
192        }
193
194        loop {
195            // Number of updates in the current round
196            let mut update_count = 0;
197            let mut propagated_categories = HashMap::new();
198
199            for (id, weakness) in self.weakness_index.iter() {
200                let categories = self.categories_by_cwe_id(*id);
201
202                let ancestors = self.direct_ancestors_by_cwe_id(*id);
203                for ancestor in ancestors.iter() {
204                    let ancestor_cats = self.categories_by_cwe_id(ancestor.id);
205                    if !ancestor_cats.is_empty() {
206                        // skip ancestors that already have categories
207                        continue;
208                    }
209                    let propagated_categories = propagated_categories.entry(ancestor.id).or_insert_with(|| PropagatedCategories {
210                        categories: HashSet::new(),
211                        no_category_count: 0,
212                    });
213                    if categories.is_empty() {
214                        propagated_categories.no_category_count += 1;
215                    } else {
216                        propagated_categories.categories.extend(categories.iter().cloned());
217                    }
218                }
219            }
220
221            for (cwe_id, categories) in propagated_categories {
222                if categories.no_category_count <= 1 {
223                    let mut existing_cats = self.cwe_categories_index.entry(cwe_id).or_insert_with(HashSet::new);
224                    let existing_cats_count = existing_cats.len();
225                    existing_cats.extend(categories.categories.iter().cloned());
226                    if existing_cats.len() > existing_cats_count {
227                        update_count += 1;
228                    }
229                }
230            }
231
232            if update_count == 0 {
233                break;
234            }
235        }
236    }
237
238    fn propagate_categories_to_subtree(&self, cwe_id: i64, categories: HashSet<Rc<Category>>, inferred_categories: &mut HashMap<i64, HashSet<Rc<Category>>>) {
239        inferred_categories.entry(cwe_id).or_insert_with(HashSet::new).extend(categories.iter().cloned());
240        let categories = inferred_categories.get(&cwe_id).expect("Should never happen").clone();
241
242        let children = self.weakness_children_by_cwe_id(cwe_id);
243        for child in children.iter() {
244            self.propagate_categories_to_subtree(child.id, categories.clone(), inferred_categories);
245        }
246    }
247
248    fn visit_weakness(&self, visitor: &mut impl WeaknessVisitor, level: usize, weakness: &Rc<Weakness>) {
249        visitor.visit(self, level, weakness.clone());
250        for child in self.weakness_children_by_cwe_id(weakness.id).iter() {
251            self.visit_weakness(visitor, level + 1, child);
252        }
253    }
254
255    fn update_indexes(&mut self, catalog: &WeaknessCatalog) {
256        self.update_category_index(catalog);
257        self.update_weakness_index(catalog);
258        self.update_weakness_children_index(catalog);
259    }
260
261    fn update_weakness_index(&mut self, catalog: &WeaknessCatalog) {
262        if let Some(catalog) = &catalog.weaknesses {
263            for weakness in catalog.weaknesses.iter() {
264                self.weakness_index.insert(weakness.id, weakness.clone());
265            }
266        }
267    }
268
269    fn update_weakness_children_index(&mut self, catalog: &WeaknessCatalog) {
270        if let Some(catalog) = &catalog.weaknesses {
271            for weakness in catalog.weaknesses.iter() {
272                let mut parent_count = 0;
273                if let Some(related_weaknesses) = &weakness.related_weaknesses {
274                    for related_weakness in &related_weaknesses.related_weaknesses {
275                        if related_weakness.nature == RelatedNature::ChildOf {
276                            self.weakness_children_index.entry(related_weakness.cwe_id).or_insert_with(HashSet::new).insert(weakness.clone());
277                            parent_count += 1;
278                        }
279                    }
280                }
281
282                if parent_count == 0 {
283                    self.weakness_roots_index.insert(weakness.id, weakness.clone());
284                }
285            }
286        }
287    }
288
289    fn update_category_index(&mut self, catalog: &WeaknessCatalog) {
290        if let Some(categories) = &catalog.categories {
291            for category in categories.categories.iter() {
292                for member in &category.relationships.has_members {
293                    self.cwe_categories_index.entry(member.cwe_id).or_insert_with(HashSet::new).insert(category.clone());
294                }
295            }
296        }
297    }
298}
299
300impl Display for CweDatabase {
301    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
302        let mut text = String::new();
303        for (catalog_name, catalog) in &self.catalogs {
304            text.push_str(&format!("Catalog: {}\n", catalog_name));
305            text.push_str(&format!("  Version: {}\n", catalog.version));
306            text.push_str(&format!("  Date: {}\n", catalog.date));
307            text.push_str(&format!("  #Weaknesses: {}\n", catalog.weaknesses.as_ref().unwrap().weaknesses.len()));
308        }
309        write!(f, "{}", text)
310    }
311}
312
313#[derive(Default)]
314struct CweIdSubTreeVisitor {
315    cwe_ids: HashSet<i64>,
316}
317
318impl WeaknessVisitor for CweIdSubTreeVisitor {
319    fn visit(&mut self, _: &CweDatabase, level: usize, weakness: Rc<Weakness>) {
320        if level > 0 {
321            self.cwe_ids.insert(weakness.id);
322        }
323    }
324}
325
326fn download_xml(file: &str) -> Result<String, Box<dyn std::error::Error>> {
327    let mut tmp_file = tempfile::tempfile()?;
328    let mut xml = String::new();
329
330    reqwest::blocking::get(file)?
331        .copy_to(&mut tmp_file)?;
332    let mut zip_archive = zip::ZipArchive::new(tmp_file)?;
333    zip_archive
334        .by_index(0)?
335        .read_to_string(&mut xml)?;
336    Ok(xml)
337}
cwe_xml/cwe/mod.rs

cwe_xml/cwe/
mod.rs
