|
| 1 | +use cached::proc_macro::cached; |
| 2 | +use cached::SizedCache; |
| 3 | +use chrono::{NaiveDateTime, Utc}; |
| 4 | +use cve::v4::{CVEContainer, CVEItem}; |
| 5 | +use diesel::mysql::MysqlConnection; |
| 6 | +use helper::init_db_pool; |
| 7 | +use nvd_api::v2::vulnerabilities::CveParameters; |
| 8 | +use nvd_api::v2::LastModDate; |
| 9 | +use nvd_api::ApiVersion; |
| 10 | +use nvd_server::error::DBResult; |
| 11 | +use nvd_server::modules::cve_db::CreateCve; |
| 12 | +use nvd_server::modules::cve_product_db::CreateCveProductByName; |
| 13 | +use nvd_server::modules::product_db::{CreateProduct, QueryProductById}; |
| 14 | +use nvd_server::modules::vendor_db::CreateVendors; |
| 15 | +use nvd_server::modules::{Cve, CveProduct, Product, Vendor}; |
| 16 | +use std::str::FromStr; |
| 17 | + |
| 18 | +// https://cwe.mitre.org/data/downloads.html |
| 19 | +// curl -s -k https://cwe.mitre.org/data/downloads.html |grep -Eo '(/[^"]*\.xml.zip)'|xargs -I % wget -c https://cwe.mitre.org% |
| 20 | + |
| 21 | +fn import_to_db(connection: &mut MysqlConnection, cve_item: CVEItem) -> DBResult<String> { |
| 22 | + let id = cve_item.cve.meta.id; |
| 23 | + let y = id.split('-').nth(1).unwrap_or_default(); |
| 24 | + let new_post = CreateCve { |
| 25 | + id: id.clone(), |
| 26 | + created_at: cve_item.published_date, |
| 27 | + updated_at: cve_item.last_modified_date, |
| 28 | + references: serde_json::json!(cve_item.cve.references.reference_data), |
| 29 | + description: serde_json::json!(cve_item.cve.description.description_data), |
| 30 | + severity: cve_item.impact.severity(), |
| 31 | + metrics: serde_json::json!(cve_item.impact), |
| 32 | + assigner: cve_item.cve.meta.assigner, |
| 33 | + configurations: serde_json::json!(cve_item.configurations.nodes), |
| 34 | + year: i32::from_str(y).unwrap_or_default(), |
| 35 | + weaknesses: serde_json::json!(cve_item.cve.problem_type.problem_type_data), |
| 36 | + timeline: Default::default(), |
| 37 | + }; |
| 38 | + // 插入到数据库 |
| 39 | + match Cve::create(connection, &new_post) { |
| 40 | + Ok(cve_id) => { |
| 41 | + // 插入cpe_match关系表 |
| 42 | + for node in cve_item.configurations.nodes { |
| 43 | + for vendor_product in node.vendor_product() { |
| 44 | + import_vendor_product_to_db(connection, vendor_product.clone()); |
| 45 | + create_cve_product( |
| 46 | + connection, |
| 47 | + cve_id.id.clone(), |
| 48 | + vendor_product.vendor, |
| 49 | + vendor_product.product, |
| 50 | + ); |
| 51 | + } |
| 52 | + } |
| 53 | + } |
| 54 | + Err(err) => { |
| 55 | + println!("Cve::create: {err:?}"); |
| 56 | + } |
| 57 | + } |
| 58 | + Ok(new_post.id) |
| 59 | +} |
| 60 | + |
| 61 | +pub fn create_cve_product( |
| 62 | + conn: &mut MysqlConnection, |
| 63 | + cve_id: String, |
| 64 | + vendor: String, |
| 65 | + product: String, |
| 66 | +) -> String { |
| 67 | + // 构建待插入对象 |
| 68 | + let cp = CreateCveProductByName { |
| 69 | + cve_id, |
| 70 | + vendor, |
| 71 | + product, |
| 72 | + }; |
| 73 | + // 插入到数据库 |
| 74 | + match CveProduct::create_by_name(conn, &cp) { |
| 75 | + Ok(_cp) => {} |
| 76 | + Err(err) => { |
| 77 | + println!("create_cve_product: {err:?}:{cp:?}"); |
| 78 | + } |
| 79 | + } |
| 80 | + String::new() |
| 81 | +} |
| 82 | + |
| 83 | +#[cached( |
| 84 | + type = "SizedCache<String, Vec<u8>>", |
| 85 | + create = "{ SizedCache::with_size(100) }", |
| 86 | + convert = r#"{ format!("{:?}", product.to_owned()) }"# |
| 87 | +)] |
| 88 | +fn import_vendor_product_to_db(connection: &mut MysqlConnection, product: cpe::Product) -> Vec<u8> { |
| 89 | + let vendor_id = create_vendor(connection, product.vendor, None); |
| 90 | + create_product(connection, vendor_id, product.product, product.part) |
| 91 | +} |
| 92 | + |
| 93 | +#[cached( |
| 94 | + type = "SizedCache<String, Vec<u8>>", |
| 95 | + create = "{ SizedCache::with_size(100) }", |
| 96 | + convert = r#"{ format!("{}", name.to_owned()) }"# |
| 97 | +)] |
| 98 | +pub fn create_vendor( |
| 99 | + conn: &mut MysqlConnection, |
| 100 | + name: String, |
| 101 | + description: Option<String>, |
| 102 | +) -> Vec<u8> { |
| 103 | + if let Ok(v) = Vendor::query_by_name(conn, &name) { |
| 104 | + return v.id; |
| 105 | + } |
| 106 | + // 构建待插入对象 |
| 107 | + let new_post = CreateVendors { |
| 108 | + id: uuid::Uuid::new_v4().as_bytes().to_vec(), |
| 109 | + name, |
| 110 | + description, |
| 111 | + official: u8::from(true), |
| 112 | + homepage: None, |
| 113 | + }; |
| 114 | + // 插入到数据库 |
| 115 | + if let Err(err) = Vendor::create(conn, &new_post) { |
| 116 | + println!("create_vendor: {err:?}"); |
| 117 | + } |
| 118 | + new_post.id |
| 119 | +} |
| 120 | + |
| 121 | +#[cached( |
| 122 | + type = "SizedCache<String, Vec<u8>>", |
| 123 | + create = "{ SizedCache::with_size(100) }", |
| 124 | + convert = r#"{ format!("{}:{:?}", name.to_owned(),vendor.to_owned()) }"# |
| 125 | +)] |
| 126 | +pub fn create_product( |
| 127 | + conn: &mut MysqlConnection, |
| 128 | + vendor: Vec<u8>, |
| 129 | + name: String, |
| 130 | + part: String, |
| 131 | +) -> Vec<u8> { |
| 132 | + let q = QueryProductById { |
| 133 | + vendor_id: vendor.clone(), |
| 134 | + name: name.clone(), |
| 135 | + }; |
| 136 | + if let Ok(v) = Product::query_by_id(conn, &q) { |
| 137 | + return v.id; |
| 138 | + } |
| 139 | + // 构建待插入对象 |
| 140 | + let new_post = CreateProduct { |
| 141 | + id: uuid::Uuid::new_v4().as_bytes().to_vec(), |
| 142 | + vendor_id: vendor, |
| 143 | + name, |
| 144 | + description: None, |
| 145 | + official: u8::from(true), |
| 146 | + part, |
| 147 | + homepage: None, |
| 148 | + }; |
| 149 | + // 插入到数据库 |
| 150 | + if let Err(err) = Product::create(conn, &new_post) { |
| 151 | + println!("create_product: {err:?}"); |
| 152 | + } |
| 153 | + new_post.id |
| 154 | +} |
| 155 | + |
| 156 | +fn main() { |
| 157 | + // let connection_pool = init_db_pool(); |
| 158 | + // let api = nvd_api::NVDApi::new(None, ApiVersion::default()).unwrap(); |
| 159 | + let now = Utc::now(); |
| 160 | + let two_h = |
| 161 | + println!("{:?}", now); |
| 162 | + // api.cve(CveParameters{ |
| 163 | + // cpe_name: None, |
| 164 | + // cve_id: None, |
| 165 | + // cvss_v2_metrics: None, |
| 166 | + // cvss_v2_severity: None, |
| 167 | + // cvss_v3_metrics: None, |
| 168 | + // cvss_v3_severity: None, |
| 169 | + // cwe_id: None, |
| 170 | + // has_cert_alerts: None, |
| 171 | + // has_cert_notes: None, |
| 172 | + // has_kev: None, |
| 173 | + // has_oval: None, |
| 174 | + // is_vulnerable: None, |
| 175 | + // keyword: None, |
| 176 | + // last_mod: Some(LastModDate{ last_mod_start_date: "".to_string(), last_mod_end_date: "".to_string() }), |
| 177 | + // no_rejected: None, |
| 178 | + // pub_date: None, |
| 179 | + // limit_offset: None, |
| 180 | + // source_identifier: None, |
| 181 | + // virtual_match: None, |
| 182 | + // }) |
| 183 | +} |
0 commit comments