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

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

Browse Source
This commit is contained in:
ParrotXray 2026-05-19 07:54:59 +00:00
parent 2d9c7f394a
commit 73ad4374bd
17 changed files with 666 additions and 436 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",

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),

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))

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()),
}
}
}

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