diff --git a/Cargo.lock b/Cargo.lock index 1ab3ee9..31d4004 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -139,6 +139,12 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "arc-swap" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457" + [[package]] name = "arrayvec" version = "0.7.4" @@ -1300,8 +1306,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -1597,13 +1605,28 @@ dependencies = [ "http 1.1.0", "hyper 1.3.1", "hyper-util", + "log", "rustls 0.22.4", + "rustls-native-certs", "rustls-pki-types", "tokio", "tokio-rustls", "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3203a961e5c83b6f5498933e78b6b263e208c197b63e9c6c53cc82ffd3f63793" +dependencies = [ + "hyper 1.3.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.5.0" @@ -1755,6 +1778,16 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3" +[[package]] +name = "iri-string" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f5f6c2df22c009ac44f6f1499308e7a3ac7ba42cd2378475cc691510e1eef1b" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.0" @@ -1794,6 +1827,21 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "9.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ae10193d25051e74945f1ea2d0b42e03cc3b890f7e4cc5faa44997d808193f" +dependencies = [ + "base64 0.21.7", + "js-sys", + "pem", + "ring", + "serde", + "serde_json", + "simple_asn1", +] + [[package]] name = "kv-log-macro" version = "1.0.7" @@ -2103,6 +2151,45 @@ dependencies = [ "memchr", ] +[[package]] +name = "octocrab" +version = "0.38.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68a8a3df00728324ad654ecd1ed449a60157c55b7ff8c109af3a35989687c367" +dependencies = [ + "arc-swap", + "async-trait", + "base64 0.22.1", + "bytes", + "cfg-if", + "chrono", + "either", + "futures", + "futures-util", + "http 1.1.0", + "http-body 1.0.0", + "http-body-util", + "hyper 1.3.1", + "hyper-rustls", + "hyper-timeout", + "hyper-util", + "jsonwebtoken", + "once_cell", + "percent-encoding", + "pin-project", + "secrecy", + "serde", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "snafu", + "tokio", + "tower", + "tower-http", + "tracing", + "url", +] + [[package]] name = "once_cell" version = "1.19.0" @@ -2236,6 +2323,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pem" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e459365e590736a54c3fa561947c84837534b8e9af6fc5bf781307e82658fae" +dependencies = [ + "base64 0.22.1", + "serde", +] + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -2941,6 +3038,19 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f1fb85efa936c42c6d5fc28d2629bb51e4b2f4b8a5211e297d599cc5a093792" +dependencies = [ + "openssl-probe", + "rustls-pemfile 2.1.2", + "rustls-pki-types", + "schannel", + "security-framework", +] + [[package]] name = "rustls-pemfile" version = "1.0.4" @@ -3208,6 +3318,15 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" +[[package]] +name = "secrecy" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bd1c54ea06cfd2f6b63219704de0b9b4f72dcc2b8fdef820be6cd799780e91e" +dependencies = [ + "zeroize", +] + [[package]] name = "security-framework" version = "2.11.0" @@ -3287,6 +3406,16 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af99884400da37c88f5e9146b7f1fd0fbcae8f6eec4e9da38b67d05486f814a6" +dependencies = [ + "itoa", + "serde", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -3398,6 +3527,18 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f27f6278552951f1f2b8cf9da965d10969b2efdea95a6ec47987ab46edfe263a" +[[package]] +name = "simple_asn1" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc4e5204eb1910f40f9cfa375f6f05b68c3abac4b6fd879c8ff5e7ae8a0a085" +dependencies = [ + "num-bigint", + "num-traits", + "thiserror", + "time", +] + [[package]] name = "siphasher" version = "0.3.11" @@ -3429,6 +3570,27 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +[[package]] +name = "snafu" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418b8136fec49956eba89be7da2847ec1909df92a9ae4178b5ff0ff092c8d95e" +dependencies = [ + "snafu-derive", +] + +[[package]] +name = "snafu-derive" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a4812a669da00d17d8266a0439eddcacbc88b17f732f927e52eeb9d196f7fb5" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.66", +] + [[package]] name = "socket2" version = "0.4.10" @@ -4161,8 +4323,30 @@ dependencies = [ "pin-project", "pin-project-lite", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5" +dependencies = [ + "bitflags 2.5.0", + "bytes", + "futures-util", + "http 1.1.0", + "http-body 1.0.0", + "http-body-util", + "iri-string", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "tracing", ] [[package]] @@ -4581,6 +4765,7 @@ dependencies = [ "hmac", "lazy_static", "migration", + "octocrab", "regex", "reqwest 0.12.4", "scraper", diff --git a/Cargo.toml b/Cargo.toml index f9677f9..7b79591 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -56,3 +56,4 @@ sha2 = "0.10.8" hmac = "0.12.1" base64 = "0.22.1" chrono-tz = "0.9.0" +octocrab = "0.38.0" diff --git a/config/development.yaml b/config/development.yaml index ed10c36..7dbabc5 100644 --- a/config/development.yaml +++ b/config/development.yaml @@ -42,3 +42,5 @@ ding_bot: lark_bot: access_token: {{ get_env(name="LARK_ACCESS_TOKEN", default="") }} secret_token: {{ get_env(name="LARK_SECRET_TOKEN", default="") }} + +github_search: true diff --git a/src/app.rs b/src/app.rs index ec3dfdb..3326278 100644 --- a/src/app.rs +++ b/src/app.rs @@ -18,6 +18,7 @@ use crate::{ msg_template::{reader_vulninfo, render_init}, BotManager, }, + search::search_github_poc, }; lazy_static! { @@ -130,8 +131,23 @@ impl WatchVulnApp { info!("{} has been pushed, skipped", vuln.key); continue; } + let key = vuln.key.clone(); let title = vuln.title.clone(); + if !vuln.cve.is_empty() && self.app_context.config.github_search { + let links = search_github_poc(&vuln.cve).await; + info!("{} found {} links from github", &vuln.cve, links.len()); + if let Err(err) = vuln_informations::Model::update_github_search_by_key( + &self.app_context.db, + &key, + links, + ) + .await + { + warn!("update vuln {} github_search error: {}", &vuln.cve, err); + } + todo!() + } let msg = match reader_vulninfo(vuln.into()) { Ok(msg) => msg, Err(err) => { diff --git a/src/config.rs b/src/config.rs index 5c4c13b..668c52f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -23,6 +23,7 @@ pub struct Config { pub tg_bot: TgBot, pub ding_bot: DingBot, pub lark_bot: LarkBot, + pub github_search: bool, } impl Config { diff --git a/src/error.rs b/src/error.rs index 52f29c9..1d8480a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -71,6 +71,9 @@ pub enum Error { #[error(transparent)] SystemTime(#[from] SystemTimeError), + #[error(transparent)] + Octocrab(#[from] octocrab::Error), + #[error(transparent)] Any(#[from] Box), diff --git a/src/grab/anti.rs b/src/grab/anti.rs index 38ea83e..311bf37 100644 --- a/src/grab/anti.rs +++ b/src/grab/anti.rs @@ -63,6 +63,7 @@ impl Grab for AntiCrawler { from, tags: vec![], reasons: vec![], + github_search: vec![], is_valuable: true, }; res.push(vuln); diff --git a/src/grab/avd.rs b/src/grab/avd.rs index 007f514..f250980 100644 --- a/src/grab/avd.rs +++ b/src/grab/avd.rs @@ -146,6 +146,7 @@ impl AVDCrawler { from: self.link.clone(), tags, reasons: vec![], + github_search: vec![], is_valuable, }; Ok(data) diff --git a/src/grab/kev.rs b/src/grab/kev.rs index 18fe46a..b91994c 100644 --- a/src/grab/kev.rs +++ b/src/grab/kev.rs @@ -57,6 +57,7 @@ impl Grab for KevCrawler { vuln.product.to_string(), "在野利用".to_string(), ], + github_search: vec![], reasons: vec![], is_valuable, }; diff --git a/src/grab/mod.rs b/src/grab/mod.rs index 9dc561e..fc649da 100644 --- a/src/grab/mod.rs +++ b/src/grab/mod.rs @@ -30,6 +30,7 @@ pub struct VulnInfo { pub from: String, pub tags: Vec, pub reasons: Vec, + pub github_search: Vec, pub is_valuable: bool, } @@ -58,6 +59,11 @@ impl From for VulnInfo { None => Vec::new(), }; + let github_search = match v.github_search { + Some(github_search) => github_search, + None => Vec::new(), + }; + VulnInfo { unique_key: v.key, title: v.title, @@ -70,6 +76,7 @@ impl From for VulnInfo { from: v.from, tags, reasons, + github_search, is_valuable: v.is_valuable, } } diff --git a/src/grab/oscs.rs b/src/grab/oscs.rs index 23dcddb..33a8833 100644 --- a/src/grab/oscs.rs +++ b/src/grab/oscs.rs @@ -146,6 +146,7 @@ impl OscCrawler { from: self.link.clone(), tags: vec![], reasons: vec![], + github_search: vec![], is_valuable, }; Ok(data) diff --git a/src/grab/seebug.rs b/src/grab/seebug.rs index e1eee92..1ad8fb9 100644 --- a/src/grab/seebug.rs +++ b/src/grab/seebug.rs @@ -143,6 +143,7 @@ impl SeeBugCrawler { from: href, tags, reasons: vec![], + github_search: vec![], is_valuable, }; res.push(data); diff --git a/src/grab/threatbook.rs b/src/grab/threatbook.rs index 4a44be8..97b4186 100644 --- a/src/grab/threatbook.rs +++ b/src/grab/threatbook.rs @@ -65,6 +65,7 @@ impl Grab for ThreadBookCrawler { from: LINK.to_string(), tags, reasons: Vec::new(), + github_search: vec![], is_valuable, }; res.push(vuln); diff --git a/src/grab/ti.rs b/src/grab/ti.rs index 97d742e..f48988f 100644 --- a/src/grab/ti.rs +++ b/src/grab/ti.rs @@ -75,6 +75,7 @@ impl TiCrawler { from: format!("https://ti.qianxin.com/vulnerability/detail/{}", detail.id), tags, reasons: vec![], + github_search: vec![], is_valuable, }; if vuln_infos diff --git a/src/lib.rs b/src/lib.rs index 155c1fa..ae4171b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,4 +7,5 @@ pub mod grab; pub mod logger; pub mod models; pub mod push; +pub mod search; pub mod utils; diff --git a/src/models/vuln_informations.rs b/src/models/vuln_informations.rs index 21648be..7945820 100644 --- a/src/models/vuln_informations.rs +++ b/src/models/vuln_informations.rs @@ -120,6 +120,29 @@ impl super::_entities::vuln_informations::Model { Ok(v) } + pub async fn update_github_search_by_key( + db: &DatabaseConnection, + key: &str, + links: Vec, + ) -> ModelResult<()> { + let txn = db.begin().await?; + let v = vuln_informations::Entity::find() + .filter(vuln_informations::Column::Key.eq(key)) + .one(&txn) + .await?; + if let Some(v) = v { + let mut v: vuln_informations::ActiveModel = v.into(); + v.github_search = ActiveValue::set(Some(links)); + v.update(&txn).await?; + txn.commit().await?; + Ok(()) + } else { + Err(ModelError::EntityUpdateNotFound { + key: key.to_string(), + }) + } + } + pub async fn update_pushed_by_key(db: &DatabaseConnection, key: String) -> ModelResult<()> { let txn = db.begin().await?; let v = vuln_informations::Entity::find() diff --git a/src/push/msg_template.rs b/src/push/msg_template.rs index 69a50cc..cab5758 100644 --- a/src/push/msg_template.rs +++ b/src/push/msg_template.rs @@ -21,7 +21,10 @@ const VULN_INFO_MSG_TEMPLATE: &str = r####" {% endif %} {% if references%}### **参考链接** {% for reference in references %}{{ loop.index }}.{{ reference }} -{% endfor %}{% endif %}"####; +{% endfor %}{% endif %} + +{% if cve %}### **开源检索** +{% if github_search | length > 0 %}{% for link in github_search %}{{ loop.index }}.{{ link }}{% endfor %}{% else %}暂未找到{% endif %}{% endif %}"####; const INIT_MSG_TEMPLATE: &str = r#" 数据初始化完成 @@ -110,6 +113,7 @@ mod tests { from: "https://avd.aliyun.com/high-risk/list".to_string(), tags, reasons, + github_search: vec![], is_valuable: false, }; let res = reader_vulninfo(v)?; diff --git a/src/search.rs b/src/search.rs new file mode 100644 index 0000000..db90fa2 --- /dev/null +++ b/src/search.rs @@ -0,0 +1,85 @@ +use regex::Regex; +use tracing::{info, warn}; + +use crate::{error::Result, utils::get_last_year_data}; + +pub async fn search_github_poc(cve_id: &str) -> Vec { + let mut res = Vec::new(); + let (nuclei_res, repo_res) = tokio::join!(search_nuclei_pr(cve_id), search_github_repo(cve_id)); + match nuclei_res { + Ok(nuclei) => res.extend(nuclei), + Err(e) => { + warn!("search nucli pr error:{}", e); + } + } + match repo_res { + Ok(repo) => res.extend(repo), + Err(e) => { + warn!("search github repo error:{}", e); + } + } + res +} + +pub async fn search_nuclei_pr(cve_id: &str) -> Result> { + info!("search nuclei PR of {}", cve_id); + let page = octocrab::instance() + .pulls("projectdiscovery", "nuclei-templates") + .list() + .per_page(100) + .page(1u32) + .send() + .await?; + let re = Regex::new(&format!(r"(?i)(?:\b|/|_){}(?:\b|/|_)", cve_id))?; + let links = page + .into_iter() + .filter(|pull| pull.title.is_some() || pull.body.is_some()) + .filter(|pull| { + re.is_match(pull.title.as_ref().unwrap_or(&String::new())) + || re.is_match(pull.body.as_ref().unwrap_or(&String::new())) + }) + .filter_map(|pull| pull.html_url) + .map(|u| u.to_string()) + .collect::>(); + Ok(links) +} + +pub async fn search_github_repo(cve_id: &str) -> Result> { + info!("search github repo of {}", cve_id); + let last_year = get_last_year_data(); + let query = format!("language:Python language:JavaScript language:C language:C++ language:Java language:PHP language:Ruby language:Rust language:C# created:>{} {}",last_year, cve_id); + let page = octocrab::instance() + .search() + .repositories(&query) + .per_page(100) + .page(1u32) + .send() + .await?; + let re = Regex::new(&format!(r"(?i)(?:\b|/|_){}(?:\b|/|_)", cve_id))?; + let links = page + .into_iter() + .filter_map(|r| r.html_url) + .filter(|url| re.captures(url.as_str()).is_some()) + .map_while(|u| Some(u.to_string())) + .collect::>(); + Ok(links) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_search_github_repo() -> Result<()> { + let res = search_github_repo("CVE-2021-4034").await?; + println!("{:?}", res); + Ok(()) + } + + #[tokio::test] + async fn test_search_nuclei_pr() -> Result<()> { + let res = search_nuclei_pr("CVE-2023-3380").await?; + println!("{:?}", res); + Ok(()) + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index afb9b1e..5fe0fa6 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,11 +1,17 @@ pub mod http_client; use crate::error::{Error, Result}; -use chrono::{DateTime, Duration, NaiveDate, Utc}; +use chrono::{DateTime, Duration, Local, NaiveDate, Utc}; use hmac::{Hmac, Mac}; use sha2::Sha256; use tera::{Context, Tera}; +pub fn get_last_year_data() -> String { + let current_date = Local::now(); + let last_year = current_date - Duration::days(365); + last_year.format("%Y-%m-%d").to_string() +} + pub fn check_over_two_week(date: &str) -> Result { let target_date = NaiveDate::parse_from_str(date, "%Y-%m-%d")?; let now = Utc::now().naive_utc().date();