feat: add init data and push init_msg

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:

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<Model> {
+    async fn crawling_task(&self, is_init: bool) -> Vec<Model> {
         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)]

src/error.rs

@@ -12,6 +12,10 @@ pub enum Error {
         backtrace: Box<std::backtrace::Backtrace>,
     },
 
+    // Model
+    #[error(transparent)]
+    Model(#[from] crate::models::ModelError),
+
     #[error("{0}")]
     Message(String),
 

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<Vec<VulnInfo>> {
         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);

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<u64> {
+        let count = vuln_informations::Entity::find().count(db).await?;
+        Ok(count)
+    }
+
     pub async fn creat_or_update(db: &DatabaseConnection, mut vuln: VulnInfo) -> ModelResult<Self> {
         let txn = db.begin().await?;
         let v = vuln_informations::Entity::find()

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<String> {
@@ -56,6 +60,18 @@ fn escape_markdown(input: String) -> String {
         .replace('!', "\\!")
 }
 
+pub fn render_init(version: String, vuln_count: u64, cron_config: String) -> Result<String> {
+    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::*;