rustsec_admin/
assigner.rs

1//! RustSec Advisory DB tool to assign ids
2
3use crate::{error::ErrorKind, prelude::*, Map};
4use rustsec::{
5    advisory::{IdKind, Parts},
6    Advisory, Collection,
7};
8use std::{
9    fs::{self, File},
10    io::{BufRead, BufReader, LineWriter, Write},
11    path::Path,
12    process::exit,
13};
14
15/// What sort of output should be generated on stdout.
16#[derive(PartialEq, Eq, Clone, Copy)]
17pub enum OutputMode {
18    /// Normal human readable logging
19    HumanReadable,
20    /// Output designed for use in the github action that runs this in prod
21    GithubAction,
22}
23
24/// assign ids to advisories in a particular repo_path
25pub fn assign_ids(repo_path: &Path, output_mode: OutputMode) {
26    let db = rustsec::Database::open(repo_path).unwrap_or_else(|e| {
27        status_err!(
28            "couldn't open advisory DB repo from {}: {}",
29            repo_path.display(),
30            e
31        );
32        exit(1);
33    });
34
35    let advisories = db.iter();
36
37    // Ensure we're parsing some advisories
38    if advisories.len() == 0 {
39        status_err!("no advisories found!");
40        exit(1);
41    }
42
43    if output_mode == OutputMode::HumanReadable {
44        status_ok!(
45            "Loaded",
46            "{} security advisories (from {})",
47            advisories.len(),
48            repo_path.display()
49        );
50    }
51
52    let mut highest_id = Map::new();
53
54    for advisory in advisories {
55        let advisory_clone = advisory.clone();
56        let metadata = advisory_clone.metadata;
57        let id = metadata.id;
58        let year = metadata.date.year();
59
60        if let IdKind::RustSec = id.kind() {
61            let id_num = id.numerical_part().unwrap();
62
63            if let Some(&number) = highest_id.get(&year) {
64                if number < id_num {
65                    highest_id.insert(year, id_num);
66                }
67            } else {
68                highest_id.insert(year, id_num);
69            }
70        }
71    }
72
73    let mut collection_strs = vec![];
74    let crates_str = Collection::Crates.to_string();
75    let rust_str = Collection::Rust.to_string();
76    collection_strs.push(crates_str);
77    collection_strs.push(rust_str);
78
79    let mut assignments = vec![];
80    for collection_str in collection_strs {
81        assign_ids_across_directory(
82            collection_str,
83            repo_path,
84            &mut highest_id,
85            output_mode,
86            &mut assignments,
87        );
88    }
89
90    if output_mode == OutputMode::GithubAction {
91        println!("Assigned {}", assignments.join(", "));
92    }
93}
94
95///Assign ids to files with placeholder IDs within the directory defined by dir_path
96fn assign_ids_across_directory(
97    collection_str: String,
98    repo_path: &Path,
99    highest_ids: &mut Map<u32, u32>,
100    output_mode: OutputMode,
101    assignments: &mut Vec<String>,
102) {
103    let dir_path = repo_path.join(collection_str);
104
105    if let Ok(collection_entry) = fs::read_dir(dir_path) {
106        for dir_entry in collection_entry {
107            let unwrapped_dir_entry = dir_entry.unwrap();
108            let dir_name = unwrapped_dir_entry.file_name().into_string().unwrap();
109            let dir_path = unwrapped_dir_entry.path();
110            let dir_path_clone = dir_path.clone();
111            for advisory_entry in fs::read_dir(dir_path).unwrap() {
112                let unwrapped_advisory = advisory_entry.unwrap();
113                let advisory_path = unwrapped_advisory.path();
114                let advisory_path_clone = advisory_path.clone();
115                let advisory_path_for_reading = advisory_path.clone();
116                let advisory_path_for_deleting = advisory_path.clone();
117                let displayed_advisory_path = advisory_path.display();
118                let advisory_filename = unwrapped_advisory.file_name();
119                let advisory_filename_str = advisory_filename.into_string().unwrap();
120                if advisory_filename_str.contains("RUSTSEC-0000-0000") {
121                    let advisory_data = fs::read_to_string(advisory_path_clone)
122                        .map_err(|e| {
123                            format_err!(
124                                ErrorKind::Io,
125                                "Couldn't open {}: {}",
126                                displayed_advisory_path,
127                                e
128                            );
129                        })
130                        .unwrap();
131
132                    let advisory_parts = Parts::parse(&advisory_data).unwrap();
133                    let advisory: Advisory = toml::from_str(advisory_parts.front_matter).unwrap();
134                    let date = advisory.metadata.date;
135                    let year = date.year();
136                    let new_id = highest_ids.get(&year).cloned().unwrap_or_default() + 1;
137                    let year_str = year.to_string();
138                    let string_id = format!("RUSTSEC-{}-{:04}", year_str, new_id);
139                    let new_filename = format!("{}.md", string_id);
140                    let new_path = dir_path_clone.join(new_filename);
141                    let original_file = File::open(advisory_path_for_reading).unwrap();
142                    let reader = BufReader::new(original_file);
143                    let new_file = File::create(new_path).unwrap();
144                    let mut writer = LineWriter::new(new_file);
145                    for line in reader.lines() {
146                        let current_line = line.unwrap();
147                        if current_line.contains("id = ") {
148                            writer
149                                .write_all(format!("id = \"{}\"\n", string_id).as_ref())
150                                .unwrap();
151                        } else {
152                            let current_line_with_newline = format!("{}\n", current_line);
153                            writer
154                                .write_all(current_line_with_newline.as_ref())
155                                .unwrap();
156                        }
157                    }
158                    highest_ids.insert(year, new_id);
159                    fs::remove_file(advisory_path_for_deleting).unwrap();
160                    if output_mode == OutputMode::HumanReadable {
161                        status_ok!("Assignment", "Assigned {} to {}", string_id, dir_name);
162                    } else {
163                        assignments.push(format!("{} to {}", string_id, dir_name))
164                    }
165                }
166            }
167        }
168    }
169}