From 56275151d9ba3d95d99419a64059101815c1a837 Mon Sep 17 00:00:00 2001 From: ParrotXray Date: Mon, 11 May 2026 14:26:14 +0800 Subject: [PATCH] Feat/improve csv record (#2) * chore: new branch * wip: todo 1 * fix: flush CSV writer after each record instead of only on shutdown --- TODO | 14 +++- net-guardia-frontend | 2 +- net-guardia/build.rs | 3 + net-guardia/src/core/infrastructure/mod.rs | 12 +-- .../src/detection/ml/traffic_logger.rs | 81 ++++++++++++++++--- net-guardia/src/model/error/ml.rs | 16 ++++ net-guardia/src/model/log/ml.rs | 3 + 7 files changed, 113 insertions(+), 18 deletions(-) diff --git a/TODO b/TODO index 687e239..12f9084 100644 --- a/TODO +++ b/TODO @@ -1,7 +1,17 @@ -1. 加入 .csv 的紀錄資料夾，實現 .csv 過一天後自動歸檔成 xxx-yyyy-oo-zz.csv +1. 加入 .csv 的紀錄資料夾，實現 .csv 過一天後自動歸檔成 xxx-yyyy-oo-zz.csv [o] 2. will change tract-onnx engine to ort-tract engine 3. 改善推論效能及速度，在大量資料時 4. 實現 ML + RULE 的共用 HashMap 實現共同決策結果 5. 代碼最佳化，檢查是否除了 build.rs 以外有無 .unwarp() and eprintln() [x] 6. 完成前端 detection 頁面 -7. 使用 sqllite 實現帳號系統、白黑名單永久記錄 \ No newline at end of file +7. 使用 sqllite 實現帳號系統、白黑名單永久記錄 + +Suricata +提升偵測準確率： +pcre 很多規則依賴正則，沒有會漏掉很多威脅 +flow:established 減少掃握手封包的無意義處理 +Flowbits 複雜攻擊鏈的偵測（例如先偵測掃描再偵測滲透） + +減少假陽性： +$HOME_NET / $EXTERNAL_NET 區分內外網，很多規則依賴這個 +flow:established 順便過濾掉不完整連線 \ No newline at end of file diff --git a/net-guardia-frontend b/net-guardia-frontend index c7f5045..671d42e 160000 --- a/net-guardia-frontend +++ b/net-guardia-frontend @@ -1 +1 @@ -Subproject commit c7f50458778c70a124ac07d2356a37c460b82745 +Subproject commit 671d42e41aa7309988f0e610df1be88ac96dde9c diff --git a/net-guardia/build.rs b/net-guardia/build.rs index 85cc70d..0c50203 100644 --- a/net-guardia/build.rs +++ b/net-guardia/build.rs @@ -479,6 +479,8 @@ fn build_vectorscan_db() { let artifact_dir = &manifest_dir.join("static").join("artifacts"); let rules_dir = &manifest_dir.join("static").join("rules"); + let csv_dir = &manifest_dir.parent().unwrap().join("records"); + println!("cargo:rerun-if-changed={}", rules_dir.display()); if let Ok(rd) = fs::read_dir(&rules_dir) { for entry in rd.flatten() { @@ -496,6 +498,7 @@ fn build_vectorscan_db() { println!("cargo:rustc-env=RULES_DB_PATH={}", out_dir.display()); println!("cargo:rustc-env=ARTIFACTCS_PATH={}", artifact_dir.display()); + println!("cargo:rustc-env=CSV_RECORD_PATH={}", csv_dir.display()); let total_contents: usize = sigs.iter().map(|s| s.chain.len()).sum(); println!( diff --git a/net-guardia/src/core/infrastructure/mod.rs b/net-guardia/src/core/infrastructure/mod.rs index c84e027..9818dbe 100644 --- a/net-guardia/src/core/infrastructure/mod.rs +++ b/net-guardia/src/core/infrastructure/mod.rs @@ -3,7 +3,7 @@ pub mod health; pub mod geoip; pub mod ml_alert; - +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -39,15 +39,15 @@ impl AppServices { let ml_models = Arc::new(MLModels::load_models(&app_config, &inference_config)?); let ml_alert = Arc::new(MLAlert::new()); - // TODO: Need Edit let traffic_logger = if app_config.traffic_logging_mode { - let csv_path = app_config.traffic_log_csv_path.clone(); + let dir = PathBuf::from(env!("CSV_RECORD_PATH")); + let basename = app_config.traffic_log_csv_path.trim_end_matches(".csv").to_string(); let mut header = vec!["Source IP".to_string(), "Destination IP".to_string(), "Timestamp".to_string()]; header.extend(FlowFeatures::all_feature_names_owned()); header.push("Label".to_string()); - let logger = TrafficLogger::new(&csv_path, header) - .map_err(|e| MiscError::TrafficLogCreateError(csv_path.clone(), e.to_string()))?; - log!(SystemLog::TrafficLoggingEnabled(csv_path)); + let logger = TrafficLogger::new(&dir, &basename, header) + .map_err(|e| MiscError::TrafficLogCreateError(dir.display().to_string(), e.to_string()))?; + log!(SystemLog::TrafficLoggingEnabled(format!("{}/{}.csv", dir.display().to_string(), basename))); Some(Arc::new(logger)) } else { None diff --git a/net-guardia/src/detection/ml/traffic_logger.rs b/net-guardia/src/detection/ml/traffic_logger.rs index cf7979a..d7506ec 100644 --- a/net-guardia/src/detection/ml/traffic_logger.rs +++ b/net-guardia/src/detection/ml/traffic_logger.rs @@ -1,35 +1,98 @@ -use std::fs::OpenOptions; +use std::fs::{self, OpenOptions}; use std::io::{BufWriter, Write}; +use std::path::{PathBuf, Path}; use std::thread; +use std::time::Duration; -use tracing; +use chrono::Local; +use crossbeam::channel::{bounded, Sender, TrySendError, RecvTimeoutError}; +use macros::log; -use crossbeam::channel::{bounded, Sender, TrySendError}; +use crate::model::error::ml::MLError; +use crate::model::log::ml::MLLog; pub struct TrafficLogger { sender: Sender>, } impl TrafficLogger { - pub fn new(csv_path: &str, header: Vec) -> Result { + pub fn new(dir: impl AsRef, basename: &str, header: Vec) -> Result { + let dir = PathBuf::from(dir.as_ref()); + fs::create_dir_all(&dir)?; + + let current_path = dir.join(format!("{}.csv", basename)); let file = OpenOptions::new() .create(true) .write(true) .truncate(true) - .open(csv_path)?; + .open(¤t_path)?; let mut writer = BufWriter::new(file); writeln!(writer, "{}", header.join(","))?; writer.flush()?; let (sender, receiver) = bounded::>(65536); + let basename = basename.to_string(); thread::Builder::new() .name("traffic-logger".to_string()) .spawn(move || { - for record in receiver { - if let Err(e) = writeln!(writer, "{}", record.join(",")) { - tracing::warn!("[traffic-logger] write error: {}", e); + let mut writer = writer; + let mut current_date = Local::now().date_naive(); + + loop { + match receiver.recv_timeout(Duration::from_secs(1)) { + Ok(record) => { + let today = Local::now().date_naive(); + if today != current_date { + let _ = writer.flush(); + drop(writer); + + let archive_name = format!("{}-{}.csv", basename, current_date.format("%Y-%m-%d")); + let archive_path = dir.join(&archive_name); + if let Err(e) = fs::rename(¤t_path, &archive_path) { + log!(MLError::TrafficLogArchiveFailed( + archive_path.display().to_string(), + e.to_string() + )); + } else { + log!(MLLog::TrafficLogArchived(archive_path.display().to_string())); + } + current_date = today; + + writer = match OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(¤t_path) + { + Ok(new_file) => { + let mut w = BufWriter::new(new_file); + if let Err(e) = writeln!(w, "{}", header.join(",")) { + log!(MLError::TrafficLogHeaderFailed(e.to_string())); + } + w + } + Err(e) => { + log!(MLError::TrafficLogOpenFailed( + current_path.display().to_string(), + e.to_string() + )); + return; + } + }; + } + + if let Err(e) = writeln!(writer, "{}", record.join(",")) { + log!(MLError::TrafficLogWriteFailed(e.to_string())); + } + } + Err(RecvTimeoutError::Timeout) => { + let _ = writer.flush(); + } + Err(RecvTimeoutError::Disconnected) => { + break; + } } } let _ = writer.flush(); @@ -45,4 +108,4 @@ impl TrafficLogger { Err(TrySendError::Disconnected(_)) => {} } } -} +} \ No newline at end of file diff --git a/net-guardia/src/model/error/ml.rs b/net-guardia/src/model/error/ml.rs index 8415541..3b547d2 100644 --- a/net-guardia/src/model/error/ml.rs +++ b/net-guardia/src/model/error/ml.rs @@ -27,5 +27,21 @@ traceable! { #[no_source] #[error("Inference flow buffers mutex poisoned; skipping operation")] InferenceLockPoisoned => tracing::Level::ERROR, + + #[no_source] + #[error("Failed to archive CSV '{path}': {reason}")] + TrafficLogArchiveFailed { path: String, reason: String } => tracing::Level::WARN, + + #[no_source] + #[error("Failed to write CSV header: {reason}")] + TrafficLogHeaderFailed { reason: String } => tracing::Level::WARN, + + #[no_source] + #[error("Failed to open new CSV file '{path}': {reason}")] + TrafficLogOpenFailed { path: String, reason: String } => tracing::Level::ERROR, + + #[no_source] + #[error("Failed to write CSV row: {reason}")] + TrafficLogWriteFailed { reason: String } => tracing::Level::WARN, } } \ No newline at end of file diff --git a/net-guardia/src/model/log/ml.rs b/net-guardia/src/model/log/ml.rs index f9d5c64..6fc1c83 100644 --- a/net-guardia/src/model/log/ml.rs +++ b/net-guardia/src/model/log/ml.rs @@ -69,5 +69,8 @@ loggable! { #[error("Failed to parse packet (length: {len})")] ParsePacketFailed { len: usize } => tracing::Level::INFO, + #[error("CSV archived to: {path}")] + TrafficLogArchived { path: String } => tracing::Level::INFO, + } } \ No newline at end of file