ParrotXray/Mantis
1
0
Fork 0
mirror of https://github.com/ParrotXray/Mantis.git synced 2026-08-24 19:00:27 +09:00
Code Issues Packages Projects Releases Wiki Activity

feat/ml-rule-fusion (#6)

Browse Source 
* wip

* wip

* feat: auto-select native ORT or ort-tract backend based on .so presence and ML+Rule fusion decision layer

* wip
This commit is contained in:
ParrotXray 2026-05-19 15:56:43 +08:00 committed by GitHub
parent e686449a98
commit 99c2d74737
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
22 changed files with 877 additions and 440 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
Cargo.lock generated
View File

@ -741,7 +741,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
dependencies = [
"glob",
"libc",
"libloading",
"libloading 0.8.9",
]
[[package]]
@ -1707,6 +1707,16 @@ dependencies = [
"windows-link",
]
[[package]]
name = "libloading"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60"
dependencies = [
"cfg-if",
"windows-link",
]
[[package]]
name = "libm"
version = "0.2.15"
@ -2249,6 +2259,7 @@ version = "2.0.0-rc.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133"
dependencies = [
"libloading 0.9.0",
"ndarray 0.17.2",
"ort-sys",
"smallvec",

7
TODO
View File

@ -1,7 +1,7 @@
1. .csv 一開始文件格式為 traffic-yyyy-oo-zz.csv, 換天後直接變日期創新檔，例如今天啟動今天為 traffic-2026-05-11.csv 隔天變成 traffic-2026-05-12.csv [x]
2. will change tract-onnx engine to ort-tract engine [x]
3. 改善推論效能及速度，在大量資料時 [x]
4. 實現 ML + RULE 的共用 HashMap 實現共同決策結果
4. 實現 ML + RULE 的共用 HashMap 實現共同決策結果 [x]
5. 代碼最佳化，檢查是否除了 build.rs 以外有無 .unwarp() and eprintln() [x]
6. 完成前端 detection 頁面
7. 使用 sqllite 實現帳號系統、白黑名單永久記錄
@ -36,13 +36,12 @@ PRIORITY BACKLOG (updated 2026-05-18)
Protocols mapped: http(1) http2(2) tls(3) dns(4) ssh(5) smtp(6)
ftp(7) mqtt(8) quic(9).
[ ] ML + Rule fusion decision layer
[x] ML + Rule fusion decision layer
Core differentiator of the project. Currently ML and rule
engine produce independent alerts with no cross-awareness.
Design: shared per-flow HashMap; three fusion modes:
Design: shared per-flow HashMap; two fusion modes:
AND - alert only when both agree
OR - alert when either fires (with source tag)
WEIGHTED - score = rule_hits * w1 + ml_mse * w2 > threshold
Files: detection/ml/engine.rs, detection/rule/rule_engine.rs,
new detection/fusion.rs, model/ml_detection.rs

11
net-guardia/Cargo.toml
View File

@ -52,9 +52,14 @@ nom7 = { version = "7.1", package = "nom" }
nom8 = { version = "8.0", package = "nom" }
aes-gcm = "0.10"
ort-tract = "0.3.0+0.22"
ort = { version = "=2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "alternative-backend"] }
ort-tract = { version = "0.3.0+0.22", optional = true }
ort = { version = "=2.0.0-rc.12", default-features = false, features = ["std", "ndarray"] }
ndarray = "0.17"
[features]
default = ["native-ort-backend"]
tract-backend = ["dep:ort-tract", "ort/alternative-backend"]
native-ort-backend = ["ort/load-dynamic", "ort/api-18"]
#csv = "1.4.0"
#anyhow = "1.0.100"
@ -65,4 +70,4 @@ rusqlite = { version = "0.32", features = ["bundled"] }
[[bin]]
name = "net-guardia"
path = "src/main.rs"
path = "src/main.rs"

796
net-guardia/build.rs
View File

File diff suppressed because it is too large Load Diff

12
net-guardia/src/core/ebpf/mod.rs
View File

@ -15,9 +15,10 @@ use crate::core::ebpf::service::Service;
use crate::core::ebpf::statistics::Statistics;
use crate::core::ebpf::xsk_manager::XskManager;
use crate::core::infrastructure::app_config::AppConfig;
use crate::detection::fusion::FusionEngine;
use crate::detection::ml::engine::Engine;
use crate::model::error::system::SystemError;
use crate::model::error::Error;
use crate::detection::ml::engine::Engine;
pub struct EbpfServices {
pub xsk_manager: Arc<XskManager>,
@ -47,14 +48,17 @@ impl EbpfServices {
Ok(ebpf_services)
}
pub async fn run(self: Arc<Self>, ml_engine: Arc<Engine>) -> Result<(), Error> {
pub async fn run(
self: Arc<Self>,
ml_engine: Arc<Engine>,
fusion_engine: Arc<FusionEngine>,
) -> Result<(), Error> {
let xsk_manager = self.xsk_manager.clone();
let statistics = self.statistics.clone();
xsk_manager.run(Some(ml_engine), &self.shutdowns)?;
xsk_manager.run(Some(ml_engine), Some(fusion_engine), &self.shutdowns)?;
let statistics_shutdown = statistics.run().await;
self.shutdowns.push(statistics_shutdown);
Ok(())
}

24
net-guardia/src/core/ebpf/xsk_manager.rs
View File

@ -17,6 +17,7 @@ use xsk_rs::config::{BindFlags, FrameSize, Interface, LibxdpFlags, QueueSize, So
use xsk_rs::{CompQueue, FillQueue, FrameDesc, RxQueue, Socket, TxQueue, Umem};
use crate::core::infrastructure::app_config::AppConfig;
use crate::detection::fusion::FusionEngine;
use crate::detection::ml::engine::Engine;
use crate::detection::rule::rule_engine::RuleEngine;
use crate::detection::rule::stream_reassembler::StreamReassembler;
@ -53,7 +54,12 @@ impl XskManager {
})
}
pub fn run(&self, engine: Option<Arc<Engine>>, shutdowns: &SegQueue<oneshot::Sender<()>>) -> Result<(), Error> {
pub fn run(
&self,
engine: Option<Arc<Engine>>,
fusion_engine: Option<Arc<FusionEngine>>,
shutdowns: &SegQueue<oneshot::Sender<()>>,
) -> Result<(), Error> {
let config = self.app_config.config.clone();
let combined_queue_count = config.combined_queue_count;
@ -82,6 +88,7 @@ impl XskManager {
Direction::Ingress,
engine.clone(),
rule_engine.clone(),
fusion_engine.clone(),
min_sig,
)?;
@ -93,6 +100,7 @@ impl XskManager {
Direction::Egress,
engine.clone(),
rule_engine.clone(),
fusion_engine.clone(),
min_sig,
)?;
@ -135,6 +143,7 @@ pub struct XskPair {
frame_pool: Arc<Mutex<Vec<FrameDesc>>>,
engine: Option<Arc<Engine>>,
rule_engine: Option<Arc<RuleEngine>>,
fusion_engine: Option<Arc<FusionEngine>>,
min_signature_matches: u32,
}
@ -147,6 +156,7 @@ impl XskPair {
direction: Direction,
engine: Option<Arc<Engine>>,
rule_engine: Option<Arc<RuleEngine>>,
fusion_engine: Option<Arc<FusionEngine>>,
min_signature_matches: u32,
) -> Result<Self, Error> {
let rx_ifname_c = CString::new(rx_ifname).map_err(|_| SystemError::UnknownError)?;
@ -206,6 +216,7 @@ impl XskPair {
frame_pool: Arc::new(Mutex::new(pool_frames)),
engine,
rule_engine,
fusion_engine,
min_signature_matches,
};
@ -340,14 +351,17 @@ impl XskPair {
if let Some(r) = &mut *reassembler {
match r.process(&packet_data) {
Ok(matches) => {
for m in matches {
for m in &matches {
log!(RuleLog::SignatureMatch(
format!("{:?}", self.direction),
m.src,
m.dst,
m.src.clone(),
m.dst.clone(),
m.sid,
m.msg,
m.msg.clone(),
));
if let Some(ref fe) = self.fusion_engine {
fe.record_rule(m);
}
}
}
Err(e) => log!(e),

38
net-guardia/src/core/infrastructure/detection_alert.rs Normal file
View File

@ -0,0 +1,38 @@
use std::sync::Arc;
use tokio::sync::broadcast;
use crate::model::ml_detection::UnifiedAlert;
pub struct DetectionAlert {
broadcast_tx: broadcast::Sender<UnifiedAlert>,
}
impl DetectionAlert {
pub fn new() -> Self {
let (broadcast_tx, _) = broadcast::channel(256);
DetectionAlert { broadcast_tx }
}
pub fn subscribe(&self) -> broadcast::Receiver<UnifiedAlert> {
self.broadcast_tx.subscribe()
}
pub fn broadcast(&self, alert: UnifiedAlert) {
if self.broadcast_tx.receiver_count() > 0 {
let _ = self.broadcast_tx.send(alert);
}
}
pub fn has_subscribers(&self) -> bool {
self.broadcast_tx.receiver_count() > 0
}
}
impl Default for DetectionAlert {
fn default() -> Self {
Self::new()
}
}
pub type SharedDetectionAlert = Arc<DetectionAlert>;

39
net-guardia/src/core/infrastructure/ml_alert.rs
View File

@ -1,39 +0,0 @@
use tokio::sync::broadcast;
use tracing::error;
use crate::model::ml_detection::{AlertMessage, DetectionResult};
pub struct MLAlert {
broadcast_tx: broadcast::Sender<AlertMessage>,
}
impl MLAlert {
pub fn new() -> Self {
let (broadcast_tx, _) = broadcast::channel(100);
MLAlert { broadcast_tx }
}
pub fn subscribe_to_alerts(&self) -> broadcast::Receiver<AlertMessage> {
self.broadcast_tx.subscribe()
}
pub fn broadcast_alert(&self, result: &DetectionResult) {
if self.broadcast_tx.receiver_count() > 0 {
let alert = AlertMessage::from_detection_result(result);
if let Err(e) = self.broadcast_tx.send(alert) {
error!("Failed to broadcast ML alert: {}", e);
}
}
}
pub fn has_subscribers(&self) -> bool {
self.broadcast_tx.receiver_count() > 0
}
}
impl Default for MLAlert {
fn default() -> Self {
Self::new()
}
}

31
net-guardia/src/core/infrastructure/mod.rs
View File

@ -1,34 +1,35 @@
pub mod app_config;
pub mod detection_alert;
pub mod health;
pub mod geoip;
pub mod ml_alert;
use std::path::PathBuf;
use chrono::Local;
use std::sync::Arc;
use std::time::Duration;
use chrono::Local;
use crossbeam::queue::SegQueue;
use macros::log;
use tokio::sync::oneshot;
use crate::core::infrastructure::app_config::AppConfig;
use crate::core::infrastructure::detection_alert::DetectionAlert;
use crate::core::infrastructure::health::SystemHealth;
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::detection::fusion::{FusionEngine, FusionMode};
use crate::detection::ml::config_loader::InferenceConfig;
use crate::detection::ml::engine::Engine;
use crate::detection::ml::feature_extractor::FlowFeatures;
use crate::detection::ml::model_loader::MLModels;
use crate::detection::ml::traffic_logger::TrafficLogger;
use crate::model::error::misc::MiscError;
use crate::model::error::Error;
use crate::model::error::system::SystemError;
use crate::model::log::system::SystemLog;
use crate::detection::ml::traffic_logger::TrafficLogger;
pub struct AppServices {
pub health: Arc<SystemHealth>,
pub ml_alert: Arc<MLAlert>,
pub detection_alert: Arc<DetectionAlert>,
pub fusion_engine: Arc<FusionEngine>,
pub ml_models: Arc<MLModels>,
pub ml_engine: Arc<Engine>,
shutdowns: SegQueue<oneshot::Sender<()>>,
@ -37,9 +38,15 @@ pub struct AppServices {
impl AppServices {
pub fn new(app_config: Arc<AppConfig>, inference_config: Arc<InferenceConfig>) -> Result<Self, Error> {
let health = SystemHealth::new(app_config.clone())?;
let ml_models = Arc::new(MLModels::load_models(&app_config)?);
let ml_alert = Arc::new(MLAlert::new());
let detection_alert = Arc::new(DetectionAlert::new());
let mode = FusionMode::from_str(&app_config.fusion_mode);
let fusion_engine = Arc::new(FusionEngine::new(
mode,
app_config.fusion_window_secs,
detection_alert.clone(),
));
let traffic_logger = if app_config.traffic_logging_mode {
let dir = PathBuf::from(env!("CSV_RECORD_PATH"));
@ -58,7 +65,7 @@ impl AppServices {
let ml_engine = Arc::new(Engine::new(
ml_models.clone(),
inference_config.clone(),
ml_alert.clone(),
fusion_engine.clone(),
app_config.max_concurrent_flows,
app_config.min_packets_for_inference,
app_config.inference_batch_size,
@ -69,10 +76,10 @@ impl AppServices {
app_config.ml_cpu,
));
Ok(Self {
health: Arc::new(health),
ml_alert,
detection_alert,
fusion_engine,
ml_models,
ml_engine,
shutdowns: SegQueue::new(),
@ -99,4 +106,4 @@ impl AppServices {
}
}
}
}
}

10
net-guardia/src/core/system.rs
View File

@ -20,7 +20,7 @@ use crate::model::error::Error;
use crate::model::log::ml::MLLog;
use crate::model::log::system::SystemLog;
use crate::utils::logging::Logging;
use crate::web::api::{control, default, health, misc, ml_alert};
use crate::web::api::{control, default, detection_alert, health, misc};
pub struct System {
pub app_config: Arc<AppConfig>,
@ -82,7 +82,7 @@ impl System {
log!(SystemLog::InitializeComplete);
self.attach_ebpf()?;
ebpf_services.run(app_services.ml_engine.clone()).await?;
ebpf_services.run(app_services.ml_engine.clone(), app_services.fusion_engine.clone()).await?;
app_services.run().await?;
self.run_http_server().await?;
Ok(())
@ -140,7 +140,7 @@ impl System {
let service = self.ebpf_services.service.clone();
let statistics = self.ebpf_services.statistics.clone();
let health = self.app_services.health.clone();
let ml_alert = self.app_services.ml_alert.clone();
let detection_alert = self.app_services.detection_alert.clone();
let port = self.app_config.http_server_bind_port;
HttpServer::new(move || {
let cors = actix_cors::Cors::default()
@ -156,9 +156,9 @@ impl System {
.app_data(web::Data::from(service.clone()))
.app_data(web::Data::from(statistics.clone()))
.app_data(web::Data::from(health.clone()))
.app_data(web::Data::from(ml_alert.clone()))
.app_data(web::Data::from(detection_alert.clone()))
.service(control::initialize())
.service(ml_alert::initialize())
.service(detection_alert::initialize())
.service(health::initialize())
.service(misc::initialize())
.default_service(route().to(default::default_route))

133
net-guardia/src/detection/fusion.rs Normal file
View File

@ -0,0 +1,133 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use crate::core::infrastructure::detection_alert::DetectionAlert;
use crate::model::ml_detection::{DetectionResult, UnifiedAlert};
use crate::model::rule_detection::RuleMatch;
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum FusionMode {
Or,
And,
}
impl FusionMode {
pub fn from_str(s: &str) -> Self {
if s.eq_ignore_ascii_case("and") { FusionMode::And } else { FusionMode::Or }
}
}
struct FusionState {
ml: Option<DetectionResult>,
rule: Option<RuleMatch>,
created_at: Instant,
}
pub struct FusionEngine {
state: Mutex<HashMap<String, FusionState>>,
mode: FusionMode,
window_secs: u64,
alert: Arc<DetectionAlert>,
}
impl FusionEngine {
pub fn new(mode: FusionMode, window_secs: u64, alert: Arc<DetectionAlert>) -> Self {
Self {
state: Mutex::new(HashMap::new()),
mode,
window_secs,
alert,
}
}
pub fn record_ml(&self, result: &DetectionResult) {
let key = ml_key(result);
let Ok(mut map) = self.state.lock() else { return };
let now = Instant::now();
let window = self.window_secs;
map.retain(|_, s| now.duration_since(s.created_at).as_secs() < window);
match self.mode {
FusionMode::Or => {
// Clone any corroborating rule match before mutating the map.
let corroboration = map.get(&key).and_then(|fs| fs.rule.clone());
if let Some(rule) = corroboration {
let alert = UnifiedAlert::from_fusion(result, &rule);
map.remove(&key);
drop(map);
self.alert.broadcast(alert);
return;
}
self.alert.broadcast(UnifiedAlert::from_ml(result));
map.insert(key, FusionState { ml: Some(result.clone()), rule: None, created_at: now });
}
FusionMode::And => {
let maybe_rule = {
let entry = map.entry(key.clone()).or_insert_with(|| FusionState {
ml: None,
rule: None,
created_at: now,
});
entry.ml = Some(result.clone());
entry.rule.clone()
};
if let Some(rule) = maybe_rule {
let alert = UnifiedAlert::from_fusion(result, &rule);
map.remove(&key);
drop(map);
self.alert.broadcast(alert);
}
}
}
}
pub fn record_rule(&self, m: &RuleMatch) {
let key = rule_key(m);
let Ok(mut map) = self.state.lock() else { return };
let now = Instant::now();
let window = self.window_secs;
map.retain(|_, s| now.duration_since(s.created_at).as_secs() < window);
match self.mode {
FusionMode::Or => {
let corroboration = map.get(&key).and_then(|fs| fs.ml.clone());
if let Some(ml) = corroboration {
let alert = UnifiedAlert::from_fusion(&ml, m);
map.remove(&key);
drop(map);
self.alert.broadcast(alert);
return;
}
self.alert.broadcast(UnifiedAlert::from_rule(m));
map.insert(key, FusionState { ml: None, rule: Some(m.clone()), created_at: now });
}
FusionMode::And => {
let maybe_ml = {
let entry = map.entry(key.clone()).or_insert_with(|| FusionState {
ml: None,
rule: None,
created_at: now,
});
entry.rule = Some(m.clone());
entry.ml.clone()
};
if let Some(ml) = maybe_ml {
let alert = UnifiedAlert::from_fusion(&ml, m);
map.remove(&key);
drop(map);
self.alert.broadcast(alert);
}
}
}
}
}
fn ml_key(result: &DetectionResult) -> String {
let r = &result.flow_key_raw;
format!("{}:{}-{}:{}", r.src_ip, r.src_port, r.dst_ip, r.dst_port)
}
fn rule_key(m: &RuleMatch) -> String {
format!("{}-{}", m.src, m.dst)
}

10
net-guardia/src/detection/ml/engine.rs
View File

@ -13,7 +13,7 @@ use super::inference::Inference;
use super::model_loader::MLModels;
use super::traffic_logger::TrafficLogger;
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::detection::fusion::FusionEngine;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::{EngineStats, InferenceStats};
@ -24,7 +24,7 @@ pub struct Engine {
tracker: Arc<Mutex<FlowTracker>>,
inference_pipeline: Arc<Inference>,
aggregator: Arc<Mutex<AttackAggregator>>,
ml_alert: Arc<MLAlert>,
fusion_engine: Arc<FusionEngine>,
min_packets: usize,
batch_size: usize,
inference_interval_secs: u64,
@ -37,7 +37,7 @@ impl Engine {
pub fn new(
models: Arc<MLModels>,
config: Arc<InferenceConfig>,
ml_alert: Arc<MLAlert>,
fusion_engine: Arc<FusionEngine>,
max_flows: usize,
min_packets: usize,
batch_size: usize,
@ -59,7 +59,7 @@ impl Engine {
tracker,
inference_pipeline,
aggregator,
ml_alert,
fusion_engine,
min_packets: effective_min_packets,
batch_size,
inference_interval_secs: interval_secs,
@ -202,7 +202,7 @@ impl Engine {
result.confidence,
result.ae_score,
));
self.ml_alert.broadcast_alert(result);
self.fusion_engine.record_ml(result);
}
}
}

16
net-guardia/src/detection/ml/model_loader.rs
View File

@ -1,10 +1,12 @@
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use macros::log;
use ort::session::{Session, builder::GraphOptimizationLevel};
use crate::core::infrastructure::app_config::AppConfig;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::RunnableModel;
pub struct MLModels {
@ -13,9 +15,17 @@ pub struct MLModels {
impl MLModels {
pub fn load_models(app_config: &Arc<AppConfig>) -> Result<Self, MLError> {
// Register ort-tract as the ORT execution backend before creating any sessions.
// Returns false only if already initialized, which is harmless.
ort::set_api(ort_tract::api());
#[cfg(feature = "native-ort-backend")]
{
log!(MLLog::BackendNativeOrt);
ort::init_from(PathBuf::from(env!("ONNXRUNTIME_PATH")))
.map_err(|_| MLError::InitializeFailed)?.commit();
}
#[cfg(feature = "tract-backend")]
{
log!(MLLog::BackendTract);
ort::set_api(ort_tract::api());
}
Ok(Self {
deep_autoencoder: Mutex::new(Self::load_lstm_ae(&app_config.deep_autoencoder_name)?),

1
net-guardia/src/detection/mod.rs
View File

@ -1,2 +1,3 @@
pub mod fusion;
pub mod ml;
pub mod rule;

16
net-guardia/src/model/config.rs
View File

@ -44,4 +44,20 @@ pub struct Config {
/// CPU core pinned to the ML inference spawn_blocking thread.
/// If absent, defaults to the last available core.
pub ml_cpu: Option<u32>,
/// Alert fusion mode: "or" (alert when either source fires) or "and" (require both).
/// Defaults to "or" when absent.
#[serde(default = "default_fusion_mode")]
pub fusion_mode: String,
/// Seconds within which both ML and Rule must fire to be correlated as Fusion.
/// Only used in "or" (corroboration window) and "and" modes. Defaults to 10.
#[serde(default = "default_fusion_window_secs")]
pub fusion_window_secs: u64,
}
fn default_fusion_mode() -> String {
"or".to_string()
}
fn default_fusion_window_secs() -> u64 {
10
}

6
net-guardia/src/model/log/ml.rs
View File

@ -6,6 +6,12 @@ loggable! {
#[error("Initializing Machine Learning with inference URL: {url}")]
Initializing { url: String } => tracing::Level::INFO,
#[error("ML backend: native ORT (libonnxruntime.so)")]
BackendNativeOrt => tracing::Level::INFO,
#[error("ML backend: ort-tract (pure Rust fallback)")]
BackendTract => tracing::Level::INFO,
#[error("Continuing without Machine Learning detection")]
Skiped => tracing::Level::WARN,

97
net-guardia/src/model/ml_detection.rs
View File

@ -3,6 +3,7 @@ use ort::session::Session;
use serde::{Deserialize, Serialize};
use crate::model::direction::Direction;
use crate::model::rule_detection::RuleMatch;
use crate::utils::packet_parser::{format_ipv4, format_ipv6};
pub type RunnableModel = Session;
@ -123,7 +124,24 @@ pub struct EngineStats {
}
#[derive(Debug, Clone, Serialize)]
pub struct AlertMessage {
#[serde(rename_all = "snake_case")]
pub enum AlertSource {
Ml,
Rule,
Fusion,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum AlertSeverity {
/// Single-source detection: either ML or Rule fired alone.
High,
/// Both ML and Rule agreed on the same flow within the fusion window.
Critical,
}
#[derive(Debug, Clone, Serialize)]
pub struct UnifiedAlert {
pub timestamp: u64,
pub flow_key: String,
pub src_ip: String,
@ -131,31 +149,94 @@ pub struct AlertMessage {
pub src_port: u16,
pub dst_port: u16,
pub protocol: u8,
pub source: AlertSource,
pub severity: AlertSeverity,
pub is_attack: bool,
pub attack_type: Option<String>,
pub confidence: f32,
pub ae_score: f32,
pub rule_sid: Option<u32>,
pub rule_msg: Option<String>,
}
impl AlertMessage {
pub fn from_detection_result(result: &DetectionResult) -> Self {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time is after UNIX_EPOCH")
.as_secs();
fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time is after UNIX_EPOCH")
.as_secs()
}
fn parse_ip_port(s: &str) -> (String, u16) {
match s.rfind(':') {
Some(pos) => {
let port = s[pos + 1..].parse::<u16>().unwrap_or(0);
(s[..pos].to_string(), port)
}
None => (s.to_string(), 0),
}
}
impl UnifiedAlert {
pub fn from_ml(result: &DetectionResult) -> Self {
Self {
timestamp,
timestamp: now_secs(),
flow_key: result.flow_key.clone(),
src_ip: result.flow_key_raw.src_ip.clone(),
dst_ip: result.flow_key_raw.dst_ip.clone(),
src_port: result.flow_key_raw.src_port,
dst_port: result.flow_key_raw.dst_port,
protocol: result.flow_key_raw.protocol,
source: AlertSource::Ml,
severity: AlertSeverity::High,
is_attack: result.is_attack,
attack_type: result.attack_type.clone(),
confidence: result.confidence,
ae_score: result.ae_score,
rule_sid: None,
rule_msg: None,
}
}
pub fn from_rule(m: &RuleMatch) -> Self {
let (src_ip, src_port) = parse_ip_port(&m.src);
let (dst_ip, dst_port) = parse_ip_port(&m.dst);
let flow_key = format!("{}->{}", m.src, m.dst);
Self {
timestamp: now_secs(),
flow_key,
src_ip,
dst_ip,
src_port,
dst_port,
protocol: 6,
source: AlertSource::Rule,
severity: AlertSeverity::High,
is_attack: true,
attack_type: None,
confidence: 1.0,
ae_score: 0.0,
rule_sid: Some(m.sid),
rule_msg: Some(m.msg.clone()),
}
}
pub fn from_fusion(result: &DetectionResult, m: &RuleMatch) -> Self {
Self {
timestamp: now_secs(),
flow_key: result.flow_key.clone(),
src_ip: result.flow_key_raw.src_ip.clone(),
dst_ip: result.flow_key_raw.dst_ip.clone(),
src_port: result.flow_key_raw.src_port,
dst_port: result.flow_key_raw.dst_port,
protocol: result.flow_key_raw.protocol,
source: AlertSource::Fusion,
severity: AlertSeverity::Critical,
is_attack: result.is_attack,
attack_type: result.attack_type.clone(),
confidence: result.confidence,
ae_score: result.ae_score,
rule_sid: Some(m.sid),
rule_msg: Some(m.msg.clone()),
}
}
}

14
net-guardia/src/utils/cpu_affinity.rs Normal file
View File

@ -0,0 +1,14 @@
use std::thread;
pub fn set_cpu_affinity(cpu: usize) {
#[cfg(target_os = "linux")]
unsafe {
let mut cpuset: libc::cpu_set_t = std::mem::zeroed();
libc::CPU_SET(cpu % num_cpus(), &mut cpuset);
libc::sched_setaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &cpuset);
}
}
pub fn num_cpus() -> usize {
thread::available_parallelism().map(|n| n.get()).unwrap_or(1)
}

23
net-guardia/src/web/api/detection_alert.rs Normal file
View File

@ -0,0 +1,23 @@
use actix_web::{get, web, HttpRequest, HttpResponse, Responder, Scope};
use crate::core::infrastructure::detection_alert::DetectionAlert;
use crate::web::websocket::alert_websocket;
pub fn initialize() -> Scope {
web::scope("/detection")
.service(websocket_alert)
}
#[get("/websocket/alert")]
async fn websocket_alert(
req: HttpRequest,
stream: web::Payload,
da: web::Data<DetectionAlert>,
) -> impl Responder {
match alert_websocket::websocket_alert(req, stream, da).await {
Ok(response) => response,
Err(err) => {
HttpResponse::InternalServerError().body(format!("WebSocket error: {}", err))
}
}
}

4
net-guardia/src/web/api/mod.rs
View File

@ -1,5 +1,5 @@
pub mod control;
pub mod default;
pub mod misc;
pub mod ml_alert;
pub mod detection_alert;
pub mod health;
pub mod misc;

14
net-guardia/src/web/websocket/alert_websocket.rs
View File

@ -4,8 +4,8 @@ use futures_util::StreamExt;
use macros::log;
use tokio::sync::broadcast;
use crate::core::infrastructure::ml_alert::MLAlert;
use crate::model::ml_detection::AlertMessage;
use crate::core::infrastructure::detection_alert::DetectionAlert;
use crate::model::ml_detection::UnifiedAlert;
use crate::model::error::http::HttpError;
use crate::model::error::misc::MiscError;
use crate::model::log::http::HttpLog;
@ -13,11 +13,11 @@ use crate::model::log::http::HttpLog;
pub async fn websocket_alert(
req: HttpRequest,
body: web::Payload,
ai: web::Data<MLAlert>,
da: web::Data<DetectionAlert>,
) -> Result<HttpResponse> {
let (response, session, msg_stream) = handle(&req, body)?;
let broadcast_rx = ai.subscribe_to_alerts();
let broadcast_rx = da.subscribe();
actix_web::rt::spawn(async move {
handle_alert_connection(session, msg_stream, broadcast_rx).await;
@ -29,7 +29,7 @@ pub async fn websocket_alert(
async fn handle_alert_connection(
mut session: Session,
mut msg_stream: MessageStream,
mut broadcast_rx: broadcast::Receiver<AlertMessage>,
mut broadcast_rx: broadcast::Receiver<UnifiedAlert>,
) {
loop {
tokio::select! {
@ -80,7 +80,7 @@ async fn handle_client_message(
}
}
async fn send_alert(session: &mut Session, alert: &AlertMessage) -> bool {
async fn send_alert(session: &mut Session, alert: &UnifiedAlert) -> bool {
match serde_json::to_string(alert) {
Ok(json) => session.text(json).await.is_ok(),
Err(err) => {
@ -88,4 +88,4 @@ async fn send_alert(session: &mut Session, alert: &AlertMessage) -> bool {
false
}
}
}
}

2
onnxruntime/info.txt Normal file
View File

@ -0,0 +1,2 @@
// this's onnxruntime library folder, so you can download the library from https://github.com/ParrotXray/onnxruntime-builder/releases
// you can select that matches the system architecture