fan-tastic-z/vulnfeed
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: update error_stack to 0.6.0

Browse Source
This commit is contained in:
fan-tastic-z
2025-09-03 16:46:25 +08:00
parent 521bfdd23d
commit cabcedb0b1
31 changed files with 209 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
Cargo.lock generated
View File

@@ -308,9 +308,9 @@ dependencies = [
[[package]]
name = "clap"
version = "4.5.45"
version = "4.5.47"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fc0e74a703892159f5ae7d3aac52c8e6c392f5ae5f359c70b5881d60aaac318"
checksum = "7eac00902d9d136acd712710d71823fb8ac8004ca445a89e73a41d45aa712931"
dependencies = [
"clap_builder",
"clap_derive",
@@ -318,9 +318,9 @@ dependencies = [
[[package]]
name = "clap_builder"
version = "4.5.44"
version = "4.5.47"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3e7f4214277f3c7aa526a59dd3fbe306a370daee1f8b7b8c987069cd8e888a8"
checksum = "2ad9bbf750e73b5884fb8a211a9424a1906c1e156724260fdae972f31d70e1d6"
dependencies = [
"anstream",
"anstyle",
@@ -330,9 +330,9 @@ dependencies = [
[[package]]
name = "clap_derive"
version = "4.5.45"
version = "4.5.47"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "14cb31bb0a7d536caef2639baa7fad459e15c3144efefa6dbd1c84562c4739f6"
checksum = "bbfd7eae0b0f1a6e63d4b13c9c478de77c2eb546fba158ad50b4203dc24b9f9c"
dependencies = [
"heck 0.5.0",
"proc-macro2",
@@ -756,9 +756,9 @@ dependencies = [
[[package]]
name = "error-stack"
version = "0.5.0"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fe413319145d1063f080f27556fd30b1d70b01e2ba10c2a6e40d4be982ffc5d1"
checksum = "b878b3fac9613c3c7f22eb70bc8a3c6ebdc03cc11479ee60fde1692d747fd45f"
dependencies = [
"anyhow",
"rustc_version",
@@ -2089,9 +2089,9 @@ dependencies = [
[[package]]
name = "mea"
version = "0.4.1"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "942eb3e014947e1a4de72d833d2df344e0c267d11407a83ebe6c767fb722aae1"
checksum = "6f3876c525f3c3f650b97b69b61eeb27c976e7725e9b59120f774e6fcf02b518"
dependencies = [
"slab",
]
@@ -3471,9 +3471,9 @@ dependencies = [
[[package]]
name = "shadow-rs"
version = "1.2.1"
version = "1.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f0b6af233ae5461c3c6b30db79190ec5fbbef048ebbd5f2cbb3043464168e00"
checksum = "b8aa5c0570cd9654158bd39f0f8caba24edbc058313946e89f4648b1de1ecf49"
dependencies = [
"const_format",
"git2",

8
Cargo.toml
View File

@@ -12,11 +12,11 @@ anstyle = "1.0.11"
argon2 = { version = "0.5.3", features = ["std"] }
async-trait = "0.1.89"
chrono = { version = "0.4.41", features = ["serde"] }
clap = { version = "4.5.45", features = ["derive"] }
clap = { version = "4.5.47", features = ["derive"] }
const_format = "0.2.34"
ctrlc = "3.4.7"
dashmap = "6.1.0"
error-stack = "0.5.0"
error-stack = "0.6.0"
fastimer = "0.9.0"
gix-discover = "0.41.0"
jsonwebtoken = "9.3.1"
@@ -30,7 +30,7 @@ logforth = { version = "0.27.0", features = [
"append-rolling-file",
"diagnostic-fastrace",
] }
mea = "0.4.1"
mea = "0.4.2"
modql = { version = "0.4.1", features = ["with-sea-query"] }
nutype = { version = "0.6.2", features = ["serde"] }
pin-project = "1.1.10"
@@ -45,7 +45,7 @@ sea-query-binder = { version = "0.7.0", features = [
] }
serde = { version = "1.0.219", features = ["derive"] }
serde_json = "1.0.143"
shadow-rs = "1.2.1"
shadow-rs = "1.3.0"
sqlx = { version = "0.8.6", features = [
"postgres",
"runtime-tokio",

8
src/bin/vulnfeed/main.rs
View File

@@ -1,8 +1,6 @@
use clap::Parser;
use error_stack::Result;
use vulnfeed::{
cli,
errors::Error,
AppResult, cli,
utils::{styled::styled, version::version},
};
@@ -14,7 +12,7 @@ struct Command {
}
impl Command {
pub fn run(self) -> Result<(), Error> {
pub fn run(self) -> AppResult<()> {
match self.cmd {
SubCommand::Server(cmd) => cmd.run(),
SubCommand::CreateSuperUser(cmd) => cmd.run(),
@@ -28,7 +26,7 @@ enum SubCommand {
CreateSuperUser(cli::CreateSuperUser),
}
fn main() -> Result<(), Error> {
fn main() -> AppResult<()> {
let cmd = Command::parse();
cmd.run()
}

21
src/cli.rs
View File

@@ -1,6 +1,7 @@
use std::{path::PathBuf, sync::Arc};
use crate::{
AppResult,
config::settings::{Config, LoadConfigResult, load_config},
domain::{
models::{
@@ -27,7 +28,7 @@ use crate::{
},
};
use clap::ValueHint;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
#[derive(Clone)]
pub struct Ctx<S: VulnService + Send + Sync + 'static> {
@@ -44,7 +45,7 @@ pub struct CommandStart {
}
impl CommandStart {
pub fn run(self) -> Result<(), Error> {
pub fn run(self) -> AppResult<()> {
error_stack::Report::set_color_mode(error_stack::fmt::ColorMode::None);
let LoadConfigResult { config, warnings } = load_config(self.config_file)?;
let telemetry_runtime = make_telemetry_runtime();
@@ -56,11 +57,14 @@ impl CommandStart {
}
log::info!("server is starting with config: {config:#?}");
let server_runtime = make_vulnfeed_runtime();
server_runtime.block_on(run_server(&server_runtime, config))
server_runtime
.block_on(run_server(&server_runtime, config))
.change_context(Error::Message("failed to start server".to_string()))?;
Ok(())
}
}
async fn run_server(server_rt: &Runtime, config: Config) -> Result<(), Error> {
async fn run_server(server_rt: &Runtime, config: Config) -> AppResult<()> {
let make_error = || Error::Message("failed to start server".to_string());
let (shutdown_tx, shutdown_rx) = mea::shutdown::new_pair();
let (acceptor, advertise_addr) = make_acceptor_and_advertise_addr(
@@ -134,7 +138,7 @@ pub struct CreateSuperUser {
}
impl CreateSuperUser {
pub fn run(self) -> Result<(), Error> {
pub fn run(self) -> AppResult<()> {
error_stack::Report::set_color_mode(error_stack::fmt::ColorMode::None);
let LoadConfigResult { config, warnings } = load_config(self.config_file)?;
let telemetry_runtime = make_telemetry_runtime();
@@ -145,11 +149,14 @@ impl CreateSuperUser {
log::warn!("{warning}");
}
let init_data_runtime = make_init_data_runtime();
init_data_runtime.block_on(run_create_super_user(config, self.password))
init_data_runtime
.block_on(run_create_super_user(config, self.password))
.change_context_lazy(|| Error::Message("failed to create super user".to_string()))?;
Ok(())
}
}
async fn run_create_super_user(config: Config, password: String) -> Result<(), Error> {
async fn run_create_super_user(config: Config, password: String) -> AppResult<()> {
let make_error = || Error::Message("failed to create super user".to_string());
let password_hash = compute_password_hash(&password).change_context_lazy(make_error)?;
let db = Pg::new(&config).await.change_context_lazy(make_error)?;

5
src/config/settings.rs
View File

@@ -1,10 +1,11 @@
use std::{path::PathBuf, str::FromStr};
use error_stack::{Result, ResultExt, bail};
use error_stack::{ResultExt, bail};
use serde::de::IntoDeserializer;
use serde::{Deserialize, Serialize};
use toml_edit::DocumentMut;
use crate::AppResult;
use crate::errors::Error;
pub struct LoadConfigResult {
@@ -12,7 +13,7 @@ pub struct LoadConfigResult {
pub warnings: Vec<String>,
}
pub fn load_config(config_file: PathBuf) -> Result<LoadConfigResult, Error> {
pub fn load_config(config_file: PathBuf) -> AppResult<LoadConfigResult> {
let content = std::fs::read_to_string(&config_file).change_context_lazy(|| {
Error::Message(format!(
"failed to read config file {}",

39
src/domain/ports.rs
View File

@@ -1,4 +1,5 @@
use crate::{
AppResult,
domain::models::{
admin_user::AdminUser,
auth::LoginRequest,
@@ -9,68 +10,58 @@ use crate::{
VulnInformation,
},
},
errors::Error,
};
use error_stack::Result;
use std::future::Future;
pub trait VulnService: Clone + Send + Sync + 'static {
fn login(&self, req: &LoginRequest) -> impl Future<Output = Result<AdminUser, Error>> + Send;
fn login(&self, req: &LoginRequest) -> impl Future<Output = AppResult<AdminUser>> + Send;
fn create_sync_data_task(
&self,
req: CreateSyncDataTaskRequest,
) -> impl Future<Output = Result<i64, Error>> + Send;
fn get_sync_data_task(
&self,
) -> impl Future<Output = Result<Option<SyncDataTask>, Error>> + Send;
) -> impl Future<Output = AppResult<i64>> + Send;
fn get_sync_data_task(&self) -> impl Future<Output = AppResult<Option<SyncDataTask>>> + Send;
fn list_vulnfusion_information(
&self,
req: ListVulnInformationRequest,
) -> impl Future<Output = Result<ListVulnInformationResponseData, Error>> + Send;
) -> impl Future<Output = AppResult<ListVulnInformationResponseData>> + Send;
fn get_vuln_information(
&self,
req: GetVulnInformationRequest,
) -> impl Future<Output = Result<Option<VulnInformation>, Error>> + Send;
) -> impl Future<Output = AppResult<Option<VulnInformation>>> + Send;
fn create_ding_bot_config(
&self,
req: CreateDingBotRequest,
) -> impl Future<Output = Result<i64, Error>> + Send;
) -> impl Future<Output = AppResult<i64>> + Send;
fn get_ding_bot_config(
&self,
) -> impl Future<Output = Result<Option<DingBotConfig>, Error>> + Send;
fn get_ding_bot_config(&self) -> impl Future<Output = AppResult<Option<DingBotConfig>>> + Send;
}
pub trait VulnRepository: Clone + Send + Sync + 'static {
fn login(&self, req: &LoginRequest) -> impl Future<Output = Result<AdminUser, Error>> + Send;
fn login(&self, req: &LoginRequest) -> impl Future<Output = AppResult<AdminUser>> + Send;
fn create_sync_data_task(
&self,
req: CreateSyncDataTaskRequest,
) -> impl Future<Output = Result<i64, Error>> + Send;
) -> impl Future<Output = AppResult<i64>> + Send;
fn get_sync_data_task(
&self,
) -> impl Future<Output = Result<Option<SyncDataTask>, Error>> + Send;
fn get_sync_data_task(&self) -> impl Future<Output = AppResult<Option<SyncDataTask>>> + Send;
fn list_vuln_information(
&self,
req: ListVulnInformationRequest,
) -> impl Future<Output = Result<ListVulnInformationResponseData, Error>> + Send;
) -> impl Future<Output = AppResult<ListVulnInformationResponseData>> + Send;
fn get_vuln_information(
&self,
req: GetVulnInformationRequest,
) -> impl Future<Output = Result<Option<VulnInformation>, Error>> + Send;
) -> impl Future<Output = AppResult<Option<VulnInformation>>> + Send;
fn create_ding_bot_config(
&self,
req: CreateDingBotRequest,
) -> impl Future<Output = Result<i64, Error>> + Send;
) -> impl Future<Output = AppResult<i64>> + Send;
fn get_ding_bot_config(
&self,
) -> impl Future<Output = Result<Option<DingBotConfig>, Error>> + Send;
fn get_ding_bot_config(&self) -> impl Future<Output = AppResult<Option<DingBotConfig>>> + Send;
}

19
src/domain/services.rs
View File

@@ -1,4 +1,5 @@
use crate::{
AppResult,
domain::{
models::{
admin_user::AdminUser,
@@ -12,9 +13,7 @@ use crate::{
},
ports::{VulnRepository, VulnService},
},
errors::Error,
};
use error_stack::Result;
#[derive(Debug, Clone)]
pub struct Service<R>
@@ -37,16 +36,16 @@ impl<R> VulnService for Service<R>
where
R: VulnRepository,
{
async fn login(&self, req: &LoginRequest) -> Result<AdminUser, Error> {
async fn login(&self, req: &LoginRequest) -> AppResult<AdminUser> {
let res = self.repo.login(req).await?;
Ok(res)
}
async fn create_sync_data_task(&self, req: CreateSyncDataTaskRequest) -> Result<i64, Error> {
async fn create_sync_data_task(&self, req: CreateSyncDataTaskRequest) -> AppResult<i64> {
let ret = self.repo.create_sync_data_task(req).await?;
Ok(ret)
}
async fn get_sync_data_task(&self) -> Result<Option<SyncDataTask>, Error> {
async fn get_sync_data_task(&self) -> AppResult<Option<SyncDataTask>> {
let ret = self.repo.get_sync_data_task().await?;
Ok(ret)
}
@@ -54,25 +53,25 @@ where
async fn list_vulnfusion_information(
&self,
req: ListVulnInformationRequest,
) -> Result<ListVulnInformationResponseData, Error> {
let ret = self.repo.list_vuln_information(req).await?; // Implement the logic here
) -> AppResult<ListVulnInformationResponseData> {
let ret = self.repo.list_vuln_information(req).await?;
Ok(ret)
}
async fn get_vuln_information(
&self,
req: GetVulnInformationRequest,
) -> Result<Option<VulnInformation>, Error> {
) -> AppResult<Option<VulnInformation>> {
let ret = self.repo.get_vuln_information(req).await?;
Ok(ret)
}
async fn get_ding_bot_config(&self) -> Result<Option<DingBotConfig>, Error> {
async fn get_ding_bot_config(&self) -> AppResult<Option<DingBotConfig>> {
let ret = self.repo.get_ding_bot_config().await?;
Ok(ret)
}
async fn create_ding_bot_config(&self, req: CreateDingBotRequest) -> Result<i64, Error> {
async fn create_ding_bot_config(&self, req: CreateDingBotRequest) -> AppResult<i64> {
let ret = self.repo.create_ding_bot_config(req).await?;
Ok(ret)
}

6
src/lib.rs
View File

@@ -1,3 +1,7 @@
use error_stack::Report;
use crate::errors::Error;
pub mod cli;
pub mod config;
pub mod domain;
@@ -5,3 +9,5 @@ pub mod errors;
pub mod input;
pub mod output;
pub mod utils;
pub type AppResult<T> = Result<T, Report<Error>>;

7
src/output/db/admin_user.rs
View File

@@ -1,9 +1,8 @@
use error_stack::Result;
use sqlx::{Postgres, Transaction};
use crate::{
AppResult,
domain::models::admin_user::{AdminUser, AdminUsername, CreateAdminUserRequest},
errors::Error,
output::db::base::{Dao, dao_fetch_by_column, dao_upsert},
};
@@ -17,14 +16,14 @@ impl AdminUserDao {
pub async fn create_super_user(
tx: &mut Transaction<'_, Postgres>,
req: CreateAdminUserRequest,
) -> Result<i64, Error> {
) -> AppResult<i64> {
dao_upsert::<Self, _>(tx, req, "name", &["password"]).await
}
pub async fn fetch_by_name(
tx: &mut Transaction<'_, Postgres>,
name: &AdminUsername,
) -> Result<Option<AdminUser>, Error> {
) -> AppResult<Option<AdminUser>> {
dao_fetch_by_column::<Self, AdminUser>(tx, "name", name.as_ref()).await
}
}

26
src/output/db/base.rs
View File

@@ -1,4 +1,4 @@
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use modql::{SIden, field::HasSeaFields};
use sea_query::{
Alias, Asterisk, Condition, Expr, Iden, IntoIden, OnConflict, PostgresQueryBuilder, Query,
@@ -7,7 +7,7 @@ use sea_query::{
use sea_query_binder::SqlxBinder;
use sqlx::{FromRow, Postgres, Row, Transaction};
use crate::errors::Error;
use crate::{AppResult, errors::Error};
pub trait Dao {
const TABLE: &'static str;
@@ -21,7 +21,7 @@ pub enum CommonIden {
Id,
}
pub async fn dao_create<D, E>(tx: &mut Transaction<'_, Postgres>, req: E) -> Result<i64, Error>
pub async fn dao_create<D, E>(tx: &mut Transaction<'_, Postgres>, req: E) -> AppResult<i64>
where
E: HasSeaFields,
D: Dao,
@@ -45,7 +45,7 @@ where
Ok(id)
}
pub async fn dao_first<D, T>(tx: &mut Transaction<'_, Postgres>) -> Result<Option<T>, Error>
pub async fn dao_first<D, T>(tx: &mut Transaction<'_, Postgres>) -> AppResult<Option<T>>
where
D: Dao,
T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Unpin + Send,
@@ -68,7 +68,7 @@ pub async fn dao_upsert<D, E>(
req: E,
conflict_column: &str,
update_columns: &[&str],
) -> Result<i64, Error>
) -> AppResult<i64>
where
E: HasSeaFields,
D: Dao,
@@ -114,7 +114,7 @@ pub async fn dao_fetch_by_column<D, T>(
tx: &mut Transaction<'_, Postgres>,
column_name: &str,
value: &str,
) -> Result<Option<T>, Error>
) -> AppResult<Option<T>>
where
D: Dao,
T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Unpin + Send,
@@ -134,11 +134,7 @@ where
Ok(result)
}
pub async fn dao_update<D, E>(
tx: &mut Transaction<'_, Postgres>,
id: i64,
req: E,
) -> Result<u64, Error>
pub async fn dao_update<D, E>(tx: &mut Transaction<'_, Postgres>, id: i64, req: E) -> AppResult<u64>
where
E: HasSeaFields,
D: Dao,
@@ -173,7 +169,7 @@ pub async fn dao_update_field<D>(
id: i64,
field_name: &str,
field_value: impl Into<sea_query::Value>,
) -> Result<u64, Error>
) -> AppResult<u64>
where
D: Dao,
{
@@ -198,7 +194,7 @@ where
pub async fn dao_fetch_by_id<D, T>(
tx: &mut Transaction<'_, Postgres>,
id: i64,
) -> Result<Option<T>, Error>
) -> AppResult<Option<T>>
where
D: Dao,
T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Unpin + Send,
@@ -284,7 +280,7 @@ impl<D: Dao> DaoQueryBuilder<D> {
self
}
pub async fn fetch_all<T>(self, tx: &mut Transaction<'_, Postgres>) -> Result<Vec<T>, Error>
pub async fn fetch_all<T>(self, tx: &mut Transaction<'_, Postgres>) -> AppResult<Vec<T>>
where
T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Unpin + Send,
{
@@ -299,7 +295,7 @@ impl<D: Dao> DaoQueryBuilder<D> {
Ok(result)
}
pub async fn count(self, tx: &mut Transaction<'_, Postgres>) -> Result<i64, Error> {
pub async fn count(self, tx: &mut Transaction<'_, Postgres>) -> AppResult<i64> {
let mut count_query = Query::select();
count_query
.from(D::table_ref())

7
src/output/db/ding_bot_config.rs
View File

@@ -1,9 +1,8 @@
use error_stack::Result;
use sqlx::{Postgres, Transaction};
use crate::{
AppResult,
domain::models::ding_bot::{CreateDingBotRequest, DingBotConfig},
errors::Error,
output::db::base::{Dao, dao_create, dao_first, dao_update},
};
@@ -17,7 +16,7 @@ impl DingBotConfigDao {
pub async fn create(
tx: &mut Transaction<'_, Postgres>,
req: CreateDingBotRequest,
) -> Result<i64, Error> {
) -> AppResult<i64> {
let ding_bot_config: Option<DingBotConfig> = dao_first::<Self, _>(tx).await?;
if let Some(config) = ding_bot_config {
dao_update::<Self, _>(tx, config.id, req).await?;
@@ -27,7 +26,7 @@ impl DingBotConfigDao {
Ok(ret)
}
pub async fn first(tx: &mut Transaction<'_, Postgres>) -> Result<Option<DingBotConfig>, Error> {
pub async fn first(tx: &mut Transaction<'_, Postgres>) -> AppResult<Option<DingBotConfig>> {
let ding_bot_config = dao_first::<Self, _>(tx).await?;
Ok(ding_bot_config)
}

6
src/output/db/pg.rs
View File

@@ -1,10 +1,10 @@
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use sqlx::{
Pool, Postgres,
postgres::{PgConnectOptions, PgPoolOptions},
};
use crate::{config::settings::Config, errors::Error};
use crate::{AppResult, config::settings::Config, errors::Error};
#[derive(Debug, Clone)]
pub struct Pg {
@@ -12,7 +12,7 @@ pub struct Pg {
}
impl Pg {
pub async fn new(config: &Config) -> Result<Self, Error> {
pub async fn new(config: &Config) -> AppResult<Self> {
let opts = PgConnectOptions::new()
.host(&config.database.host)
.port(config.database.port)

17
src/output/db/repository_impl.rs
View File

@@ -1,6 +1,7 @@
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use crate::{
AppResult,
domain::{
models::{
admin_user::AdminUser,
@@ -23,7 +24,7 @@ use crate::{
};
impl VulnRepository for Pg {
async fn login(&self, req: &LoginRequest) -> Result<AdminUser, Error> {
async fn login(&self, req: &LoginRequest) -> AppResult<AdminUser> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -40,7 +41,7 @@ impl VulnRepository for Pg {
Err(Error::BadRequest("invalid account or password".to_string()).into())
}
async fn create_sync_data_task(&self, req: CreateSyncDataTaskRequest) -> Result<i64, Error> {
async fn create_sync_data_task(&self, req: CreateSyncDataTaskRequest) -> AppResult<i64> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -52,7 +53,7 @@ impl VulnRepository for Pg {
Ok(sync_data_task_id)
}
async fn get_sync_data_task(&self) -> Result<Option<SyncDataTask>, Error> {
async fn get_sync_data_task(&self) -> AppResult<Option<SyncDataTask>> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -67,7 +68,7 @@ impl VulnRepository for Pg {
async fn list_vuln_information(
&self,
req: ListVulnInformationRequest,
) -> Result<ListVulnInformationResponseData, Error> {
) -> AppResult<ListVulnInformationResponseData> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -91,7 +92,7 @@ impl VulnRepository for Pg {
async fn get_vuln_information(
&self,
req: GetVulnInformationRequest,
) -> Result<Option<VulnInformation>, Error> {
) -> AppResult<Option<VulnInformation>> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -103,7 +104,7 @@ impl VulnRepository for Pg {
Ok(vuln_information)
}
async fn create_ding_bot_config(&self, req: CreateDingBotRequest) -> Result<i64, Error> {
async fn create_ding_bot_config(&self, req: CreateDingBotRequest) -> AppResult<i64> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -115,7 +116,7 @@ impl VulnRepository for Pg {
Ok(id)
}
async fn get_ding_bot_config(&self) -> Result<Option<DingBotConfig>, Error> {
async fn get_ding_bot_config(&self) -> AppResult<Option<DingBotConfig>> {
let mut tx =
self.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())

9
src/output/db/sync_data_task.rs
View File

@@ -1,10 +1,9 @@
use error_stack::Result;
use sea_query::Value;
use sqlx::{Postgres, Transaction};
use crate::{
AppResult,
domain::models::sync_data_task::{CreateSyncDataTaskRequest, SyncDataTask},
errors::Error,
output::db::base::{Dao, dao_create, dao_first, dao_update, dao_update_field},
};
@@ -15,7 +14,7 @@ impl Dao for SyncDataTaskDao {
}
impl SyncDataTaskDao {
pub async fn first(tx: &mut Transaction<'_, Postgres>) -> Result<Option<SyncDataTask>, Error> {
pub async fn first(tx: &mut Transaction<'_, Postgres>) -> AppResult<Option<SyncDataTask>> {
let task = dao_first::<Self, _>(tx).await?;
Ok(task)
}
@@ -23,7 +22,7 @@ impl SyncDataTaskDao {
pub async fn create(
tx: &mut Transaction<'_, Postgres>,
req: CreateSyncDataTaskRequest,
) -> Result<i64, Error> {
) -> AppResult<i64> {
let task: Option<SyncDataTask> = dao_first::<Self, _>(tx).await?;
if let Some(t) = task {
dao_update::<Self, _>(tx, t.id, req).await?;
@@ -37,7 +36,7 @@ impl SyncDataTaskDao {
tx: &mut Transaction<'_, Postgres>,
id: i64,
job_id: String,
) -> Result<u64, Error> {
) -> AppResult<u64> {
let row = dao_update_field::<Self>(tx, id, "job_id", Value::String(Some(Box::new(job_id))))
.await?;
Ok(row)

17
src/output/db/vuln_information.rs
View File

@@ -1,13 +1,12 @@
use error_stack::Result;
use sea_query::Value;
use sqlx::{Postgres, Transaction};
use crate::{
AppResult,
domain::models::{
page_utils::PageFilter,
vuln_information::{CreateVulnInformation, SearchParams, VulnInformation},
},
errors::Error,
output::db::base::{
Dao, DaoQueryBuilder, dao_create, dao_fetch_by_column, dao_fetch_by_id, dao_update,
dao_update_field,
@@ -29,7 +28,7 @@ impl VulnInformationDao {
tx: &mut Transaction<'_, Postgres>,
id: i64,
status: bool,
) -> Result<u64, Error> {
) -> AppResult<u64> {
let row = dao_update_field::<Self>(tx, id, "pushed", Value::Bool(Some(status))).await?;
Ok(row)
}
@@ -37,7 +36,7 @@ impl VulnInformationDao {
pub async fn create(
tx: &mut Transaction<'_, Postgres>,
req: CreateVulnInformation,
) -> Result<i64, Error> {
) -> AppResult<i64> {
let id = dao_create::<Self, _>(tx, req).await?;
Ok(id)
}
@@ -91,7 +90,7 @@ impl VulnInformationDao {
pub async fn create_or_update(
tx: &mut Transaction<'_, Postgres>,
mut req: CreateVulnInformation,
) -> Result<(i64, bool), Error> {
) -> AppResult<(i64, bool)> {
let mut as_new_vuln = false;
if let Some(mut vuln) =
dao_fetch_by_column::<Self, VulnInformation>(tx, "key", &req.key).await?
@@ -118,7 +117,7 @@ impl VulnInformationDao {
tx: &mut Transaction<'_, Postgres>,
page_filter: &PageFilter,
search_params: &SearchParams,
) -> Result<Vec<VulnInformation>, Error> {
) -> AppResult<Vec<VulnInformation>> {
let mut query_builder = DaoQueryBuilder::<Self>::new();
if let Some(title) = &search_params.title {
@@ -150,7 +149,7 @@ impl VulnInformationDao {
pub async fn filter_vulnfusion_information_count(
tx: &mut Transaction<'_, Postgres>,
search_params: &SearchParams,
) -> Result<i64, Error> {
) -> AppResult<i64> {
let mut query_builder = DaoQueryBuilder::<Self>::new();
if let Some(title) = &search_params.title {
@@ -175,14 +174,14 @@ impl VulnInformationDao {
pub async fn fetch_by_id(
tx: &mut Transaction<'_, Postgres>,
id: i64,
) -> Result<Option<VulnInformation>, Error> {
) -> AppResult<Option<VulnInformation>> {
dao_fetch_by_id::<Self, VulnInformation>(tx, id).await
}
pub async fn fetch_by_key(
tx: &mut Transaction<'_, Postgres>,
key: &str,
) -> Result<Option<VulnInformation>, Error> {
) -> AppResult<Option<VulnInformation>> {
dao_fetch_by_column::<Self, VulnInformation>(tx, "key", key).await
}
}

43
src/output/plugins/avd.rs
View File

@@ -1,5 +1,5 @@
use async_trait::async_trait;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use regex::Regex;
use reqwest::Url;
@@ -8,6 +8,7 @@ use scraper::{Html, Selector};
use serde_json::{Value, json};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -41,7 +42,7 @@ impl VulnPlugin for AVDPlugin {
self.link.to_string()
}
async fn update(&self, page_limit: i32) -> Result<(), Error> {
async fn update(&self, page_limit: i32) -> AppResult<()> {
let mut page_count = self.get_page_count().await?;
if page_count > page_limit {
page_count = page_limit;
@@ -67,7 +68,7 @@ impl VulnPlugin for AVDPlugin {
}
impl AVDPlugin {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> Result<AVDPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<AVDPlugin> {
let http_client = HttpClient::try_new()?;
let avd = AVDPlugin {
name: "AVDPlugin".to_string(),
@@ -80,7 +81,7 @@ impl AVDPlugin {
Ok(avd)
}
pub async fn get_page_count(&self) -> Result<i32, Error> {
pub async fn get_page_count(&self) -> AppResult<i32> {
let new_url = self.waf_bypass(&self.link).await?;
let content = self.http_client.get_html_content(&new_url).await?;
@@ -101,7 +102,7 @@ impl AVDPlugin {
}
}
pub async fn parse_page(&self, page: i32) -> Result<Vec<CreateVulnInformation>, Error> {
pub async fn parse_page(&self, page: i32) -> AppResult<Vec<CreateVulnInformation>> {
let page_url = format!("{}?page={}", self.link, page);
let document = self.get_document(&page_url).await?;
let detail_links = self.get_detail_links(document)?;
@@ -116,7 +117,7 @@ impl AVDPlugin {
Ok(res)
}
fn get_detail_links(&self, document: Html) -> Result<Vec<String>, Error> {
fn get_detail_links(&self, document: Html) -> AppResult<Vec<String>> {
let src_url_selector = Selector::parse("tbody tr td a")
.map_err(|err| Error::Message(format!("selector parse error: {}", err)))?;
@@ -128,7 +129,7 @@ impl AVDPlugin {
Ok(detail_links)
}
pub async fn parse_detail_page(&self, href: &str) -> Result<CreateVulnInformation, Error> {
pub async fn parse_detail_page(&self, href: &str) -> AppResult<CreateVulnInformation> {
let detail_url = format!("https://avd.aliyun.com{}", href);
let document = self.get_document(&detail_url).await?;
@@ -180,7 +181,7 @@ impl AVDPlugin {
Ok(data)
}
fn get_avd_id(&self, detail_url: &str) -> Result<String, Error> {
fn get_avd_id(&self, detail_url: &str) -> AppResult<String> {
let url = Url::parse(detail_url)
.map_err(|err| Error::Message(format!("avd get detail url parse error {}", err)))?;
let avd_id = url
@@ -204,7 +205,7 @@ impl AVDPlugin {
Ok(references)
}
fn get_solutions(&self, document: &Html) -> Result<String, Error> {
fn get_solutions(&self, document: &Html) -> AppResult<String> {
let solutions_selector = Selector::parse(".text-detail").map_err(|err| {
Error::Message(format!("avd get solutions selector parse error {}", err))
})?;
@@ -219,7 +220,7 @@ impl AVDPlugin {
Ok(solutions)
}
fn get_description(&self, document: &Html) -> Result<String, Error> {
fn get_description(&self, document: &Html) -> AppResult<String> {
let description_selector = Selector::parse(".text-detail div").map_err(|err| {
Error::Message(format!("avd get description selector parse error {}", err))
})?;
@@ -231,7 +232,7 @@ impl AVDPlugin {
Ok(description)
}
fn get_title(&self, document: &Html) -> Result<String, Error> {
fn get_title(&self, document: &Html) -> AppResult<String> {
let title_selector = Selector::parse("h5[class='header__title'] .header__title__text")
.map_err(|err| Error::Message(format!("avd get title selector parse error {}", err)))?;
let title = document
@@ -244,7 +245,7 @@ impl AVDPlugin {
Ok(title)
}
fn get_severity(&self, document: &Html) -> Result<Severity, Error> {
fn get_severity(&self, document: &Html) -> AppResult<Severity> {
let level_selector = Selector::parse("h5[class='header__title'] .badge")
.map_err(|err| Error::Message(format!("avd get level selector parse error {}", err)))?;
let level = document
@@ -264,7 +265,7 @@ impl AVDPlugin {
Ok(severity)
}
fn get_mertric_value(&self, document: &Html, index: usize) -> Result<String, Error> {
fn get_mertric_value(&self, document: &Html, index: usize) -> AppResult<String> {
let value_selector = Selector::parse(".metric-value").map_err(|e| {
Error::Message(format!("avd get metric value selector parse error {}", e))
})?;
@@ -278,7 +279,7 @@ impl AVDPlugin {
Ok(metric_value)
}
fn get_cve_id(&self, document: &Html) -> Result<String, Error> {
fn get_cve_id(&self, document: &Html) -> AppResult<String> {
let mut cve_id = self.get_mertric_value(document, 0)?;
if !Regex::new(CVEID_REGEXP)
.change_context_lazy(|| Error::Message("avd get cve id regex parse error".to_string()))?
@@ -289,22 +290,22 @@ impl AVDPlugin {
Ok(cve_id)
}
fn get_utilization(&self, document: &Html) -> Result<String, Error> {
fn get_utilization(&self, document: &Html) -> AppResult<String> {
self.get_mertric_value(document, 1)
}
fn get_disclosure(&self, document: &Html) -> Result<String, Error> {
fn get_disclosure(&self, document: &Html) -> AppResult<String> {
self.get_mertric_value(document, 3)
}
async fn get_document(&self, url: &str) -> Result<Html, Error> {
async fn get_document(&self, url: &str) -> AppResult<Html> {
let new_url = self.waf_bypass(url).await?;
let content = self.http_client.get_html_content(&new_url).await?;
let document = Html::parse_document(&content);
Ok(document)
}
async fn waf_bypass(&self, target_url: &str) -> Result<String, Error> {
async fn waf_bypass(&self, target_url: &str) -> AppResult<String> {
let script_content = self.get_script_content(target_url).await?;
if script_content.is_empty() {
return Err(Error::Message("waf bypass script not found".to_string()).into());
@@ -355,12 +356,12 @@ impl AVDPlugin {
_window: Value,
_document: Value,
location: Value,
) -> Result<String, Error> {
) -> AppResult<String> {
let runtime = Runtime::new()
.map_err(|e| Error::Message(format!("execution script runtime new error {}", e)))?;
let context = Context::full(&runtime)
.map_err(|e| Error::Message(format!("execution script context full error {}", e)))?;
context.with(|ctx| -> Result<String,Error> {
context.with(|ctx| -> AppResult<String> {
let href = location["href"].as_str().unwrap_or("");
let parsed_url = Url::parse(href)
.map_err(|e| Error::Message(format!("url parsing error {}", e)))?;
@@ -479,7 +480,7 @@ impl AVDPlugin {
})
}
async fn get_script_content(&self, target_url: &str) -> Result<String, Error> {
async fn get_script_content(&self, target_url: &str) -> AppResult<String> {
let origin_content = self.http_client.get_html_content(target_url).await?;
let script_regex = Regex::new(SCRIPT_REGEXP)
.map_err(|e| Error::Message(format!("avd get script regex parse error {}", e)))?;

7
src/output/plugins/kev.rs
View File

@@ -1,10 +1,11 @@
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use serde::{Deserialize, Serialize};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -25,7 +26,7 @@ pub struct KevPlugin {
}
impl KevPlugin {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> Result<KevPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<KevPlugin> {
let http_client = HttpClient::try_new()?;
let kv = KevPlugin {
name: "KevPlugin".to_string(),
@@ -53,7 +54,7 @@ impl VulnPlugin for KevPlugin {
self.link.to_string()
}
async fn update(&self, page_limit: i32) -> Result<(), Error> {
async fn update(&self, page_limit: i32) -> AppResult<()> {
let kev_list_resp: KevResp = self
.http_client
.get_json(KEV_URL)

7
src/output/plugins/mod.rs
View File

@@ -7,14 +7,13 @@ pub mod ti;
use async_trait::async_trait;
use dashmap::DashMap;
use error_stack::Result;
use lazy_static::lazy_static;
use mea::mpsc::UnboundedSender;
use std::sync::Arc;
use crate::{
AppResult,
domain::models::vuln_information::CreateVulnInformation,
errors::Error,
output::plugins::{
avd::AVDPlugin, kev::KevPlugin, oscs::OscsPlugin, seekbug::SeekBugPlugin,
threatbook::ThreatBookPlugin, ti::TiPlugin,
@@ -25,7 +24,7 @@ lazy_static! {
static ref PLUGINS: Arc<DashMap<String, Box<dyn VulnPlugin>>> = Arc::new(DashMap::new());
}
pub fn init(sender: UnboundedSender<CreateVulnInformation>) -> Result<(), Error> {
pub fn init(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<()> {
KevPlugin::try_new(sender.clone())?;
AVDPlugin::try_new(sender.clone())?;
OscsPlugin::try_new(sender.clone())?;
@@ -40,7 +39,7 @@ pub trait VulnPlugin: Send + Sync + 'static {
fn get_name(&self) -> String;
fn get_display_name(&self) -> String;
fn get_link(&self) -> String;
async fn update(&self, page_limit: i32) -> Result<(), Error>;
async fn update(&self, page_limit: i32) -> AppResult<()>;
}
pub fn register_plugin(name: String, plugin: Box<dyn VulnPlugin>) {

17
src/output/plugins/oscs.rs
View File

@@ -1,10 +1,11 @@
use async_trait::async_trait;
use chrono::{DateTime, FixedOffset};
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use serde::{Deserialize, Serialize};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -40,7 +41,7 @@ impl VulnPlugin for OscsPlugin {
self.link.to_string()
}
async fn update(&self, page_limit: i32) -> Result<(), Error> {
async fn update(&self, page_limit: i32) -> AppResult<()> {
let mut page_count = self.get_page_count().await?;
if page_count > page_limit {
page_count = page_limit;
@@ -53,7 +54,7 @@ impl VulnPlugin for OscsPlugin {
}
impl OscsPlugin {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> Result<OscsPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<OscsPlugin> {
let http_client = HttpClient::try_new()?;
let oscs = OscsPlugin {
name: "OscsPlugin".to_string(),
@@ -66,7 +67,7 @@ impl OscsPlugin {
Ok(oscs)
}
pub async fn get_list_resp(&self, page: i32, per_page: i32) -> Result<OscsListResp, Error> {
pub async fn get_list_resp(&self, page: i32, per_page: i32) -> AppResult<OscsListResp> {
let params = serde_json::json!({
"page": page,
"per_page": per_page,
@@ -81,7 +82,7 @@ impl OscsPlugin {
Ok(oscs_list_resp)
}
pub async fn get_page_count(&self) -> Result<i32, Error> {
pub async fn get_page_count(&self) -> AppResult<i32> {
let oscs_list_resp = self
.get_list_resp(OSCS_PAGE_DEFAULT, OSCS_PER_PAGE_DEFAULT)
.await?;
@@ -99,7 +100,7 @@ impl OscsPlugin {
Ok(page_count)
}
pub async fn parse_page(&self, page: i32) -> Result<(), Error> {
pub async fn parse_page(&self, page: i32) -> AppResult<()> {
let oscs_list_resp = self.get_list_resp(page, OSCS_PAGE_SIZE).await?;
for item in oscs_list_resp.data.data {
let mut tags = Vec::new();
@@ -126,7 +127,7 @@ impl OscsPlugin {
Ok(())
}
pub async fn parse_detail(&self, mps: &str) -> Result<CreateVulnInformation, Error> {
pub async fn parse_detail(&self, mps: &str) -> AppResult<CreateVulnInformation> {
let detail = self.get_detail_resp(mps).await?;
if detail.code != 200 || !detail.success || detail.data.is_empty() {
return Err(Error::Message(format!("oscs get: {} detail error", mps)).into());
@@ -179,7 +180,7 @@ impl OscsPlugin {
}
}
async fn get_detail_resp(&self, mps: &str) -> Result<OscsDetailResp, Error> {
async fn get_detail_resp(&self, mps: &str) -> AppResult<OscsDetailResp> {
let params = serde_json::json!({
"vuln_no": mps,
});

25
src/output/plugins/seekbug.rs
View File

@@ -1,9 +1,10 @@
use async_trait::async_trait;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use scraper::{ElementRef, Html, Selector};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -35,7 +36,7 @@ impl VulnPlugin for SeekBugPlugin {
self.link.to_string()
}
async fn update(&self, page_limit: i32) -> Result<(), Error> {
async fn update(&self, page_limit: i32) -> AppResult<()> {
let mut page_count = self.get_page_count().await?;
if page_count > page_limit {
page_count = page_limit;
@@ -48,7 +49,7 @@ impl VulnPlugin for SeekBugPlugin {
}
impl SeekBugPlugin {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> Result<SeekBugPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<SeekBugPlugin> {
let http_client = HttpClient::try_new()?;
let seebug = SeekBugPlugin {
name: "SeeBugPlugin".to_string(),
@@ -61,7 +62,7 @@ impl SeekBugPlugin {
Ok(seebug)
}
pub async fn get_page_count(&self) -> Result<i32, Error> {
pub async fn get_page_count(&self) -> AppResult<i32> {
let document = self.get_document(SEEBUG_LIST_URL).await?;
let selector = Selector::parse("ul.pagination li a")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
@@ -78,7 +79,7 @@ impl SeekBugPlugin {
Ok(total)
}
pub async fn parse_page(&self, page: i32) -> Result<(), Error> {
pub async fn parse_page(&self, page: i32) -> AppResult<()> {
let url = format!("{}?page={}", SEEBUG_LIST_URL, page);
let document = self.get_document(&url).await?;
let selector = Selector::parse(".sebug-table tbody tr")
@@ -169,13 +170,13 @@ impl SeekBugPlugin {
}
}
async fn get_document(&self, url: &str) -> Result<Html, Error> {
async fn get_document(&self, url: &str) -> AppResult<Html> {
let content = self.http_client.get_html_content(url).await?;
let document = Html::parse_document(&content);
Ok(document)
}
fn get_href(&self, el: ElementRef) -> Result<(String, String), Error> {
fn get_href(&self, el: ElementRef) -> AppResult<(String, String)> {
let selector = Selector::parse("td a")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let a_element = el
@@ -193,7 +194,7 @@ impl SeekBugPlugin {
Ok((href.to_owned(), unique_key.to_owned()))
}
fn get_disclosure(&self, el: ElementRef) -> Result<String, Error> {
fn get_disclosure(&self, el: ElementRef) -> AppResult<String> {
let selector = Selector::parse("td")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let disclosure = el
@@ -205,7 +206,7 @@ impl SeekBugPlugin {
Ok(disclosure)
}
fn get_severity_title(&self, el: ElementRef) -> Result<String, Error> {
fn get_severity_title(&self, el: ElementRef) -> AppResult<String> {
let selector = Selector::parse("td div")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let td_element = el
@@ -220,7 +221,7 @@ impl SeekBugPlugin {
Ok(severity_title.to_owned())
}
fn get_title(&self, el: ElementRef) -> Result<String, Error> {
fn get_title(&self, el: ElementRef) -> AppResult<String> {
let selector = Selector::parse("td a[class='vul-title']")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let title = el
@@ -231,7 +232,7 @@ impl SeekBugPlugin {
Ok(title)
}
fn get_cve_id(&self, el: ElementRef) -> Result<String, Error> {
fn get_cve_id(&self, el: ElementRef) -> AppResult<String> {
let selector = Selector::parse("td i[class='fa fa-id-card ']")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let cve_ids = el
@@ -252,7 +253,7 @@ impl SeekBugPlugin {
Ok(cve_ids.to_string())
}
fn get_tag(&self, el: ElementRef) -> Result<String, Error> {
fn get_tag(&self, el: ElementRef) -> AppResult<String> {
let selector = Selector::parse("td .fa-file-text-o")
.map_err(|err| Error::Message(format!("parse html error {}", err)))?;
let tag = el

9
src/output/plugins/threatbook.rs
View File

@@ -1,10 +1,11 @@
use async_trait::async_trait;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use reqwest::header;
use serde::{Deserialize, Serialize};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -37,7 +38,7 @@ impl VulnPlugin for ThreatBookPlugin {
self.link.to_string()
}
async fn update(&self, _page_limit: i32) -> Result<(), Error> {
async fn update(&self, _page_limit: i32) -> AppResult<()> {
let home_page_resp: ThreadBookHomePage = self
.http_client
.get_json(HOME_PAGE_URL)
@@ -94,9 +95,7 @@ impl VulnPlugin for ThreatBookPlugin {
}
impl ThreatBookPlugin {
pub fn try_new(
sender: UnboundedSender<CreateVulnInformation>,
) -> Result<ThreatBookPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<ThreatBookPlugin> {
let mut headers: reqwest::header::HeaderMap = header::HeaderMap::new();
headers.insert("Referer", header::HeaderValue::from_static(LINK));
headers.insert(

13
src/output/plugins/ti.rs
View File

@@ -1,10 +1,11 @@
use async_trait::async_trait;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedSender;
use reqwest::header;
use serde::{Deserialize, Serialize};
use crate::{
AppResult,
domain::models::vuln_information::{CreateVulnInformation, Severity},
errors::Error,
output::plugins::{VulnPlugin, register_plugin},
@@ -36,13 +37,13 @@ impl VulnPlugin for TiPlugin {
self.link.to_string()
}
async fn update(&self, _page_limit: i32) -> Result<(), Error> {
async fn update(&self, _page_limit: i32) -> AppResult<()> {
self.get_vuln_infos().await
}
}
impl TiPlugin {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> Result<TiPlugin, Error> {
pub fn try_new(sender: UnboundedSender<CreateVulnInformation>) -> AppResult<TiPlugin> {
let mut headers = header::HeaderMap::new();
headers.insert(
"Referer",
@@ -65,7 +66,7 @@ impl TiPlugin {
Ok(ti)
}
pub async fn get_vuln_infos(&self) -> Result<(), Error> {
pub async fn get_vuln_infos(&self) -> AppResult<()> {
let ti_one_day_resp = self.get_ti_one_day_resp().await?;
for detail in ti_one_day_resp.data.key_vuln_add {
let tags = self.get_tags(detail.tag);
@@ -95,7 +96,7 @@ impl TiPlugin {
Ok(())
}
pub async fn get_ti_one_day_resp(&self) -> Result<TiOneDayResp, Error> {
pub async fn get_ti_one_day_resp(&self) -> AppResult<TiOneDayResp> {
let resp: TiOneDayResp = self
.http_client
.get_json(ONE_URL)
@@ -141,9 +142,7 @@ pub struct Data {
pub vuln_update_count: i32,
pub key_vuln_add_count: i32,
pub poc_exp_add_count: i32,
// pub patch_add_count: i32,
pub key_vuln_add: Vec<TiVulnDetail>,
// pub poc_exp_add: Vec<TiVulnDetail>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]

8
src/output/push/ding_bot.rs
View File

@@ -2,11 +2,11 @@ use std::time::SystemTime;
use async_trait::async_trait;
use base64::{Engine, prelude::BASE64_STANDARD};
use error_stack::Result;
use reqwest::header;
use serde::{Deserialize, Serialize};
use crate::{
AppResult,
errors::Error,
output::push::MessageBot,
utils::{http_client::HttpClient, util::calc_hmac_sha256},
@@ -24,7 +24,7 @@ pub struct DingBot {
#[async_trait]
impl MessageBot for DingBot {
async fn push_markdown(&self, title: String, msg: String) -> Result<(), Error> {
async fn push_markdown(&self, title: String, msg: String) -> AppResult<()> {
let msg = msg.replace("\n\n", "\n\n&nbsp;\n");
let message = serde_json::json!({
"msgtype": MSG_TYPE,
@@ -62,7 +62,7 @@ impl MessageBot for DingBot {
}
impl DingBot {
pub fn try_new(access_token: String, secret_token: String) -> Result<Self, Error> {
pub fn try_new(access_token: String, secret_token: String) -> AppResult<Self> {
let mut headers = header::HeaderMap::new();
headers.insert("Accept-Charset", header::HeaderValue::from_static("utf8"));
let http_client = HttpClient::try_new_with_headers(headers)?;
@@ -73,7 +73,7 @@ impl DingBot {
})
}
pub fn generate_sign(&self) -> Result<Sign, Error> {
pub fn generate_sign(&self) -> AppResult<Sign> {
let timestamp = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_err(|e| Error::Message(format!("get timestamp failed: {}", e)))?

8
src/output/push/mod.rs
View File

@@ -2,12 +2,12 @@ use std::sync::Arc;
use async_trait::async_trait;
use dashmap::DashMap;
use error_stack::Result;
use lazy_static::lazy_static;
use serde_json::Value;
use crate::{
domain::models::vuln_information::VulnInformation, errors::Error, utils::util::render_string,
AppResult, domain::models::vuln_information::VulnInformation, errors::Error,
utils::util::render_string,
};
pub mod ding_bot;
@@ -18,7 +18,7 @@ lazy_static! {
#[async_trait]
pub trait MessageBot: Send + Sync {
async fn push_markdown(&self, title: String, msg: String) -> Result<(), Error>;
async fn push_markdown(&self, title: String, msg: String) -> AppResult<()>;
}
const VULN_INFO_MSG_TEMPLATE: &str = r####"
@@ -47,7 +47,7 @@ const VULN_INFO_MSG_TEMPLATE: &str = r####"
const MAX_REFERENCE_LENGTH: usize = 8;
pub fn reader_vulninfo(mut vuln: VulnInformation) -> Result<String, Error> {
pub fn reader_vulninfo(mut vuln: VulnInformation) -> AppResult<String> {
if vuln.reference_links.len() > MAX_REFERENCE_LENGTH {
vuln.reference_links = vuln.reference_links[..MAX_REFERENCE_LENGTH].to_vec();
}

15
src/output/scheduler.rs
View File

@@ -1,11 +1,12 @@
use std::{str::FromStr, sync::Arc, time::Instant};
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use tokio::task::JoinSet;
use tokio_cron_scheduler::JobScheduler;
use uuid::Uuid;
use crate::{
AppResult,
errors::Error,
output::{
db::{pg::Pg, sync_data_task::SyncDataTaskDao},
@@ -19,7 +20,7 @@ pub struct Scheduler {
}
impl Scheduler {
pub async fn try_new(pg: Pg) -> Result<Self, Error> {
pub async fn try_new(pg: Pg) -> AppResult<Self> {
let sched = JobScheduler::new()
.await
.change_context_lazy(|| Error::Message("Failed to create scheduler".to_string()))?;
@@ -29,7 +30,7 @@ impl Scheduler {
})
}
fn create_job(&self, interval_minutes: i32) -> Result<tokio_cron_scheduler::Job, Error> {
fn create_job(&self, interval_minutes: i32) -> AppResult<tokio_cron_scheduler::Job> {
let cron_syntax = format!("0 */{} * * * *", interval_minutes);
log::debug!("Creating job with cron syntax: {}", cron_syntax);
let job =
@@ -49,7 +50,7 @@ impl Scheduler {
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
task: &crate::domain::models::sync_data_task::SyncDataTask,
job: tokio_cron_scheduler::Job,
) -> Result<(), Error> {
) -> AppResult<()> {
let new_job_id = self.sched.add(job).await.change_context_lazy(|| {
Error::Message("Failed to add new job to scheduler".to_string())
})?;
@@ -57,7 +58,7 @@ impl Scheduler {
Ok(())
}
async fn remove_existing_job(&self, job_id: &str) -> Result<Option<Uuid>, Error> {
async fn remove_existing_job(&self, job_id: &str) -> AppResult<Option<Uuid>> {
let job_id = Uuid::from_str(job_id)
.change_context_lazy(|| Error::Message("Failed to parse job ID".to_string()))?;
self.sched.remove(&job_id).await.change_context_lazy(|| {
@@ -70,7 +71,7 @@ impl Scheduler {
Ok(Some(job_id))
}
pub async fn update(&self, id: i64) -> Result<(), Error> {
pub async fn update(&self, id: i64) -> AppResult<()> {
let mut tx =
self.pg.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())
@@ -96,7 +97,7 @@ impl Scheduler {
}
}
pub async fn init_from_db(self) -> Result<Self, Error> {
pub async fn init_from_db(self) -> AppResult<Self> {
let mut tx =
self.pg.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())

9
src/output/worker.rs
View File

@@ -1,9 +1,10 @@
use std::sync::Arc;
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use mea::mpsc::UnboundedReceiver;
use crate::{
AppResult,
domain::models::vuln_information::CreateVulnInformation,
errors::Error,
output::{
@@ -26,7 +27,7 @@ impl Worker {
}
}
pub async fn run(&mut self) -> Result<(), Error> {
pub async fn run(&mut self) -> AppResult<()> {
while let Some(req) = self.receiver.recv().await {
match self.store(req).await {
Err(e) => {
@@ -43,7 +44,7 @@ impl Worker {
Ok(())
}
pub async fn ding_bot_push(&self, id: i64) -> Result<(), Error> {
pub async fn ding_bot_push(&self, id: i64) -> AppResult<()> {
log::info!("ding bot push start! id: {}", id);
let mut tx =
self.pg.pool.begin().await.change_context_lazy(|| {
@@ -81,7 +82,7 @@ impl Worker {
Ok(())
}
pub async fn store(&self, mut req: CreateVulnInformation) -> Result<(i64, bool), Error> {
pub async fn store(&self, mut req: CreateVulnInformation) -> AppResult<(i64, bool)> {
let mut tx =
self.pg.pool.begin().await.change_context_lazy(|| {
Error::Message("failed to begin transaction".to_string())

8
src/utils/auth/jwt.rs
View File

@@ -1,4 +1,4 @@
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use jsonwebtoken::{
Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation, decode, encode,
get_current_timestamp,
@@ -6,7 +6,7 @@ use jsonwebtoken::{
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::errors::Error;
use crate::{AppResult, errors::Error};
const JWT_ALGORITHM: Algorithm = Algorithm::HS512;
@@ -42,7 +42,7 @@ impl JWT {
expiration: u64,
user_id: i64,
claims: Map<String, Value>,
) -> Result<String, Error> {
) -> AppResult<String> {
let exp = get_current_timestamp().saturating_add(expiration);
let claims = UserClaims {
@@ -62,7 +62,7 @@ impl JWT {
Ok(token)
}
pub fn validate(&self, token: &str) -> Result<TokenData<UserClaims>, Error> {
pub fn validate(&self, token: &str) -> AppResult<TokenData<UserClaims>> {
let mut validate = Validation::new(self.algorithm);
validate.leeway = 0;

14
src/utils/http_client.rs
View File

@@ -1,7 +1,7 @@
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use reqwest::header::{self, HeaderMap};
use crate::errors::Error;
use crate::{AppResult, errors::Error};
#[derive(Debug, Clone)]
pub struct HttpClient {
@@ -9,7 +9,7 @@ pub struct HttpClient {
}
impl HttpClient {
pub fn try_new_with_headers(mut headers: HeaderMap) -> Result<Self, Error> {
pub fn try_new_with_headers(mut headers: HeaderMap) -> AppResult<Self> {
headers.insert("User-Agent", header::HeaderValue::from_static("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"));
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
@@ -21,7 +21,7 @@ impl HttpClient {
http_client: client,
})
}
pub fn try_new() -> Result<Self, Error> {
pub fn try_new() -> AppResult<Self> {
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.danger_accept_invalid_certs(true)
@@ -32,7 +32,7 @@ impl HttpClient {
})
}
pub async fn get_json(&self, url: &str) -> Result<reqwest::Response, Error> {
pub async fn get_json(&self, url: &str) -> AppResult<reqwest::Response> {
let content = self
.http_client
.get(url)
@@ -41,7 +41,7 @@ impl HttpClient {
.change_context_lazy(|| Error::Message("Failed to send HTTP request".to_string()))?;
Ok(content)
}
pub async fn get_html_content(&self, url: &str) -> Result<String, Error> {
pub async fn get_html_content(&self, url: &str) -> AppResult<String> {
let content = self
.http_client
.get(url)
@@ -56,7 +56,7 @@ impl HttpClient {
Ok(content)
}
pub async fn post_json<Body>(&self, url: &str, body: &Body) -> Result<reqwest::Response, Error>
pub async fn post_json<Body>(&self, url: &str, body: &Body) -> AppResult<reqwest::Response>
where
Body: serde::Serialize,
{

6
src/utils/password_hash.rs
View File

@@ -2,11 +2,11 @@ use argon2::{
Argon2, Params, PasswordHash, PasswordHasher, PasswordVerifier,
password_hash::{SaltString, rand_core::OsRng},
};
use error_stack::{Result, ResultExt};
use error_stack::ResultExt;
use crate::errors::Error;
use crate::{AppResult, errors::Error};
pub fn compute_password_hash(password: &str) -> Result<String, Error> {
pub fn compute_password_hash(password: &str) -> AppResult<String> {
let arg2 = Argon2::new(
argon2::Algorithm::Argon2id,
argon2::Version::V0x13,

6
src/utils/search.rs
View File

@@ -1,6 +1,6 @@
use regex::Regex;
use crate::{errors::Error, utils::util::get_last_year_data};
use crate::{AppResult, errors::Error, utils::util::get_last_year_data};
pub async fn search_github_poc(cve_id: &str) -> Vec<String> {
let mut res = Vec::new();
@@ -20,7 +20,7 @@ pub async fn search_github_poc(cve_id: &str) -> Vec<String> {
res
}
pub async fn search_nuclei_pr(cve_id: &str) -> Result<Vec<String>, Error> {
pub async fn search_nuclei_pr(cve_id: &str) -> AppResult<Vec<String>> {
log::info!("search nuclei PR of {}", cve_id);
let page = octocrab::instance()
.pulls("projectdiscovery", "nuclei-templates")
@@ -45,7 +45,7 @@ pub async fn search_nuclei_pr(cve_id: &str) -> Result<Vec<String>, Error> {
Ok(links)
}
pub async fn search_github_repo(cve_id: &str) -> Result<Vec<String>, Error> {
pub async fn search_github_repo(cve_id: &str) -> AppResult<Vec<String>> {
log::info!("search github repo of {}", cve_id);
let last_year = get_last_year_data();
let query = format!(

7
src/utils/util.rs
View File

@@ -1,12 +1,11 @@
use chrono::{DateTime, Duration, Local, Months, NaiveDate, Utc};
use error_stack::Result;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use tera::{Context, Tera};
use crate::errors::Error;
use crate::{AppResult, errors::Error};
pub fn timestamp_to_date(timestamp: i64) -> Result<String, Error> {
pub fn timestamp_to_date(timestamp: i64) -> AppResult<String> {
let dt = DateTime::from_timestamp_millis(timestamp);
if let Some(dt) = dt {
return Ok(dt.format("%Y-%m-%d").to_string());
@@ -38,7 +37,7 @@ pub fn get_last_year_data() -> String {
last_year.format("%Y-%m-%d").to_string()
}
pub fn check_over_two_week(date: &str) -> Result<bool, Error> {
pub fn check_over_two_week(date: &str) -> AppResult<bool> {
let target_date = NaiveDate::parse_from_str(date, "%Y-%m-%d")
.map_err(|e| Error::Message(format!("parse date error: {:?}", e)))?;
let now = Utc::now().naive_utc().date();