rustsec_admin/
assigner.rs1use 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#[derive(PartialEq, Eq, Clone, Copy)]
17pub enum OutputMode {
18 HumanReadable,
20 GithubAction,
22}
23
24pub 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 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
95fn 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}