From 53f9f50d7cd6c2dfde4b658140f0459e9fe1bd8c Mon Sep 17 00:00:00 2001 From: fan-tastic-z Date: Thu, 23 May 2024 14:59:33 +0800 Subject: [PATCH] feat: add init data and push init_msg --- config/development.yaml | 4 ++-- src/app.rs | 38 +++++++++++++++++++++++++++++---- src/error.rs | 4 ++++ src/grab/ti.rs | 5 ----- src/models/vuln_informations.rs | 10 ++++++--- src/push/msg_template.rs | 16 ++++++++++++++ 6 files changed, 63 insertions(+), 14 deletions(-) diff --git a/config/development.yaml b/config/development.yaml index 9073412..cceaa62 100644 --- a/config/development.yaml +++ b/config/development.yaml @@ -13,8 +13,8 @@ database: max_connections: 1 task: - # every day 7:00-22:00 interval 1 minute Execute task - cron_config: "0 */1 2-23 * * *" + # every day 7:00-22:00 interval 30 minute Execute task + cron_config: "0 */30 7-21 * * *" # Application logging configuration logger: diff --git a/src/app.rs b/src/app.rs index 8acc9c4..a7b62f6 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1,3 +1,4 @@ +use lazy_static::lazy_static; use sea_orm::DatabaseConnection; use std::{collections::HashMap, sync::Arc, time::Duration}; use tokio::{task::JoinSet, time}; @@ -11,10 +12,18 @@ use crate::{ error::Result, grab::{self, Grab, Severity}, models::_entities::vuln_informations::{self, Model}, - push::{msg_template::reader_vulninfo, telegram::Telegram}, + push::{ + msg_template::{reader_vulninfo, render_init}, + telegram::Telegram, + }, }; +lazy_static! { + static ref VERSION: &'static str = env!("CARGO_PKG_VERSION"); +} + const PAGE_LIMIT: i32 = 1; +const INIT_PAGE_LIMIT: i32 = 2; #[derive(Clone)] pub struct WatchVulnApp { @@ -36,12 +45,19 @@ impl WatchVulnApp { pub async fn run(&self) -> Result<()> { let self_arc = Arc::new(self.clone()); + + // init data + self_arc.crawling_task(true).await; + let local_count = vuln_informations::Model::query_count(&self.app_context.db).await?; + info!("init finished, local database has {} vulns", local_count); + self.push_init_msg(local_count).await?; + let sched = JobScheduler::new().await?; let schedule = self.app_context.config.task.cron_config.as_str(); let job = Job::new_async(schedule, move |_uuid, _lock| { let self_clone = self_arc.clone(); Box::pin(async move { - let res = self_clone.crawling_task().await; + let res = self_clone.crawling_task(false).await; info!("crawling over all count is: {}", res.len()); self_clone.push(res).await; }) @@ -54,12 +70,16 @@ impl WatchVulnApp { } } - async fn crawling_task(&self) -> Vec { + async fn crawling_task(&self, is_init: bool) -> Vec { tracing::info!("{:?}", self.app_context.config); let mut set = JoinSet::new(); for v in self.grabs.as_ref().values() { let grab = v.to_owned(); - set.spawn(async move { grab.get_update(PAGE_LIMIT).await }); + if is_init { + set.spawn(async move { grab.get_update(INIT_PAGE_LIMIT).await }); + } else { + set.spawn(async move { grab.get_update(PAGE_LIMIT).await }); + } } let mut new_vulns = Vec::new(); while let Some(set_res) = set.join_next().await { @@ -117,6 +137,16 @@ impl WatchVulnApp { } } } + + async fn push_init_msg(&self, local_count: u64) -> Result<()> { + let init_msg = render_init( + VERSION.to_string(), + local_count, + self.app_context.config.task.cron_config.clone(), + )?; + self.app_context.tg_bot.push_markdown(init_msg).await?; + Ok(()) + } } #[derive(Clone, Debug)] diff --git a/src/error.rs b/src/error.rs index d1536b4..f312165 100644 --- a/src/error.rs +++ b/src/error.rs @@ -12,6 +12,10 @@ pub enum Error { backtrace: Box, }, + // Model + #[error(transparent)] + Model(#[from] crate::models::ModelError), + #[error("{0}")] Message(String), diff --git a/src/grab/ti.rs b/src/grab/ti.rs index bb29604..b766f3c 100644 --- a/src/grab/ti.rs +++ b/src/grab/ti.rs @@ -1,7 +1,6 @@ use async_trait::async_trait; use reqwest::header::{self}; use serde::{Deserialize, Serialize}; -use tracing::info; use super::{Grab, Provider, Severity, VulnInfo}; use crate::error::Result; @@ -69,10 +68,6 @@ impl TiCrawler { pub async fn get_vuln_infos(&self) -> Result> { let ti_one_day_resp = self.get_ti_one_day_resp().await?; - info!( - "key_vuln_add len is {}", - ti_one_day_resp.data.key_vuln_add.len() - ); let mut vuln_infos = Vec::with_capacity(ti_one_day_resp.data.key_vuln_add.len()); for detail in ti_one_day_resp.data.key_vuln_add { let tags = self.get_tags(detail.tag); diff --git a/src/models/vuln_informations.rs b/src/models/vuln_informations.rs index 51b1398..e395ba6 100644 --- a/src/models/vuln_informations.rs +++ b/src/models/vuln_informations.rs @@ -1,7 +1,6 @@ use sea_orm::{ - ActiveModelTrait, - ActiveValue::{self}, - ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, TransactionTrait, + ActiveModelTrait, ActiveValue, ColumnTrait, DatabaseConnection, EntityTrait, PaginatorTrait, + QueryFilter, TransactionTrait, }; use tracing::info; @@ -22,6 +21,11 @@ impl super::_entities::vuln_informations::Model { vuln.ok_or_else(|| ModelError::EntityNotFound) } + pub async fn query_count(db: &DatabaseConnection) -> ModelResult { + let count = vuln_informations::Entity::find().count(db).await?; + Ok(count) + } + pub async fn creat_or_update(db: &DatabaseConnection, mut vuln: VulnInfo) -> 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 34b82bc..10f1614 100644 --- a/src/push/msg_template.rs +++ b/src/push/msg_template.rs @@ -22,6 +22,10 @@ const VULN_INFO_MSG_TEMPLATE: &str = r####" {% for reference in references %}{{ loop.index }}.{{ reference }} {% endfor %}{% endif %}"####; +const INIT_MSG_TEMPLATE: &str = r#" +数据初始化完成，当前版本 {{ version }} 本地漏洞数量: {{ vuln_count }} 检查周期配置： {{ cron_config }} +"#; + const MAX_REFERENCE_LENGTH: usize = 8; pub fn reader_vulninfo(mut vuln: VulnInfo) -> Result { @@ -56,6 +60,18 @@ fn escape_markdown(input: String) -> String { .replace('!', "\\!") } +pub fn render_init(version: String, vuln_count: u64, cron_config: String) -> Result { + let json_value = serde_json::json!( + { + "version": version, + "vuln_count": vuln_count, + "cron_config": cron_config, + } + ); + let markdown = render_string(INIT_MSG_TEMPLATE, &json_value)?; + Ok(escape_markdown(markdown)) +} + #[cfg(test)] mod tests { use super::*;