Revert "feat: Complete basic ml inference (#12)"

This reverts commit ae6e3aea7647800e6253565114dc2b24568bcd42.
This commit is contained in:
ParrotXray 2026-02-21 10:52:12 +00:00
parent ae6e3aea76
commit efacf42268
35 changed files with 20356 additions and 23416 deletions

3
.gitmodules vendored
View File

@ -1,3 +0,0 @@
[submodule "net-guardia-frontend"]
path = net-guardia-frontend
url = https://github.com/DaLaw2/NetGuardia-FrontEnd.git

843
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,3 @@
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use network_types::ip::IpProto;
use crate::model::ip_address::{AddrPortV4, AddrPortV6};
@ -10,134 +9,27 @@ pub enum Event {
IPv6(IPv6Event),
}
impl Event {
pub fn timestamp_us(&self) -> u64 {
match self {
Event::IPv4(e) => e.timestamp_us,
Event::IPv6(e) => e.timestamp_us,
}
}
pub fn packet_length(&self) -> u32 {
match self {
Event::IPv4(e) => e.packet_length,
Event::IPv6(e) => e.packet_length,
}
}
pub fn header_length(&self) -> u16 {
match self {
Event::IPv4(e) => e.header_length,
Event::IPv6(e) => e.header_length,
}
}
pub fn payload_length(&self) -> u32 {
match self {
Event::IPv4(e) => e.payload_length,
Event::IPv6(e) => e.payload_length,
}
}
pub fn tcp_flags(&self) -> &TcpFlags {
match self {
Event::IPv4(e) => &e.tcp_flags,
Event::IPv6(e) => &e.tcp_flags,
}
}
pub fn tcp_window_size(&self) -> u16 {
match self {
Event::IPv4(e) => e.tcp_window_size,
Event::IPv6(e) => e.tcp_window_size,
}
}
pub fn is_forward(&self) -> bool {
match self {
Event::IPv4(e) => e.is_forward,
Event::IPv6(e) => e.is_forward,
}
}
pub fn protocol(&self) -> &IpProto {
match self {
Event::IPv4(e) => &e.protocol,
Event::IPv6(e) => &e.protocol,
}
}
pub fn src_ip(&self) -> IpAddr {
match self {
Event::IPv4(e) => {
IpAddr::V4(Ipv4Addr::from(e.src_ip))
}
Event::IPv6(e) => {
IpAddr::V6(Ipv6Addr::from(e.src_ip))
}
}
}
pub fn dst_ip(&self) -> IpAddr {
match self {
Event::IPv4(e) => {
IpAddr::V4(Ipv4Addr::from(e.dst_ip))
}
Event::IPv6(e) => {
IpAddr::V6(Ipv6Addr::from(e.dst_ip))
}
}
}
pub fn src_port(&self) -> u16 {
match self {
Event::IPv4(e) => e.src_port,
Event::IPv6(e) => e.src_port,
}
}
pub fn dst_port(&self) -> u16 {
match self {
Event::IPv4(e) => e.dst_port,
Event::IPv6(e) => e.dst_port,
}
}
pub fn set_is_forward(&mut self, value: bool) {
match self {
Event::IPv4(e) => e.is_forward = value,
Event::IPv6(e) => e.is_forward = value,
}
}
}
#[repr(C, align(8))]
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct IPv4Event {
pub protocol: IpProto,
pub src_ip: u32,
pub dst_ip: u32,
pub src_port: u16,
pub dst_port: u16,
pub packet_length: u32,
pub payload_length: u32,
pub header_length: u16,
pub timestamp_us: u64,
pub tcp_flags: TcpFlags,
pub tcp_window_size: u16,
pub is_forward: bool,
pub source_ip: u32,
pub destination_ip: u32,
pub source_port: u16,
pub destination_port: u16,
pub len: u32,
pub timestamp: u64,
}
impl IPv4Event {
#[inline(always)]
pub fn source_addr(&self) -> AddrPortV4 {
AddrPortV4::new(self.src_ip, self.src_port)
AddrPortV4::new(self.source_ip, self.source_port)
}
#[inline(always)]
pub fn destination_addr(&self) -> AddrPortV4 {
AddrPortV4::new(self.dst_ip, self.dst_port)
AddrPortV4::new(self.destination_ip, self.destination_port)
}
}
@ -145,54 +37,22 @@ impl IPv4Event {
#[derive(Clone)]
pub struct IPv6Event {
pub protocol: IpProto,
pub src_ip: u128,
pub dst_ip: u128,
pub src_port: u16,
pub dst_port: u16,
pub packet_length: u32,
pub payload_length: u32,
pub header_length: u16,
pub timestamp_us: u64,
pub tcp_flags: TcpFlags,
pub tcp_window_size: u16,
pub is_forward: bool,
pub source_ip: u128,
pub destination_ip: u128,
pub source_port: u16,
pub destination_port: u16,
pub len: u32,
pub timestamp: u64,
}
impl IPv6Event {
#[inline(always)]
pub fn source_addr(&self) -> AddrPortV6 {
AddrPortV6::new(self.src_ip, self.src_port)
AddrPortV6::new(self.source_ip, self.source_port)
}
#[inline(always)]
pub fn destination_addr(&self) -> AddrPortV6 {
AddrPortV6::new(self.dst_ip, self.dst_port)
}
}
#[derive(Debug, Clone, Default)]
pub struct TcpFlags {
pub fin: bool,
pub syn: bool,
pub rst: bool,
pub psh: bool,
pub ack: bool,
pub urg: bool,
pub ece: bool,
pub cwr: bool,
}
impl TcpFlags {
pub fn from_byte(flags: u8) -> Self {
Self {
fin: (flags & 0x01) != 0,
syn: (flags & 0x02) != 0,
rst: (flags & 0x04) != 0,
psh: (flags & 0x08) != 0,
ack: (flags & 0x10) != 0,
urg: (flags & 0x20) != 0,
ece: (flags & 0x40) != 0,
cwr: (flags & 0x80) != 0,
}
AddrPortV6::new(self.destination_ip, self.destination_port)
}
}

View File

@ -1,13 +1,9 @@
[Config]
ingress_ifname = "enp4s0f1" # Ingress NIC Name
egress_ifname = "enp4s0f0" # Egress NIC Name
geoip_db_name = "GeoLite2-City.mmdb"
deep_autoencoder_name = "deep_autoencoder.onnx"
random_forest_name = "random_forest.onnx"
mlp_name = "mlp.onnx"
models_config_name = "inference_config.json"
geoip_db_path = "net-guardia/static/geo/GeoLite2-City.mmdb"
combined_queue_count = 8 # NIC Combined Queue Count (ethtool -l <NIC>)
channel_size = 4096
xsk_channel_size = 4096
fill_queue_size = 4096 # Umem Used (Should not modify)
comp_queue_size = 4096 # Umem Used (Should not modify)
tx_queue_size = 4096 # Umem Used (Should not modify)
@ -15,8 +11,4 @@ rx_queue_size = 4096 # Umem Used (Should not modify)
frame_size = 4096 # Umem Used (Should not modify)
frame_count = 4096 # Umem Used (Should not modify)
http_server_bind_port = 8080 # Http Server Listen Port
refresh_interval = 5 # Statistics Refresh Time
max_concurrent_flows = 10000 # max_flows: track up to 10000 concurrent flows
min_packets_for_inference = 5 # min_packets: minimum 10 packets per flow for inference
inference_interval_secs = 5 # interval_secs: run inference every 5 seconds
refresh_interval = 5 # Statistics Refresh Time

File diff suppressed because it is too large Load Diff

View File

@ -1,692 +0,0 @@
{
"created_at": "2026-01-27T13:05:39.258903",
"model": {
"deep_autoencoder": {
"file": "deep_autoencoder.onnx",
"input_dim": 78,
"encoding_dim": 16
},
"random_forest": {
"file": "random_forest.onnx",
"n_estimators": 100,
"n_features": 78
},
"mlp_classifier": {
"file": "mlp.onnx",
"input_dim": 78,
"n_classes": 10
}
},
"preprocessing": {
"clip_params": {
"Destination Port": {
"lower": 22.0,
"upper": 63734.0
},
"Flow Duration": {
"lower": 1.0,
"upper": 118756600.72000001
},
"Total Fwd Packets": {
"lower": 1.0,
"upper": 98.0
},
"Total Backward Packets": {
"lower": 0.0,
"upper": 126.0
},
"Total Length of Fwd Packets": {
"lower": 0.0,
"upper": 13119.520000000019
},
"Total Length of Bwd Packets": {
"lower": 0.0,
"upper": 186394.64000000013
},
"Fwd Packet Length Max": {
"lower": 0.0,
"upper": 5840.0
},
"Fwd Packet Length Min": {
"lower": 0.0,
"upper": 98.0
},
"Fwd Packet Length Mean": {
"lower": 0.0,
"upper": 1932.5
},
"Fwd Packet Length Std": {
"lower": 0.0,
"upper": 2376.49858
},
"Bwd Packet Length Max": {
"lower": 0.0,
"upper": 4380.0
},
"Bwd Packet Length Min": {
"lower": 0.0,
"upper": 308.0
},
"Bwd Packet Length Mean": {
"lower": 0.0,
"upper": 1715.628244226656
},
"Bwd Packet Length Std": {
"lower": 0.0,
"upper": 1167.9936145600002
},
"Flow Bytes\/s": {
"lower": 0.0,
"upper": 37000000.0
},
"Flow Packets\/s": {
"lower": 0.041637660508,
"upper": 2000000.0
},
"Flow IAT Mean": {
"lower": 1.0,
"upper": 30800000.0
},
"Flow IAT Std": {
"lower": 0.0,
"upper": 44253994.319940574
},
"Flow IAT Max": {
"lower": 1.0,
"upper": 96199268.00000003
},
"Flow IAT Min": {
"lower": 0.0,
"upper": 1999969.04
},
"Fwd IAT Total": {
"lower": 0.0,
"upper": 119000000.0
},
"Fwd IAT Mean": {
"lower": 0.0,
"upper": 74300000.0
},
"Fwd IAT Std": {
"lower": 0.0,
"upper": 35894248.669783354
},
"Fwd IAT Max": {
"lower": 0.0,
"upper": 97800000.0
},
"Fwd IAT Min": {
"lower": 0.0,
"upper": 74300000.0
},
"Bwd IAT Total": {
"lower": 0.0,
"upper": 118000000.0
},
"Bwd IAT Mean": {
"lower": 0.0,
"upper": 73734410.76000024
},
"Bwd IAT Std": {
"lower": 0.0,
"upper": 30100000.0
},
"Bwd IAT Max": {
"lower": 0.0,
"upper": 93800000.0
},
"Bwd IAT Min": {
"lower": 0.0,
"upper": 73734410.76000024
},
"Fwd PSH Flags": {
"lower": 0.0,
"upper": 1.0
},
"Bwd PSH Flags": {
"lower": 0.0,
"upper": 0.0
},
"Fwd URG Flags": {
"lower": 0.0,
"upper": 0.0
},
"Bwd URG Flags": {
"lower": 0.0,
"upper": 0.0
},
"Fwd Header Length": {
"lower": 20.0,
"upper": 2376.0
},
"Bwd Header Length": {
"lower": 0.0,
"upper": 3312.0
},
"Fwd Packets\/s": {
"lower": 0.02417299188,
"upper": 2000000.0
},
"Bwd Packets\/s": {
"lower": 0.0,
"upper": 142857.1429
},
"Min Packet Length": {
"lower": 0.0,
"upper": 89.0
},
"Max Packet Length": {
"lower": 0.0,
"upper": 5840.0
},
"Packet Length Mean": {
"lower": 0.0,
"upper": 1153.4362532
},
"Packet Length Std": {
"lower": 0.0,
"upper": 1845.154148000012
},
"Packet Length Variance": {
"lower": 0.0,
"upper": 3404593.931120044
},
"FIN Flag Count": {
"lower": 0.0,
"upper": 1.0
},
"SYN Flag Count": {
"lower": 0.0,
"upper": 1.0
},
"RST Flag Count": {
"lower": 0.0,
"upper": 0.0
},
"PSH Flag Count": {
"lower": 0.0,
"upper": 1.0
},
"ACK Flag Count": {
"lower": 0.0,
"upper": 1.0
},
"URG Flag Count": {
"lower": 0.0,
"upper": 1.0
},
"CWE Flag Count": {
"lower": 0.0,
"upper": 0.0
},
"ECE Flag Count": {
"lower": 0.0,
"upper": 0.0
},
"Down\/Up Ratio": {
"lower": 0.0,
"upper": 5.0
},
"Average Packet Size": {
"lower": 0.0,
"upper": 1182.362785582229
},
"Avg Fwd Segment Size": {
"lower": 0.0,
"upper": 1932.5
},
"Avg Bwd Segment Size": {
"lower": 0.0,
"upper": 1715.628244226656
},
"Fwd Header Length.1": {
"lower": 20.0,
"upper": 2376.0
},
"Fwd Avg Bytes\/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Fwd Avg Packets\/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Fwd Avg Bulk Rate": {
"lower": 0.0,
"upper": 0.0
},
"Bwd Avg Bytes\/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Bwd Avg Packets\/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Bwd Avg Bulk Rate": {
"lower": 0.0,
"upper": 0.0
},
"Subflow Fwd Packets": {
"lower": 1.0,
"upper": 98.0
},
"Subflow Fwd Bytes": {
"lower": 0.0,
"upper": 13119.520000000019
},
"Subflow Bwd Packets": {
"lower": 0.0,
"upper": 126.0
},
"Subflow Bwd Bytes": {
"lower": 0.0,
"upper": 186394.64000000013
},
"Init_Win_bytes_forward": {
"lower": -1.0,
"upper": 65535.0
},
"Init_Win_bytes_backward": {
"lower": -1.0,
"upper": 65535.0
},
"act_data_pkt_fwd": {
"lower": 0.0,
"upper": 49.0
},
"min_seg_size_forward": {
"lower": 20.0,
"upper": 40.0
},
"Active Mean": {
"lower": 0.0,
"upper": 3410152.040000004
},
"Active Std": {
"lower": 0.0,
"upper": 1944723.8431600018
},
"Active Max": {
"lower": 0.0,
"upper": 5717825.840000004
},
"Active Min": {
"lower": 0.0,
"upper": 2918887.7600000002
},
"Idle Mean": {
"lower": 0.0,
"upper": 94100000.0
},
"Idle Std": {
"lower": 0.0,
"upper": 20712312.015200634
},
"Idle Max": {
"lower": 0.0,
"upper": 94800000.0
},
"Idle Min": {
"lower": 0.0,
"upper": 94100000.0
}
},
"scaler": {
"mean": [
9415.502043247605,
11217541.457604988,
5.41509711200182,
5.010050165039152,
508.74335117243135,
3302.166099044608,
217.24886707430434,
20.105687086824716,
65.40745127720037,
68.76070937185386,
394.97690859650953,
49.88046220640826,
159.85952962743096,
121.94653847258593,
872411.6556228747,
64701.11201063558,
840255.3525863544,
1516373.8724808302,
4422301.697900706,
24666.656195771673,
10897580.686685171,
1785372.7881935516,
1206278.3042178946,
4276119.216091527,
1025069.5899462275,
10158901.07420493,
1666490.9735317046,
939816.6614413982,
3689910.1546691586,
1003727.4123986625,
0.054886351088404936,
0.0,
0.0,
0.0,
140.54050707030981,
133.63606216540694,
58358.168874140014,
4983.791329731815,
19.812641079549177,
480.4855221752525,
108.29833052637488,
144.4437888623996,
99406.03379366906,
0.01811977227544623,
0.054886351088404936,
0.0,
0.25800834720207716,
0.28694244020382764,
0.11596161536441252,
0.0,
0.0,
0.6969992921551522,
122.46836156918589,
65.40745127720037,
159.85952962743843,
140.54050707030981,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
5.41509711200182,
508.74335117243135,
5.010050165039152,
3302.166099044608,
6891.551141020379,
2404.6474695096604,
2.6520584031389776,
25.79473555242033,
61328.50855670946,
36311.697575423306,
137445.81057140988,
37549.65545208145,
3689065.3106244486,
167385.3556173429,
3883594.4613142335,
3458258.8992590285
],
"std": [
19743.81786383566,
30111390.85802201,
10.612016675700064,
12.872110268702288,
1653.7982707262067,
17387.56157755931,
645.437387791972,
22.356844266679282,
193.61621129020307,
244.137250250046,
801.628533600607,
65.65797274491328,
277.7559653593753,
269.57514761088515,
3622907.235127058,
237431.11489053545,
3293477.440659972,
5309087.889319858,
14143198.815741453,
160458.44923675407,
29985211.254190512,
8130184.699376457,
4247601.6170286415,
14271787.10745024,
7769581.505106356,
29154202.839787327,
8031691.1024753135,
3622760.030182759,
13368925.151392205,
7684129.1959599955,
0.22775829195136954,
1.0,
1.0,
1.0,
271.11142624961946,
342.17520432711024,
231133.09437362824,
14026.707464402314,
22.005832334437958,
970.4322795025132,
179.9690556845557,
280.2535015764328,
343919.1355100174,
0.13338457979891152,
0.22775829195136954,
1.0,
0.4375386154114052,
0.4523344738284939,
0.32017888613474904,
1.0,
1.0,
0.6456077980707177,
185.26118186235436,
193.61621129020307,
277.7559653594075,
271.11142624961946,
1.0,
1.0,
1.0,
1.0,
1.0,
1.0,
10.612016675700064,
1653.7982707262067,
12.872110268702288,
17387.56157755931,
14922.88749872485,
9322.158869764009,
5.784504103070997,
6.252894021045679,
322010.01591284445,
201975.98444339563,
641116.7578408231,
245957.32066446543,
13068297.83388166,
1617503.1921646637,
13724469.425904194,
12765671.772576509
],
"feature_names": [
"Destination Port",
"Flow Duration",
"Total Fwd Packets",
"Total Backward Packets",
"Total Length of Fwd Packets",
"Total Length of Bwd Packets",
"Fwd Packet Length Max",
"Fwd Packet Length Min",
"Fwd Packet Length Mean",
"Fwd Packet Length Std",
"Bwd Packet Length Max",
"Bwd Packet Length Min",
"Bwd Packet Length Mean",
"Bwd Packet Length Std",
"Flow Bytes\/s",
"Flow Packets\/s",
"Flow IAT Mean",
"Flow IAT Std",
"Flow IAT Max",
"Flow IAT Min",
"Fwd IAT Total",
"Fwd IAT Mean",
"Fwd IAT Std",
"Fwd IAT Max",
"Fwd IAT Min",
"Bwd IAT Total",
"Bwd IAT Mean",
"Bwd IAT Std",
"Bwd IAT Max",
"Bwd IAT Min",
"Fwd PSH Flags",
"Bwd PSH Flags",
"Fwd URG Flags",
"Bwd URG Flags",
"Fwd Header Length",
"Bwd Header Length",
"Fwd Packets\/s",
"Bwd Packets\/s",
"Min Packet Length",
"Max Packet Length",
"Packet Length Mean",
"Packet Length Std",
"Packet Length Variance",
"FIN Flag Count",
"SYN Flag Count",
"RST Flag Count",
"PSH Flag Count",
"ACK Flag Count",
"URG Flag Count",
"CWE Flag Count",
"ECE Flag Count",
"Down\/Up Ratio",
"Average Packet Size",
"Avg Fwd Segment Size",
"Avg Bwd Segment Size",
"Fwd Header Length.1",
"Fwd Avg Bytes\/Bulk",
"Fwd Avg Packets\/Bulk",
"Fwd Avg Bulk Rate",
"Bwd Avg Bytes\/Bulk",
"Bwd Avg Packets\/Bulk",
"Bwd Avg Bulk Rate",
"Subflow Fwd Packets",
"Subflow Fwd Bytes",
"Subflow Bwd Packets",
"Subflow Bwd Bytes",
"Init_Win_bytes_forward",
"Init_Win_bytes_backward",
"act_data_pkt_fwd",
"min_seg_size_forward",
"Active Mean",
"Active Std",
"Active Max",
"Active Min",
"Idle Mean",
"Idle Std",
"Idle Max",
"Idle Min"
]
},
"post_scaling_clip": {
"min": -5.0,
"max": 5.0
}
},
"ensemble": {
"strategy_name": "Max",
"threshold": 0.3603835832054485,
"tpr": 0.9985636048676042,
"fpr": 0.0029998719808261593,
"precision": 0.9879023673677936,
"f1": 0.9932043770233031
},
"ae_normalization": {
"min": 1.480184057280627e-5,
"max": 1.1978646574008398,
"mean": 0.0023598657051511107,
"std": 0.008826375460685048,
"median": 0.00031163828850514576,
"p90": 0.004909278753678756,
"p95": 0.010077479060604477,
"p99": 0.03354822433745485
},
"attack_labels": {
"0": "Bot",
"1": "DDoS",
"2": "DoS GoldenEye",
"3": "DoS Hulk",
"4": "DoS Slowhttptest",
"5": "DoS slowloris",
"6": "FTP-Patator",
"7": "PortScan",
"8": "SSH-Patator",
"9": "Web Attack"
},
"feature_order": [
"Destination Port",
"Flow Duration",
"Total Fwd Packets",
"Total Backward Packets",
"Total Length of Fwd Packets",
"Total Length of Bwd Packets",
"Fwd Packet Length Max",
"Fwd Packet Length Min",
"Fwd Packet Length Mean",
"Fwd Packet Length Std",
"Bwd Packet Length Max",
"Bwd Packet Length Min",
"Bwd Packet Length Mean",
"Bwd Packet Length Std",
"Flow Bytes\/s",
"Flow Packets\/s",
"Flow IAT Mean",
"Flow IAT Std",
"Flow IAT Max",
"Flow IAT Min",
"Fwd IAT Total",
"Fwd IAT Mean",
"Fwd IAT Std",
"Fwd IAT Max",
"Fwd IAT Min",
"Bwd IAT Total",
"Bwd IAT Mean",
"Bwd IAT Std",
"Bwd IAT Max",
"Bwd IAT Min",
"Fwd PSH Flags",
"Bwd PSH Flags",
"Fwd URG Flags",
"Bwd URG Flags",
"Fwd Header Length",
"Bwd Header Length",
"Fwd Packets\/s",
"Bwd Packets\/s",
"Min Packet Length",
"Max Packet Length",
"Packet Length Mean",
"Packet Length Std",
"Packet Length Variance",
"FIN Flag Count",
"SYN Flag Count",
"RST Flag Count",
"PSH Flag Count",
"ACK Flag Count",
"URG Flag Count",
"CWE Flag Count",
"ECE Flag Count",
"Down\/Up Ratio",
"Average Packet Size",
"Avg Fwd Segment Size",
"Avg Bwd Segment Size",
"Fwd Header Length.1",
"Fwd Avg Bytes\/Bulk",
"Fwd Avg Packets\/Bulk",
"Fwd Avg Bulk Rate",
"Bwd Avg Bytes\/Bulk",
"Bwd Avg Packets\/Bulk",
"Bwd Avg Bulk Rate",
"Subflow Fwd Packets",
"Subflow Fwd Bytes",
"Subflow Bwd Packets",
"Subflow Bwd Bytes",
"Init_Win_bytes_forward",
"Init_Win_bytes_backward",
"act_data_pkt_fwd",
"min_seg_size_forward",
"Active Mean",
"Active Std",
"Active Max",
"Active Min",
"Idle Mean",
"Idle Std",
"Idle Max",
"Idle Min"
]
}

Binary file not shown.

View File

@ -1,5 +1,5 @@
{
"threshold": 0.3603835832054485,
"threshold": 0.15133114984430773,
"strategy_name": "Max",
"clip_params": {
"Destination Port": {
@ -58,11 +58,11 @@
"lower": 0.0,
"upper": 1167.9936145600002
},
"Flow Bytes\/s": {
"Flow Bytes/s": {
"lower": 0.0,
"upper": 37000000.0
},
"Flow Packets\/s": {
"Flow Packets/s": {
"lower": 0.041637660508,
"upper": 2000000.0
},
@ -146,11 +146,11 @@
"lower": 0.0,
"upper": 3312.0
},
"Fwd Packets\/s": {
"Fwd Packets/s": {
"lower": 0.02417299188,
"upper": 2000000.0
},
"Bwd Packets\/s": {
"Bwd Packets/s": {
"lower": 0.0,
"upper": 142857.1429
},
@ -206,7 +206,7 @@
"lower": 0.0,
"upper": 0.0
},
"Down\/Up Ratio": {
"Down/Up Ratio": {
"lower": 0.0,
"upper": 5.0
},
@ -226,11 +226,11 @@
"lower": 20.0,
"upper": 2376.0
},
"Fwd Avg Bytes\/Bulk": {
"Fwd Avg Bytes/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Fwd Avg Packets\/Bulk": {
"Fwd Avg Packets/Bulk": {
"lower": 0.0,
"upper": 0.0
},
@ -238,11 +238,11 @@
"lower": 0.0,
"upper": 0.0
},
"Bwd Avg Bytes\/Bulk": {
"Bwd Avg Bytes/Bulk": {
"lower": 0.0,
"upper": 0.0
},
"Bwd Avg Packets\/Bulk": {
"Bwd Avg Packets/Bulk": {
"lower": 0.0,
"upper": 0.0
},
@ -478,26 +478,31 @@
"post_clip_min": -5.0,
"post_clip_max": 5.0,
"ae_normalization": {
"min": 1.480184057280627e-5,
"max": 1.1978646574008398,
"mean": 0.0023598657051511107,
"std": 0.008826375460685048,
"median": 0.00031163828850514576,
"p90": 0.004909278753678756,
"p95": 0.010077479060604477,
"p99": 0.03354822433745485
"min": 4.147972858692375e-06,
"max": 1.292201307125899,
"mean": 0.0026365920400427795,
"std": 0.011092374621751085,
"median": 0.00023979156462091921,
"p90": 0.005233132974117131,
"p95": 0.01092618446409825,
"p99": 0.04010412447293844
},
"attack_labels": {
"0": "Bot",
"1": "DDoS",
"2": "DoS GoldenEye",
"3": "DoS Hulk",
"4": "DoS Slowhttptest",
"5": "DoS slowloris",
"6": "FTP-Patator",
"7": "PortScan",
"8": "SSH-Patator",
"9": "Web Attack"
"0": "BENIGN",
"1": "Bot",
"2": "DDoS",
"3": "DoS GoldenEye",
"4": "DoS Hulk",
"5": "DoS Slowhttptest",
"6": "DoS slowloris",
"7": "FTP-Patator",
"8": "Heartbleed",
"9": "Infiltration",
"10": "PortScan",
"11": "SSH-Patator",
"12": "Web Attack <20> Brute Force",
"13": "Web Attack <20> Sql Injection",
"14": "Web Attack <20> XSS"
},
"feature_names": [
"Destination Port",
@ -514,8 +519,8 @@
"Bwd Packet Length Min",
"Bwd Packet Length Mean",
"Bwd Packet Length Std",
"Flow Bytes\/s",
"Flow Packets\/s",
"Flow Bytes/s",
"Flow Packets/s",
"Flow IAT Mean",
"Flow IAT Std",
"Flow IAT Max",
@ -536,8 +541,8 @@
"Bwd URG Flags",
"Fwd Header Length",
"Bwd Header Length",
"Fwd Packets\/s",
"Bwd Packets\/s",
"Fwd Packets/s",
"Bwd Packets/s",
"Min Packet Length",
"Max Packet Length",
"Packet Length Mean",
@ -551,16 +556,16 @@
"URG Flag Count",
"CWE Flag Count",
"ECE Flag Count",
"Down\/Up Ratio",
"Down/Up Ratio",
"Average Packet Size",
"Avg Fwd Segment Size",
"Avg Bwd Segment Size",
"Fwd Header Length.1",
"Fwd Avg Bytes\/Bulk",
"Fwd Avg Packets\/Bulk",
"Fwd Avg Bytes/Bulk",
"Fwd Avg Packets/Bulk",
"Fwd Avg Bulk Rate",
"Bwd Avg Bytes\/Bulk",
"Bwd Avg Packets\/Bulk",
"Bwd Avg Bytes/Bulk",
"Bwd Avg Packets/Bulk",
"Bwd Avg Bulk Rate",
"Subflow Fwd Packets",
"Subflow Fwd Bytes",

Binary file not shown.

@ -1 +0,0 @@
Subproject commit bf27f315f2536025fa9068b9027a151a160bb70d

View File

@ -32,10 +32,9 @@ tracing-appender = "0.2.3"
tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
url = "2.5.7"
xsk-rs = { workspace = true }
maxminddb = "0.27.1"
maxminddb = "0.26.0"
lru = "0.16.2"
futures = "0.3.31"
tract-onnx = "0.22.0"
[build-dependencies]
cargo_metadata = { workspace = true }

View File

@ -19,7 +19,6 @@ use crate::core::infrastructure::health::SystemHealth;
use crate::core::infrastructure::app_config::AppConfig;
use crate::model::error::system::SystemError;
use crate::model::error::Error;
use crate::ml::engine::Engine;
pub struct EbpfServices {
pub xsk_manager: Arc<XskManager>,
@ -38,9 +37,9 @@ impl EbpfServices {
) -> Result<Self, Error> {
let xsk_manager = XskManager::new(app_config.clone(), ingress_ebpf)?;
let access_control = AccessControl::new(ingress_ebpf)?;
let health = SystemHealth::new(app_config.clone())?;
let health = SystemHealth::new(&app_config)?;
let service = Service::new(ingress_ebpf)?;
let statistics = Statistics::new(app_config.clone(), ingress_ebpf, egress_ebpf)?;
let statistics = Statistics::new(app_config, ingress_ebpf, egress_ebpf)?;
let ebpf_services = Self {
xsk_manager: Arc::new(xsk_manager),
access_control: Arc::new(access_control),
@ -52,12 +51,12 @@ 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>) -> Result<(), Error> {
let xsk_manager = self.xsk_manager.clone();
let statistics = self.statistics.clone();
let health = self.health.clone();
xsk_manager.run(Some(ml_engine))?;
xsk_manager.run()?;
let statistics_shutdown = statistics.run().await;
let health_shutdown = health.run(Duration::from_secs(3)).await;

View File

@ -94,7 +94,7 @@ impl Statistics {
let boot_time = boot_time();
let mut ipv4_maps = HashMap::new();
let mut ipv6_maps = HashMap::new();
let geo_ip = match GeoIpService::new(&app_config.geoip_db_name) {
let geo_ip = match GeoIpService::new(&app_config.geoip_db_path) {
Ok(service) => Some(Arc::new(service)),
Err(err) => {
log!(MiscError::InvalidGeoIPConfiguration(err));

View File

@ -23,7 +23,6 @@ use crate::model::error::Error;
use crate::model::error::ebpf::EbpfError;
use crate::model::error::system::SystemError;
use crate::model::log::ebpf::EbpfLog;
use crate::ml::engine::{Engine, PacketProcessor};
pub struct XskManager {
app_config: Arc<AppConfig>,
@ -46,23 +45,20 @@ impl XskManager {
})
}
pub fn run(&self, ml_engine: Option<Arc<Engine>>) -> Result<(), Error> {
pub fn run(&self) -> Result<(), Error> {
let config = self.app_config.config.clone();
let combined_queue_count = config.combined_queue_count;
let packet_processor = ml_engine.map(|engine| Arc::new(PacketProcessor::new(engine)));
for queue_id in 0..combined_queue_count {
let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(config.channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(config.channel_size);
let (ingress_to_egress_tx, ingress_to_egress_rx) = bounded(config.xsk_channel_size);
let (egress_to_ingress_tx, egress_to_ingress_rx) = bounded(config.xsk_channel_size);
let ingress_xsk = XskPair::new(
config.clone(),
queue_id,
&config.ingress_ifname,
&config.egress_ifname,
Direction::Ingress,
packet_processor.clone(),
Direction::Ingress
)?;
let egress_xsk = XskPair::new(
@ -71,7 +67,6 @@ impl XskManager {
&config.egress_ifname,
&config.ingress_ifname,
Direction::Egress,
None,
)?;
let mut xsk_map = self.xsk_map.lock();
@ -110,7 +105,6 @@ pub struct XskPair {
tx: TxQueue,
rx: RxQueue,
frame_pool: Arc<Mutex<Vec<FrameDesc>>>,
packet_processor: Option<Arc<PacketProcessor>>,
}
impl XskPair {
@ -120,7 +114,6 @@ impl XskPair {
rx_ifname: &str,
tx_ifname: &str,
direction: Direction,
packet_processor: Option<Arc<PacketProcessor>>,
) -> Result<Self, Error> {
let rx_ifname_c = CString::new(rx_ifname).map_err(|_| SystemError::UnknownError)?;
@ -177,7 +170,6 @@ impl XskPair {
tx,
rx,
frame_pool: Arc::new(Mutex::new(pool_frames)),
packet_processor,
};
Ok(xsk_pair)
@ -196,7 +188,6 @@ impl XskPair {
.name(thread_name.clone())
.spawn(move || {
let mut shutdown_rx = Some(shutdown_rx);
let mut idle_count: u32 = 0;
loop {
if let Some(ref mut rx) = shutdown_rx {
@ -208,48 +199,40 @@ impl XskPair {
}
}
let mut total_activity = 0;
match self.process_comp_queue() {
Ok(count) => total_activity += count,
Err(e) => log!(EbpfLog::CompQueueError(format!("{:?}", e))),
if let Err(e) = self.process_comp_queue() {
log!(EbpfLog::CompQueueError {
error: format!("{:?}", e)
});
}
match self.process_rx_queue(&forward_tx) {
Ok(count) => total_activity += count,
Err(e) => log!(EbpfLog::RXQueueError(format!("{:?}", e))),
if let Err(e) = self.process_rx_queue(&forward_tx) {
log!(EbpfLog::RXQueueError {
error: format!("{:?}", e)
});
}
match self.process_tx_queue(&forward_rx) {
Ok(count) => total_activity += count,
Err(e) => log!(EbpfLog::TXQueueError(format!("{:?}", e))),
if let Err(e) = self.process_tx_queue(&forward_rx) {
log!(EbpfLog::TXQueueError {
error: format!("{:?}", e)
});
}
if total_activity == 0 {
idle_count = idle_count.saturating_add(1);
} else {
idle_count = 0;
}
let sleep_us = match idle_count {
0..=10 => 1,
11..=100 => 10,
_ => 100,
};
thread::sleep(Duration::from_micros(sleep_us));
thread::sleep(Duration::from_micros(1));
}
log!(EbpfLog::XSKShutdown);
})
.map(|_| shutdown_tx)
.map_err(|e| {
log!(EbpfLog::ThreadSpawnFailed(thread_name.clone(), e.to_string()));
log!(EbpfLog::ThreadSpawnFailed {
thread_name: thread_name.clone(),
error: e.to_string()
});
EbpfError::ThreadSpawnFailed(e)
})
}
fn process_comp_queue(&mut self) -> Result<usize, EbpfError> {
fn process_comp_queue(&mut self) -> Result<(), EbpfError> {
let mut comp_descs = vec![FrameDesc::default(); 256];
let nb_completed = unsafe { self.comp_queue.consume(&mut comp_descs) };
@ -262,10 +245,10 @@ impl XskPair {
}
}
Ok(nb_completed)
Ok(())
}
fn process_rx_queue(&mut self, forward_tx: &Sender<Vec<u8>>) -> Result<usize, EbpfError> {
fn process_rx_queue(&mut self, forward_tx: &Sender<Vec<u8>>) -> Result<(), EbpfError> {
let mut rx_descs = vec![FrameDesc::default(); 64];
let rx_count = unsafe { self.rx.consume(&mut rx_descs) };
@ -277,10 +260,6 @@ impl XskPair {
let data = unsafe { self.umem.data(rx_desc) };
let packet_data = data.contents()[..packet_len].to_vec();
if let Some(ref processor) = self.packet_processor {
processor.process(&packet_data);
}
if let Err(e) = forward_tx.try_send(packet_data) {
match e {
crossbeam::channel::TrySendError::Full(_) => {
@ -296,15 +275,18 @@ impl XskPair {
unsafe {
let produced = self.fill_queue.produce(&rx_descs[..rx_count]);
if produced != rx_count {
log!(EbpfLog::FillQueueIncomplete(produced, rx_count));
log!(EbpfLog::FillQueueIncomplete {
produced,
expected: rx_count
});
}
}
}
Ok(rx_count)
Ok(())
}
fn process_tx_queue(&mut self, forward_rx: &Receiver<Vec<u8>>) -> Result<usize, EbpfError> {
fn process_tx_queue(&mut self, forward_rx: &Receiver<Vec<u8>>) -> Result<(), EbpfError> {
let mut packets_to_send = Vec::with_capacity(64);
while let Ok(packet) = forward_rx.try_recv() {
packets_to_send.push(packet);
@ -314,7 +296,7 @@ impl XskPair {
}
if packets_to_send.is_empty() {
return Ok(0);
return Ok(());
}
let _ = self.process_comp_queue();
@ -325,8 +307,10 @@ impl XskPair {
};
if pool_size == 0 {
log!(EbpfLog::FramePoolExhausted(packets_to_send.len()));
return Ok(0);
log!(EbpfLog::FramePoolExhausted {
send_len: packets_to_send.len()
});
return Ok(());
}
let mut frames = Vec::with_capacity(packets_to_send.len());
@ -343,7 +327,7 @@ impl XskPair {
if frames.is_empty() {
log!(EbpfLog::NoFramesAvailable);
return Ok(0);
return Ok(());
}
for (frame, packet) in frames.iter_mut().zip(packets_to_send.iter()) {
@ -356,14 +340,16 @@ impl XskPair {
}
}
let nb_submitted = unsafe { self.tx.produce(&frames) };
let _nb_submitted = unsafe { self.tx.produce(&frames) };
if let Err(e) = self.tx.wakeup() {
if e.kind() != std::io::ErrorKind::WouldBlock {
log!(EbpfLog::TXWakeupFailed(e.to_string()));
log!(EbpfLog::TXWakeupFailed {
error: e.to_string()
});
}
}
Ok(nb_submitted)
Ok(())
}
}

View File

@ -1,8 +1,9 @@
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::path::Path;
use std::sync::Arc;
use maxminddb::{geoip2, MaxMindDbError, Reader};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use lru::LruCache;
use std::num::NonZeroUsize;
@ -16,8 +17,7 @@ pub struct GeoIpService {
}
impl GeoIpService {
pub fn new(db_name: &str) -> Result<Self, MaxMindDbError> {
let db_path = PathBuf::from("net-guardia/static/geo").join(db_name);
pub fn new<P: AsRef<Path>>(db_path: P) -> Result<Self, MaxMindDbError> {
Self::with_cache_size(db_path, 10000)
}
@ -48,10 +48,7 @@ impl GeoIpService {
Self::lookup_from_db_blocking(&reader, ip)
})
.await
.map_err(|e| MaxMindDbError::InvalidDatabase {
message: format!("Task join error: {}", e),
offset: None,
})??;
.map_err(|e| MaxMindDbError::InvalidDatabase(format!("Task join error: {}", e)))??;
{
let mut cache = self.cache.write().await;
@ -65,22 +62,38 @@ impl GeoIpService {
reader: &Reader<Vec<u8>>,
ip: IpAddr,
) -> Result<Option<GeoLocation>, MaxMindDbError> {
let lookup_result = reader.lookup(ip)?;
let city_option: Option<geoip2::City> = lookup_result.decode()?;
let city_option: Option<geoip2::City> = reader.lookup(ip)?;
Ok(city_option.map(|city| {
let country_name = city.country.names.english
let country_name = city
.country
.as_ref()
.and_then(|c| c.names.as_ref())
.and_then(|n| n.get("en"))
.map(|s| s.to_string());
let country_code = city.country.iso_code
let country_code = city
.country
.as_ref()
.and_then(|c| c.iso_code)
.map(|s| s.to_string());
let city_name = city.city.names.english
let city_name = city
.city
.as_ref()
.and_then(|c| c.names.as_ref())
.and_then(|n| n.get("en"))
.map(|s| s.to_string());
let latitude = city.location.latitude.or(Some(0.0));
let longitude = city.location.longitude.or(Some(0.0));
let timezone = city.location.time_zone.map(|s| s.to_string());
let latitude = city.location.as_ref().and_then(|l| l.latitude);
let longitude = city.location.as_ref().and_then(|l| l.longitude);
let timezone = city
.location
.as_ref()
.and_then(|l| l.time_zone)
.map(|s| s.to_string());
GeoLocation {
country: country_name,

View File

@ -102,7 +102,7 @@ pub struct SystemHealthStatus {
}
impl SystemHealth {
pub fn new(config: Arc<AppConfig>) -> Result<Self, Error> {
pub fn new(config: &Arc<AppConfig>) -> Result<Self, Error> {
let (broadcast_tx, _) = broadcast::channel(100);
let health = SystemHealth {

View File

@ -15,22 +15,17 @@ use crate::model::error::ebpf::EbpfError;
use crate::model::error::http::HttpError;
use crate::model::error::misc::MiscError;
use crate::model::error::Error;
use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::log::system::SystemLog;
use crate::utils::logging::Logging;
use crate::web::api::{control, default, misc};
use crate::ml::model_loader::MLModels;
use crate::ml::config_loader::InferenceConfig;
use crate::ml::engine::Engine;
pub struct System {
pub app_config: Arc<AppConfig>,
pub ebpf_services: Arc<EbpfServices>,
pub ingress_ebpf: Ebpf,
pub egress_ebpf: Ebpf,
pub ml_models: Arc<MLModels>,
pub inference_config: Arc<InferenceConfig>,
pub ml_engine: Arc<Engine>,
#[allow(dead_code)]
ingress_program_array: ProgramArray<MapData>,
#[allow(dead_code)]
@ -43,38 +38,17 @@ impl System {
let (mut egress_ebpf, egress_program_array) = System::get_egress_ebpf()?;
let app_config = Arc::new(AppConfig::new()?);
let inference_config = Arc::new(InferenceConfig::load_file(&app_config.models_config_name)?);
let ml_models = Arc::new(MLModels::load_models(&app_config, inference_config.num_features())?);
let ebpf_services = Arc::new(EbpfServices::new(
app_config.clone(),
&mut ingress_ebpf,
&mut egress_ebpf,
)?);
let ml_engine = Arc::new(Engine::new(
ml_models.clone(),
inference_config.clone(),
app_config.max_concurrent_flows,
app_config.min_packets_for_inference,
app_config.inference_interval_secs,
));
log!(MLLog::EngineStarted {
max_flows: app_config.max_concurrent_flows,
min_packets: app_config.min_packets_for_inference,
interval_secs: app_config.inference_interval_secs
});
let system = System {
app_config,
ebpf_services,
ingress_ebpf,
egress_ebpf,
ml_models,
inference_config,
ml_engine,
ingress_program_array,
egress_program_array,
};
@ -85,24 +59,11 @@ impl System {
let ebpf_services = self.ebpf_services.clone();
Logging::initialize()?;
log!(SystemLog::Initializing);
log!(MLLog::ModelsLoaded { info: self.ml_models.get_model_info("deep_autoencoder") });
log!(MLLog::ModelsLoaded { info: self.ml_models.get_model_info("random_forest") });
log!(MLLog::ModelsLoaded { info: self.ml_models.get_model_info("mlp") });
log!(MLLog::ConfigLoaded {
features: self.inference_config.num_features(),
attacks: self.inference_config.num_attack_types()
});
self.aya_log_init()?;
log!(SystemLog::InitializeComplete);
self.attach_ebpf()?;
let _ml_handle = self.ml_engine.clone().start();
ebpf_services.run(self.ml_engine.clone()).await?;
ebpf_services.run().await?;
self.run_http_server().await?;
Ok(())
}
@ -156,8 +117,6 @@ impl System {
let service = self.ebpf_services.service.clone();
let statistics = self.ebpf_services.statistics.clone();
let health = self.ebpf_services.health.clone();
let ml_models = self.ml_models.clone();
let inference_config = self.inference_config.clone();
let port = self.app_config.http_server_bind_port;
HttpServer::new(move || {
let cors = actix_cors::Cors::default()
@ -172,8 +131,6 @@ 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_models.clone()))
.app_data(web::Data::from(inference_config.clone()))
.service(control::initialize())
.service(misc::initialize())
.default_service(route().to(default::default_route))
@ -241,4 +198,4 @@ impl System {
}
Ok(())
}
}
}

View File

@ -1,51 +0,0 @@
use std::collections::HashMap;
use std::time::{Duration, Instant};
use crate::model::ml_detection::FlowKey;
pub struct AttackAggregator {
detections: HashMap<FlowKey, Vec<(Instant, f32)>>,
window_duration: Duration,
min_detections: usize,
alert_threshold_multiplier: f32,
}
impl AttackAggregator {
pub fn new(window_secs: u64, min_detections: usize) -> Self {
Self {
detections: HashMap::new(),
window_duration: Duration::from_secs(window_secs),
min_detections,
alert_threshold_multiplier: 1.2,
}
}
pub fn should_alert(&mut self, flow_key: &FlowKey, ensemble_score: f32, threshold: f32) -> bool {
let now = Instant::now();
let detections = self.detections.entry(flow_key.clone()).or_default();
detections.retain(|(time, _)| now.duration_since(*time) < self.window_duration);
detections.push((now, ensemble_score));
if detections.len() >= self.min_detections {
let avg_score: f32 =
detections.iter().map(|(_, s)| s).sum::<f32>() / detections.len() as f32;
return avg_score > threshold * self.alert_threshold_multiplier;
}
false
}
pub fn cleanup(&mut self) {
let now = Instant::now();
self.detections.retain(|_, detections| {
detections.retain(|(time, _)| now.duration_since(*time) < self.window_duration);
!detections.is_empty()
});
}
pub fn tracked_flows(&self) -> usize {
self.detections.len()
}
}

View File

@ -1,44 +0,0 @@
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
use crate::model::error::ml::MLError;
use crate::model::ml_detection::{AENormalization, ClipParams};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceConfig {
pub threshold: f64,
pub strategy_name: String,
pub clip_params: HashMap<String, ClipParams>,
pub scaler_mean: Vec<f64>,
pub scaler_std: Vec<f64>,
pub post_clip_min: f64,
pub post_clip_max: f64,
pub ae_normalization: AENormalization,
pub attack_labels: HashMap<String, String>,
pub feature_names: Vec<String>,
}
impl InferenceConfig {
pub fn load_file(file: &str) -> Result<Self, MLError> {
let path = PathBuf::from("models").join(file);
let content = fs::read_to_string(&path)
.map_err(|_| MLError::ConfigLoadFailed { path: path.to_path_buf() })?;
let config: InferenceConfig = serde_json::from_str(&content)
.map_err(|e| MLError::ConfigParseFailed { reason: e.to_string() })?;
Ok(config)
}
pub fn num_features(&self) -> usize {
self.feature_names.len()
}
pub fn num_attack_types(&self) -> usize {
self.attack_labels.len()
}
pub fn get_attack_label(&self, id: usize) -> Option<&String> {
self.attack_labels.get(&id.to_string())
}
}

View File

@ -1,185 +0,0 @@
use std::time::{Duration, Instant};
use std::sync::{Arc, Mutex};
use macros::log;
use tokio::time::interval;
use super::config_loader::InferenceConfig;
use super::flow_tracker::FlowTracker;
use super::inference::Inference;
use super::model_loader::MLModels;
use super::aggregator::AttackAggregator;
use crate::utils::packet_parser::parse_packet;
use crate::model::log::ml::MLLog;
use crate::model::ml_detection::{EngineStats, InferenceStats};
pub struct Engine {
flow_tracker: Arc<FlowTracker>,
inference_pipeline: Arc<Inference>,
aggregator: Arc<Mutex<AttackAggregator>>,
min_packets: usize,
inference_interval_secs: u64,
}
impl Engine {
pub fn new(
models: Arc<MLModels>,
config: Arc<InferenceConfig>,
max_flows: usize,
min_packets: usize,
interval_secs: u64,
) -> Self {
let flow_tracker = Arc::new(FlowTracker::new(max_flows));
let inference_pipeline = Arc::new(Inference::new(models, config));
let aggregator = Arc::new(Mutex::new(AttackAggregator::new(30, 10)));
Self {
flow_tracker,
inference_pipeline,
aggregator,
min_packets,
inference_interval_secs: interval_secs,
}
}
pub fn start(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
self.run_inference_loop().await;
})
}
pub fn get_flow_tracker(&self) -> Arc<FlowTracker> {
self.flow_tracker.clone()
}
async fn run_inference_loop(&self) {
let mut ticker = interval(Duration::from_secs(self.inference_interval_secs));
loop {
ticker.tick().await;
let start = Instant::now();
let total_flows = self.flow_tracker.flow_count();
let all_flows = self.flow_tracker.get_flows_snapshot();
let packet_counts: Vec<usize> = all_flows.iter().map(|f| f.packet_count()).collect();
let flows = self
.flow_tracker
.get_flows_for_inference(self.min_packets);
log!(
MLLog::FlowStats(
total_flows,
flows.len(),
self.min_packets,
format!("{:?}", packet_counts)
)
);
if flows.is_empty() {
log!(
MLLog::InferenceSkipped(
format!(
"No flows with sufficient packets (total flows: {}, min packets: {})",
total_flows,
self.min_packets
)
)
);
continue;
}
let batch_size = flows.len().min(200);
let batch = &flows[..batch_size];
log!(MLLog::RunningInference(batch_size));
let results = self.inference_pipeline.infer_batch(batch);
let elapsed_us = start.elapsed().as_micros() as u64;
let stats = InferenceStats::from_results(&results, elapsed_us);
if results.len() != batch_size {
log!(MLLog::InferenceResults(batch_size, results.len()));
}
log!(
MLLog::InferenceCompleted(
stats.total_flows,
stats.malicious_flows,
stats.benign_flows,(elapsed_us as f64 / 1000.0) as u32,
stats.flows_per_second
)
);
if let Ok(mut aggregator) = self.aggregator.lock() {
for result in &results {
if result.is_attack {
let should_alert = aggregator.should_alert(
&result.flow_key_raw,
result.ensemble_score,
self.inference_pipeline.config.threshold as f32,
);
if should_alert {
log!(
MLLog::ThreatDetected(
result.flow_key.clone(),
result.attack_type.clone().unwrap_or_else(|| "UNKNOWN".to_string()),
result.confidence,
result.ae_score,
result.rf_score,
result.ensemble_score,
)
);
}
}
}
aggregator.cleanup();
}
self.flow_tracker.cleanup_old_flows(60_000_000);
}
}
pub fn process_packet(&self, packet_data: &[u8]) {
match parse_packet(packet_data) {
Some(packet_info) => {
self.flow_tracker.process_packet(packet_info);
}
None => {
log!(MLLog::ParsePacketFailed(packet_data.len()))
}
}
}
pub fn get_stats(&self) -> EngineStats {
EngineStats {
active_flows: self.flow_tracker.flow_count(),
}
}
}
pub struct PacketProcessor {
ml_engine: Arc<Engine>,
}
impl PacketProcessor {
pub fn new(ml_engine: Arc<Engine>) -> Self {
Self { ml_engine }
}
pub fn process(&self, packet_data: &[u8]) {
self.ml_engine.process_packet(packet_data);
}
pub fn process_batch(&self, packets: &[Vec<u8>]) {
for packet in packets {
self.process(packet);
}
}
}

View File

@ -1,264 +0,0 @@
use std::collections::HashMap;
use crate::model::ml_detection::{PacketData, ClipParams};
use super::flow_tracker::FlowData;
#[derive(Debug, Clone)]
pub struct FlowFeatures {
pub features: Vec<f64>,
pub feature_num: usize,
}
impl FlowFeatures {
pub fn extract(flow: &FlowData, feature_names: &[String]) -> Self {
let feature_num = feature_names.len();
let mut features = Vec::with_capacity(feature_num);
for name in feature_names {
let value = Self::get_feature_by_name(flow, name.trim());
features.push(value);
}
Self { features, feature_num }
}
fn get_feature_by_name(flow: &FlowData, feature_name: &str) -> f64 {
let safe_div = |a: f64, b: f64| if b > 0.0 { a / b } else { 0.0 };
// 1-5
let fwd_count = flow.fwd_packets.len() as f64;
let bwd_count = flow.bwd_packets.len() as f64;
let total_count = fwd_count + bwd_count;
let duration_us = flow.duration_us() as f64;
let duration_s = duration_us / 1_000_000.0;
let duration_s = if duration_s > 0.0 { duration_s } else { 1e-6 };
// 6-9
let fwd_lengths: Vec<f64> = flow.fwd_packets.iter().map(|p| p.length as f64).collect();
let (fwd_max, fwd_min, fwd_mean, fwd_std) = compute_stats(&fwd_lengths);
// 10-13
let bwd_lengths: Vec<f64> = flow.bwd_packets.iter().map(|p| p.length as f64).collect();
let (bwd_max, bwd_min, bwd_mean, bwd_std) = compute_stats(&bwd_lengths);
// 14-15
let total_bytes = (flow.fwd_total_bytes + flow.bwd_total_bytes) as f64;
// 16-19
let flow_iats = compute_flow_iats(&flow.fwd_packets, &flow.bwd_packets);
let (flow_iat_mean, flow_iat_std, flow_iat_max, flow_iat_min) = compute_stats(&flow_iats);
// 20-24
let fwd_iats = compute_iats(&flow.fwd_packets);
let fwd_iat_total: f64 = fwd_iats.iter().sum();
let (fwd_iat_mean, fwd_iat_std, fwd_iat_max, fwd_iat_min) = compute_stats(&fwd_iats);
// 25-29
let bwd_iats = compute_iats(&flow.bwd_packets);
let bwd_iat_total: f64 = bwd_iats.iter().sum();
let (bwd_iat_mean, bwd_iat_std, bwd_iat_max, bwd_iat_min) = compute_stats(&bwd_iats);
// 30-37
let fwd_psh = flow.fwd_packets.iter().filter(|p| p.flags.psh).count() as f64;
let bwd_psh = flow.bwd_packets.iter().filter(|p| p.flags.psh).count() as f64;
let fwd_urg = flow.fwd_packets.iter().filter(|p| p.flags.urg).count() as f64;
let bwd_urg = flow.bwd_packets.iter().filter(|p| p.flags.urg).count() as f64;
// 38-55
let all_lengths: Vec<f64> = flow
.fwd_packets
.iter()
.chain(flow.bwd_packets.iter())
.map(|p| p.length as f64)
.collect();
let (max_len, min_len, mean_len, std_len) = compute_stats(&all_lengths);
// 56-67
let fwd_bulk = &flow.fwd_bulk_state;
let bwd_bulk = &flow.bwd_bulk_state;
// 68-69
let fwd_seg_sizes: Vec<f64> = flow.fwd_packets
.iter()
.filter(|p| p.payload_length > 0)
.map(|p| p.header_length as f64)
.collect();
// 70-73
let (active_mean, active_std, active_max, active_min) = compute_stats(
&flow
.active_periods
.iter()
.map(|&x| x as f64)
.collect::<Vec<_>>(),
);
// 74-77
let (idle_mean, idle_std, idle_max, idle_min) = compute_stats(
&flow
.idle_periods
.iter()
.map(|&x| x as f64)
.collect::<Vec<_>>(),
);
match feature_name {
"Destination Port" => flow.flow_key.dst_port as f64,
"Flow Duration" => duration_us,
"Total Fwd Packets" => fwd_count,
"Total Backward Packets" => bwd_count,
"Total Length of Fwd Packets" => flow.fwd_total_bytes as f64,
"Total Length of Bwd Packets" => flow.bwd_total_bytes as f64,
"Fwd Packet Length Max" => fwd_max,
"Fwd Packet Length Min" => fwd_min,
"Fwd Packet Length Mean" => fwd_mean,
"Fwd Packet Length Std" => fwd_std,
"Bwd Packet Length Max" => bwd_max,
"Bwd Packet Length Min" => bwd_min,
"Bwd Packet Length Mean" => bwd_mean,
"Bwd Packet Length Std" => bwd_std,
"Flow Bytes/s" => safe_div(total_bytes, duration_s),
"Flow Packets/s" => safe_div(total_count, duration_s),
"Flow IAT Mean" => flow_iat_mean,
"Flow IAT Std" => flow_iat_std,
"Flow IAT Max" => flow_iat_max,
"Flow IAT Min" => flow_iat_min,
"Fwd IAT Total" => fwd_iat_total,
"Fwd IAT Mean" => fwd_iat_mean,
"Fwd IAT Std" => fwd_iat_std,
"Fwd IAT Max" => fwd_iat_max,
"Fwd IAT Min" => fwd_iat_min,
"Bwd IAT Total" => bwd_iat_total,
"Bwd IAT Mean" => bwd_iat_mean,
"Bwd IAT Std" => bwd_iat_std,
"Bwd IAT Max" => bwd_iat_max,
"Bwd IAT Min" => bwd_iat_min,
"Fwd PSH Flags" => fwd_psh,
"Bwd PSH Flags" => bwd_psh,
"Fwd URG Flags" => fwd_urg,
"Bwd URG Flags" => bwd_urg,
"Fwd Header Length" => flow.fwd_header_bytes as f64,
"Bwd Header Length" => flow.bwd_header_bytes as f64,
"Fwd Packets/s" => safe_div(fwd_count, duration_s),
"Bwd Packets/s" => safe_div(bwd_count, duration_s),
"Min Packet Length" => min_len,
"Max Packet Length" => max_len,
"Packet Length Mean" => mean_len,
"Packet Length Std" => std_len,
"Packet Length Variance" => std_len * std_len,
"FIN Flag Count" => flow.fin_count as f64,
"SYN Flag Count" => flow.syn_count as f64,
"RST Flag Count" => flow.rst_count as f64,
"PSH Flag Count" => flow.psh_count as f64,
"ACK Flag Count" => flow.ack_count as f64,
"URG Flag Count" => flow.urg_count as f64,
"CWE Flag Count" => flow.cwe_count as f64,
"ECE Flag Count" => flow.ece_count as f64,
"Down/Up Ratio" => safe_div(bwd_count, fwd_count),
"Average Packet Size" => safe_div(total_bytes, total_count),
"Avg Fwd Segment Size" => safe_div(flow.fwd_total_bytes as f64, fwd_count),
"Avg Bwd Segment Size" => safe_div(flow.bwd_total_bytes as f64, bwd_count),
"Fwd Header Length.1" => flow.fwd_header_bytes as f64,
"Fwd Avg Bytes/Bulk" => safe_div(fwd_bulk.total_bytes as f64, fwd_bulk.bulk_count as f64),
"Fwd Avg Packets/Bulk" => safe_div(fwd_bulk.total_packets as f64, fwd_bulk.bulk_count as f64),
"Fwd Avg Bulk Rate" => safe_div(fwd_bulk.total_bytes as f64, duration_s),
"Bwd Avg Bytes/Bulk" => safe_div(bwd_bulk.total_bytes as f64, bwd_bulk.bulk_count as f64),
"Bwd Avg Packets/Bulk" => safe_div(bwd_bulk.total_packets as f64, bwd_bulk.bulk_count as f64),
"Bwd Avg Bulk Rate" => safe_div(bwd_bulk.total_bytes as f64, duration_s),
"Subflow Fwd Packets" => fwd_count,
"Subflow Fwd Bytes" => flow.fwd_total_bytes as f64,
"Subflow Bwd Packets" => bwd_count,
"Subflow Bwd Bytes" => flow.bwd_total_bytes as f64,
"Init_Win_bytes_forward" => flow.init_win_bytes_fwd as f64,
"Init_Win_bytes_backward" => flow.init_win_bytes_bwd as f64,
"act_data_pkt_fwd" => fwd_seg_sizes.len() as f64,
"min_seg_size_forward" => fwd_seg_sizes.iter()
.min_by(|a, b| a.total_cmp(b))
.copied()
.unwrap_or(0.0),
"Active Mean" => active_mean,
"Active Std" => active_std,
"Active Max" => active_max,
"Active Min" => active_min,
"Idle Mean" => idle_mean,
"Idle Std" => idle_std,
"Idle Max" => idle_max,
"Idle Min" => idle_min,
_ => {
0.0
}
}
}
pub fn normalize(&mut self, means: &[f64], stds: &[f64]) {
for i in 0..self.feature_num {
if stds[i] > 0.0 {
self.features[i] = (self.features[i] - means[i]) / stds[i];
} else {
self.features[i] = 0.0;
}
}
}
pub fn clip(&mut self, clip_min: f64, clip_max: f64) {
for i in 0..self.feature_num {
self.features[i] = self.features[i].max(clip_min).min(clip_max);
}
}
/// Winsorization: Clip each feature according to its specific bounds from clip_params
pub fn winsorize(&mut self, clip_params: &HashMap<String, ClipParams>, feature_names: &[String]) {
for (i, feature_name) in feature_names.iter().enumerate() {
if i < self.feature_num {
if let Some(params) = clip_params.get(feature_name) {
self.features[i] = self.features[i].clamp(params.lower, params.upper);
}
}
}
}
}
fn compute_stats(values: &[f64]) -> (f64, f64, f64, f64) {
if values.is_empty() {
return (0.0, 0.0, 0.0, 0.0);
}
let n = values.len() as f64;
let sum: f64 = values.iter().sum();
let mean = sum / n;
let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
let variance: f64 = values.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / n;
let std = variance.sqrt();
(max, min, mean, std)
}
fn compute_iats(packets: &[PacketData]) -> Vec<f64> {
if packets.len() < 2 {
return vec![0.0];
}
packets
.windows(2)
.map(|w| (w[1].timestamp_us - w[0].timestamp_us) as f64)
.collect()
}
fn compute_flow_iats(fwd_packets: &[PacketData], bwd_packets: &[PacketData]) -> Vec<f64> {
let mut all_packets: Vec<&PacketData> = fwd_packets.iter().chain(bwd_packets.iter()).collect();
all_packets.sort_by_key(|p| p.timestamp_us);
if all_packets.len() < 2 {
return vec![0.0];
}
all_packets
.windows(2)
.map(|w| (w[1].timestamp_us - w[0].timestamp_us) as f64)
.collect()
}

View File

@ -1,262 +0,0 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time;
use crate::model::ml_detection::{BulkState, FlowKey, PacketData};
use common::model::event::Event;
#[derive(Debug, Clone)]
pub struct FlowData {
pub flow_key: FlowKey,
pub start_time_us: u64,
pub last_time_us: u64,
pub fwd_packets: Vec<PacketData>,
pub fwd_total_bytes: u64,
pub fwd_header_bytes: u64,
pub bwd_packets: Vec<PacketData>,
pub bwd_total_bytes: u64,
pub bwd_header_bytes: u64,
pub fin_count: u32,
pub syn_count: u32,
pub rst_count: u32,
pub psh_count: u32,
pub ack_count: u32,
pub urg_count: u32,
pub cwe_count: u32,
pub ece_count: u32,
pub init_win_bytes_fwd: u16,
pub init_win_bytes_bwd: u16,
pub active_periods: Vec<u64>,
pub idle_periods: Vec<u64>,
pub last_packet_time: u64,
pub fwd_bulk_state: BulkState,
pub bwd_bulk_state: BulkState,
}
impl FlowData {
pub fn new(flow_key: FlowKey, first_packet: &Event) -> Self {
Self {
flow_key,
start_time_us: first_packet.timestamp_us(),
last_time_us: first_packet.timestamp_us(),
fwd_packets: Vec::new(),
fwd_total_bytes: 0,
fwd_header_bytes: 0,
bwd_packets: Vec::new(),
bwd_total_bytes: 0,
bwd_header_bytes: 0,
fin_count: 0,
syn_count: 0,
rst_count: 0,
psh_count: 0,
ack_count: 0,
urg_count: 0,
cwe_count: 0,
ece_count: 0,
init_win_bytes_fwd: if first_packet.is_forward() {
first_packet.tcp_window_size()
} else {
0
},
init_win_bytes_bwd: if !first_packet.is_forward() {
first_packet.tcp_window_size()
} else {
0
},
active_periods: Vec::new(),
idle_periods: Vec::new(),
last_packet_time: first_packet.timestamp_us(),
fwd_bulk_state: BulkState::default(),
bwd_bulk_state: BulkState::default(),
}
}
pub fn add_packet(&mut self, packet: &Event) {
let packet_data = PacketData {
timestamp_us: packet.timestamp_us(),
length: packet.packet_length(),
header_length: packet.header_length(),
payload_length: packet.payload_length(),
flags: packet.tcp_flags().clone(),
};
if packet.tcp_flags().fin {
self.fin_count += 1;
}
if packet.tcp_flags().syn {
self.syn_count += 1;
}
if packet.tcp_flags().rst {
self.rst_count += 1;
}
if packet.tcp_flags().psh {
self.psh_count += 1;
}
if packet.tcp_flags().ack {
self.ack_count += 1;
}
if packet.tcp_flags().urg {
self.urg_count += 1;
}
if packet.tcp_flags().cwr {
self.cwe_count += 1;
}
if packet.tcp_flags().ece {
self.ece_count += 1;
}
let iat = packet.timestamp_us().saturating_sub(self.last_packet_time);
const IDLE_THRESHOLD_US: u64 = 1_000_000;
if iat > IDLE_THRESHOLD_US {
self.idle_periods.push(iat);
} else if iat > 0 {
self.active_periods.push(iat);
}
self.last_packet_time = packet.timestamp_us();
self.last_time_us = packet.timestamp_us();
if packet.is_forward() {
self.fwd_packets.push(packet_data.clone());
self.fwd_total_bytes += packet.packet_length() as u64;
self.fwd_header_bytes += packet.header_length() as u64;
if self.init_win_bytes_fwd == 0 {
self.init_win_bytes_fwd = packet.tcp_window_size();
}
Self::update_bulk_state(&mut self.fwd_bulk_state, &packet_data);
} else {
self.bwd_packets.push(packet_data.clone());
self.bwd_total_bytes += packet.packet_length() as u64;
self.bwd_header_bytes += packet.header_length() as u64;
if self.init_win_bytes_bwd == 0 {
self.init_win_bytes_bwd = packet.tcp_window_size();
}
Self::update_bulk_state(&mut self.bwd_bulk_state, &packet_data);
}
}
fn update_bulk_state(bulk_state: &mut BulkState, packet: &PacketData) {
const BULK_MIN_PACKETS: u64 = 4;
const BULK_MIN_BYTES: u64 = 1000;
if packet.payload_length > 0 {
if !bulk_state.in_bulk {
bulk_state.in_bulk = true;
bulk_state.last_bulk_bytes = packet.length as u64;
bulk_state.last_bulk_packets = 1;
} else {
bulk_state.last_bulk_bytes += packet.length as u64;
bulk_state.last_bulk_packets += 1;
}
} else {
if bulk_state.in_bulk
&& bulk_state.last_bulk_packets >= BULK_MIN_PACKETS
&& bulk_state.last_bulk_bytes >= BULK_MIN_BYTES
{
bulk_state.bulk_count += 1;
bulk_state.total_bytes += bulk_state.last_bulk_bytes;
bulk_state.total_packets += bulk_state.last_bulk_packets;
}
bulk_state.in_bulk = false;
bulk_state.last_bulk_bytes = 0;
bulk_state.last_bulk_packets = 0;
}
}
pub fn duration_us(&self) -> u64 {
self.last_time_us.saturating_sub(self.start_time_us)
}
pub fn packet_count(&self) -> usize {
self.fwd_packets.len() + self.bwd_packets.len()
}
}
pub struct FlowTracker {
flows: Arc<Mutex<HashMap<FlowKey, FlowData>>>,
max_flows: usize,
}
impl FlowTracker {
pub fn new(max_flows: usize) -> Self {
Self {
flows: Arc::new(Mutex::new(HashMap::new())),
max_flows,
}
}
pub fn process_packet(&self, mut packet: Event) {
let flow_key = FlowKey::from_packet(&packet);
let reverse_key = flow_key.reverse();
let Ok(mut flows) = self.flows.lock() else {
return;
};
let (actual_key, is_forward) = if flows.contains_key(&flow_key) {
(flow_key, true)
} else if flows.contains_key(&reverse_key) {
(reverse_key, false)
} else {
(flow_key, true)
};
packet.set_is_forward(is_forward);
let flow = flows.entry(actual_key.clone()).or_insert_with(|| {
FlowData::new(actual_key, &packet)
});
flow.add_packet(&packet);
if flows.len() > self.max_flows {
if let Some(key) = flows.keys().next().cloned() {
flows.remove(&key);
}
}
}
pub fn get_flows_snapshot(&self) -> Vec<FlowData> {
let Ok(flows) = self.flows.lock() else {
return Vec::new();
};
flows.values().cloned().collect()
}
pub fn get_flows_for_inference(&self, min_packets: usize) -> Vec<FlowData> {
let Ok(flows) = self.flows.lock() else {
return Vec::new();
};
flows
.values()
.filter(|flow| flow.packet_count() >= min_packets)
.cloned()
.collect()
}
pub fn cleanup_old_flows(&self, max_age_us: u64) {
let now = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
.map(|d| d.as_micros() as u64)
.unwrap_or(0);
let Ok(mut flows) = self.flows.lock() else {
return;
};
flows.retain(|_, flow| {
now.saturating_sub(flow.last_time_us) < max_age_us
});
}
pub fn flow_count(&self) -> usize {
let Ok(flows) = self.flows.lock() else {
return 0;
};
flows.len()
}
}

View File

@ -1,226 +0,0 @@
use std::sync::Arc;
use tract_onnx::prelude::*;
use macros::log;
use super::flow_tracker::FlowData;
use super::config_loader::InferenceConfig;
use super::feature_extractor::FlowFeatures;
use super::model_loader::MLModels;
use crate::model::ml_detection::DetectionResult;
use crate::model::log::ml::MLLog;
pub struct Inference {
pub models: Arc<MLModels>,
pub config: Arc<InferenceConfig>,
}
impl Inference {
pub fn new(models: Arc<MLModels>, config: Arc<InferenceConfig>) -> Self {
Self { models, config }
}
pub fn infer_batch(&self, flows: &[FlowData]) -> Vec<DetectionResult> {
flows
.iter()
.filter_map(|flow| self.infer_single(flow))
.collect()
}
pub fn infer_single(&self, flow: &FlowData) -> Option<DetectionResult> {
// extract
let mut features = FlowFeatures::extract(
flow,
&self.config.feature_names
);
// pre-process
features.winsorize(&self.config.clip_params, &self.config.feature_names);
features.normalize(&self.config.scaler_mean, &self.config.scaler_std);
features.clip(self.config.post_clip_min, self.config.post_clip_max);
// input tensor
let input = tract_ndarray::Array2::from_shape_fn((1, self.config.num_features()), |(_, j)| {
features.features[j] as f32
});
// Deep Autoencoder
let ae_score = match self.run_autoencoder(&input) {
Ok(score) => score,
Err(e) => {
log!(MLLog::InferenceFailed("DeepAutoEncoder".to_string(), e.to_string()));
return None;
}
};
// Random Forest
let rf_score = match self.run_random_forest(&input) {
Ok(score) => score,
Err(e) => {
log!(MLLog::InferenceFailed("RandomForest".to_string(), e.to_string()));
return None;
}
};
// Ensemble score
let ensemble_score = self.compute_ensemble_score(ae_score, rf_score);
let is_anomaly = ensemble_score > self.config.threshold as f32;
let flow_key = format!(
"{}:{} -> {}:{} (proto {})",
flow.flow_key.src_ip,
flow.flow_key.src_port,
flow.flow_key.dst_ip,
flow.flow_key.dst_port,
flow.flow_key.protocol
);
let flow_key_raw = flow.flow_key.clone();
if is_anomaly {
// MLP
let (attack_type, confidence) = match self.run_mlp(&input) {
Ok((attack_type, conf)) => (attack_type, conf),
Err(e) => {
log!(MLLog::InferenceFailed("MLP".to_string(), e.to_string()));
("UNKNOWN".to_string(), ensemble_score)
}
};
if confidence < 0.75 {
return Some(DetectionResult {
flow_key,
flow_key_raw,
is_attack: false,
attack_type: None,
confidence: 1.0 - ensemble_score,
ae_score,
rf_score,
ensemble_score,
});
}
Some(DetectionResult {
flow_key,
flow_key_raw,
is_attack: true,
attack_type: Some(attack_type),
confidence,
ae_score,
rf_score,
ensemble_score,
})
} else {
Some(DetectionResult {
flow_key,
flow_key_raw,
is_attack: false,
attack_type: None,
confidence: 1.0 - ensemble_score,
ae_score,
rf_score,
ensemble_score,
})
}
}
fn run_autoencoder(&self, input: &tract_ndarray::Array2<f32>) -> TractResult<f32> {
let input_tensor = input.clone().into_tensor();
let result = self
.models
.deep_autoencoder
.run(tvec![input_tensor.into()])?;
let output = result[0]
.to_array_view::<f32>()?
.into_dimensionality::<tract_ndarray::Ix2>()?;
let diff = input - &output;
let squared_errors = &diff * &diff;
let mse = squared_errors.sum() / self.config.num_features() as f32;
let ae_norm = &self.config.ae_normalization;
let ae_score = (mse - ae_norm.min as f32) / (ae_norm.max as f32 - ae_norm.min as f32 + 1e-10);
let ae_score = ae_score.clamp(0.0, 1.0);
Ok(ae_score)
}
fn run_random_forest(&self, input: &tract_ndarray::Array2<f32>) -> TractResult<f32> {
let input_tensor = input.clone().into_tensor();
let result = self.models.random_forest.run(tvec![input_tensor.into()])?;
// output[0] = output_label (i64)
// output[1] = output_probability (sequence of maps)
if result.len() > 1 {
if let Ok(proba) = result[1].to_array_view::<f32>() {
if proba.len() > 1 {
return Ok(proba.iter().nth(1).copied().unwrap_or(0.0));
} else if proba.len() == 1 {
return Ok(proba.iter().next().copied().unwrap_or(0.0));
}
}
}
let label = result[0].to_array_view::<i64>()?;
let prediction = label.iter().next().copied().unwrap_or(0);
Ok(if prediction != 0 { 1.0 } else { 0.0 })
}
fn run_mlp(&self, input: &tract_ndarray::Array2<f32>) -> TractResult<(String, f32)> {
let input_tensor = input.clone().into_tensor();
let result = self.models.mlp.run(tvec![input_tensor.into()])?;
let output = result[0].to_array_view::<f32>()?;
let mut max_prob: f32 = 0.0;
let mut predicted_class: usize = 0;
for (i, &prob) in output.iter().enumerate() {
if prob > max_prob {
max_prob = prob;
predicted_class = i;
}
}
let attack_type = self.config
.attack_labels
.get(&predicted_class.to_string())
.cloned()
.unwrap_or_else(|| "UNKNOWN".to_string());
Ok((attack_type, max_prob))
}
fn compute_ensemble_score(&self, ae_score: f32, rf_score: f32) -> f32 {
let strategy = &self.config.strategy_name;
if strategy.starts_with("W_") {
if let Some(weights_str) = strategy.strip_prefix("W_") {
let parts: Vec<&str> = weights_str.split(':').collect();
if parts.len() == 2 {
if let (Ok(w1), Ok(w2)) = (parts[0].parse::<f32>(), parts[1].parse::<f32>()) {
let w1 = w1 / 10.0;
let w2 = w2 / 10.0;
return w1 * ae_score + w2 * rf_score;
}
}
}
(ae_score + rf_score) / 2.0
} else {
match strategy.as_str() {
"Average" => (ae_score + rf_score) / 2.0,
"Max" => ae_score.max(rf_score),
"Min" => ae_score.min(rf_score),
"Product" => ae_score * rf_score,
_ => (ae_score + rf_score) / 2.0,
}
}
}
}

View File

@ -1,7 +0,0 @@
pub mod model_loader;
pub mod config_loader;
pub mod flow_tracker;
pub mod feature_extractor;
pub mod inference;
pub mod engine;
pub mod aggregator;

View File

@ -1,61 +0,0 @@
use tract_onnx::prelude::*;
use std::path::PathBuf;
use crate::core::infrastructure::app_config::AppConfig;
use crate::model::error::ml::MLError;
use crate::model::ml_detection::RunnableModel;
pub struct MLModels {
pub deep_autoencoder: RunnableModel,
pub random_forest: RunnableModel,
pub mlp: RunnableModel,
}
impl MLModels {
pub fn load_models(app_config: &Arc<AppConfig>, features: usize) -> Result<Self, MLError> {
Ok(Self {
deep_autoencoder: Self::loader(&app_config.deep_autoencoder_name, features)?,
random_forest: Self::loader(&app_config.random_forest_name, features)?,
mlp: Self::loader(&app_config.mlp_name, features)?,
})
}
pub fn loader(model: &str, features: usize) -> Result<RunnableModel, MLError> {
let model_path = PathBuf::from("models").join(model);
let mut model = onnx()
.model_for_path(&model_path)
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?;
model.set_input_fact(0, f32::fact(&[1, features]).into())
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?;
let runnable_model = model
.into_optimized()
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path.clone() }
})?
.into_runnable()
.map_err(|_| {
MLError::ModelLoadFailed { path: model_path }
})?;
Ok(runnable_model)
}
pub fn get_model_info(&self, name: &str) -> String {
let model = match name {
"deep_autoencoder" => &self.deep_autoencoder,
"random_forest" => &self.random_forest,
"mlp" => &self.mlp,
_ => return "unknown model".to_string(),
};
let inputs = model.model().inputs.len();
let outputs = model.model().outputs.len();
format!("{}: inputs: {}, outputs: {}", name, inputs, outputs)
}
}

View File

@ -1,4 +1,5 @@
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
#[derive(Debug, Deserialize)]
pub struct ConfigTable {
@ -10,13 +11,9 @@ pub struct ConfigTable {
pub struct Config {
pub ingress_ifname: String,
pub egress_ifname: String,
pub geoip_db_name: String,
pub deep_autoencoder_name: String,
pub random_forest_name: String,
pub mlp_name: String,
pub models_config_name: String,
pub geoip_db_path: PathBuf,
pub combined_queue_count: u32,
pub channel_size: usize,
pub xsk_channel_size: usize,
pub fill_queue_size: u32,
pub comp_queue_size: u32,
pub tx_queue_size: u32,
@ -24,8 +21,5 @@ pub struct Config {
pub frame_size: u32,
pub frame_count: u32,
pub refresh_interval: u64,
pub http_server_bind_port: u16,
pub max_concurrent_flows: usize,
pub min_packets_for_inference: usize,
pub inference_interval_secs: u64,
pub http_server_bind_port: u16
}

View File

@ -1,4 +1,4 @@
use macros::traceable;
use macros::{loggable, traceable};
use tracing;
traceable! {

View File

@ -4,20 +4,7 @@ use macros::traceable;
traceable! {
MLError {
#[no_source]
#[error("Initialize Machine Learning detection failed")]
InitializeFailed => tracing::Level::ERROR,
#[no_source]
#[error("Failed to load ONNX model from: {path:?}")]
ModelLoadFailed { path: PathBuf } => tracing::Level::ERROR,
#[no_source]
#[error("Failed to load inference configuration from: {path:?}")]
ConfigLoadFailed { path: PathBuf } => tracing::Level::ERROR,
#[no_source]
#[error("Failed to parse inference configuration: {reason}")]
ConfigParseFailed { reason: String } => tracing::Level::ERROR,
}
}
}

View File

@ -59,9 +59,3 @@ impl From<SystemError> for Error {
Self::System(error)
}
}
impl From<MLError> for Error {
fn from(error: MLError) -> Self {
Self::ML(error)
}
}

View File

@ -13,10 +13,12 @@ impl NativeConvert for IPv4 {
type Native = Ipv4Addr;
fn into_native(self) -> Self::Native {
// eBPF 儲存的是 big-endian需要轉換成 host order
Ipv4Addr::from(u32::from_be(self))
}
fn from_native(native: Self::Native) -> Self {
// 轉回 big-endian 給 eBPF
native.to_bits().to_be()
}
}

View File

@ -41,38 +41,5 @@ loggable! {
#[error("Queue pair {queue_id} started successfully")]
QueuePairStarted { queue_id: u32 } => tracing::Level::INFO,
#[error("ML models loaded - {info}")]
ModelsLoaded { info: String } => tracing::Level::INFO,
#[error("Inference configuration loaded: {features} features, {attacks} attack types")]
ConfigLoaded { features: usize, attacks: usize } => tracing::Level::INFO,
#[error("ML Engine started: max_flows={max_flows}, min_packets={min_packets}, interval={interval_secs}s")]
EngineStarted { max_flows: usize, min_packets: usize, interval_secs: u64 } => tracing::Level::INFO,
#[error("Inference completed: {total_flows} flows ({anomaly} anomaly, {benign} benign) in {duration_ms}ms ({throughput:.1} flows/s)")]
InferenceCompleted { total_flows: usize, anomaly: usize, benign: usize, duration_ms: u32, throughput: f32 } => tracing::Level::INFO,
#[error("Inference skipped: {reason}")]
InferenceSkipped { reason: String } => tracing::Level::INFO,
#[error("Threat detected: {flow} -> {attack_type} (confidence: {confidence:.2}, ae_score: {ae_score:.4}, rf_score: {rf_score:.4}, ensemble: {ensemble_score:.4})")]
ThreatDetected { flow: String, attack_type: String, confidence: f32, ae_score: f32, rf_score: f32, ensemble_score: f32 } => tracing::Level::WARN,
#[error("Flow stats: total={total_flows}, qualified={flows_len}, min_packets={min_packets}, packet_counts: {counts}")]
FlowStats { total_flows: usize, flows_len: usize, min_packets: usize, counts: String } => tracing::Level::INFO,
#[error("Running inference on {size} flows")]
RunningInference { size: usize } => tracing::Level::INFO,
#[error("Inference returned fewer results: expected {size}, got {len}")]
InferenceResults { size: usize, len: usize } => tracing::Level::INFO,
#[error("{model} inference failed: {error}")]
InferenceFailed { model: String, error: String } => tracing::Level::INFO,
#[error("Failed to parse packet (length: {len})")]
ParsePacketFailed { len: usize } => tracing::Level::INFO,
}
}
}

View File

@ -2,4 +2,4 @@ pub mod ebpf;
pub mod http;
pub mod ml;
pub mod system;
pub mod misc;
mod misc;

View File

@ -1,131 +0,0 @@
use serde::{Deserialize, Serialize};
use tract_onnx::prelude::{Graph, SimplePlan, TypedFact, TypedOp};
use common::model::event::{Event, TcpFlags};
use crate::utils::packet_parser::{format_ipv4, format_ipv6};
pub type RunnableModel = SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClipParams {
pub lower: f64,
pub upper: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AENormalization {
pub min: f64,
pub max: f64,
pub mean: f64,
pub std: f64,
pub median: f64,
pub p90: f64,
pub p95: f64,
pub p99: f64,
}
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct FlowKey {
pub src_ip: String,
pub dst_ip: String,
pub src_port: u16,
pub dst_port: u16,
pub protocol: u8,
}
impl FlowKey {
pub fn from_packet(packet: &Event) -> Self {
match packet {
Event::IPv4(ipv4) => Self {
src_ip: format_ipv4(ipv4.src_ip),
dst_ip: format_ipv4(ipv4.dst_ip),
src_port: ipv4.src_port,
dst_port: ipv4.dst_port,
protocol: ipv4.protocol as u8,
},
Event::IPv6(ipv6) => Self {
src_ip: format_ipv6(ipv6.src_ip),
dst_ip: format_ipv6(ipv6.dst_ip),
src_port: ipv6.src_port,
dst_port: ipv6.dst_port,
protocol: ipv6.protocol as u8,
},
}
}
pub fn reverse(&self) -> Self {
Self {
src_ip: self.dst_ip.clone(),
dst_ip: self.src_ip.clone(),
src_port: self.dst_port,
dst_port: self.src_port,
protocol: self.protocol,
}
}
}
#[derive(Debug, Clone)]
pub struct PacketData {
pub timestamp_us: u64,
pub length: u32,
pub header_length: u16,
pub payload_length: u32,
pub flags: TcpFlags,
}
#[derive(Debug, Clone, Default)]
pub struct BulkState {
pub bulk_count: u32,
pub total_bytes: u64,
pub total_packets: u64,
pub last_bulk_bytes: u64,
pub last_bulk_packets: u64,
pub in_bulk: bool,
}
#[derive(Debug, Clone)]
pub struct DetectionResult {
pub flow_key: String,
pub flow_key_raw: FlowKey,
pub is_attack: bool,
pub attack_type: Option<String>,
pub confidence: f32,
pub ae_score: f32,
pub rf_score: f32,
pub ensemble_score: f32,
}
#[derive(Debug, Clone, Default)]
pub struct InferenceStats {
pub total_flows: usize,
pub malicious_flows: usize,
pub benign_flows: usize,
pub inference_time_us: u64,
pub flows_per_second: f32,
}
impl InferenceStats {
pub fn from_results(results: &[DetectionResult], elapsed_us: u64) -> Self {
let total = results.len();
let malicious = results.iter().filter(|r| r.is_attack).count();
let benign = total - malicious;
let fps = if elapsed_us > 0 {
(total as f64 / (elapsed_us as f64 / 1_000_000.0)) as f32
} else {
0.0
};
Self {
total_flows: total,
malicious_flows: malicious,
benign_flows: benign,
inference_time_us: elapsed_us,
flows_per_second: fps,
}
}
}
#[derive(Debug, Clone)]
pub struct EngineStats {
pub active_flows: usize,
}

View File

@ -6,4 +6,3 @@ pub mod ip_address;
pub mod list_type;
pub mod log;
pub mod time_type;
pub mod ml_detection;

View File

@ -1,9 +1,7 @@
use std::time;
use std::mem;
use common::model::event::{Event, IPv4Event, IPv6Event, TcpFlags};
use common::model::event::{Event, IPv4Event, IPv6Event};
use network_types::ip::IpProto;
/// Parse raw packet bytes into an Event
pub fn parse_packet(packet_data: &[u8]) -> Option<Event> {
if packet_data.len() < 14 {
return None;
@ -11,158 +9,104 @@ pub fn parse_packet(packet_data: &[u8]) -> Option<Event> {
let eth_type = u16::from_be_bytes([packet_data[12], packet_data[13]]);
let timestamp_us = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_micros() as u64;
.as_secs();
match eth_type {
0x0800 => parse_ipv4(packet_data, timestamp_us),
0x86DD => parse_ipv6(packet_data, timestamp_us),
0x0800 => parse_ipv4(packet_data, timestamp),
0x86DD => parse_ipv6(packet_data, timestamp),
_ => None,
}
}
fn parse_ipv4(packet_data: &[u8], timestamp_us: u64) -> Option<Event> {
fn parse_ipv4(packet_data: &[u8], timestamp: u64) -> Option<Event> {
// Ethernet header (14) + minimum IPv4 header (20) = 34 bytes
if packet_data.len() < 34 {
return None;
}
let ip_header = &packet_data[14..];
let protocol_byte = ip_header[9];
let protocol = unsafe { mem::transmute::<u8, IpProto>(protocol_byte) };
let src_ip = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]);
let dst_ip = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]);
// Parse IPv4 header
let protocol = ip_header[9];
let source_ip = u32::from_be_bytes([ip_header[12], ip_header[13], ip_header[14], ip_header[15]]);
let destination_ip = u32::from_be_bytes([ip_header[16], ip_header[17], ip_header[18], ip_header[19]]);
// Get IP header length
let ihl = (ip_header[0] & 0x0F) as usize * 4;
// Total length
let total_len = u16::from_be_bytes([ip_header[2], ip_header[3]]) as u32;
if packet_data.len() < 14 + ihl + 4 {
return None;
}
let transport_header = &ip_header[ihl..];
let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]);
let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]);
let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 {
if packet_data.len() < 14 + ihl + 20 {
return None;
}
let data_offset = (transport_header[12] >> 4) as u16 * 4;
let flags = TcpFlags::from_byte(transport_header[13]);
let window = u16::from_be_bytes([transport_header[14], transport_header[15]]);
(flags, window, data_offset)
} else if protocol_byte == 17 {
(TcpFlags::default(), 0, 8)
// Parse transport layer (TCP/UDP)
let (source_port, destination_port) = if packet_data.len() >= 14 + ihl + 4 {
let transport_header = &ip_header[ihl..];
let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]);
let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]);
(src_port, dst_port)
} else {
(TcpFlags::default(), 0, 0)
(0, 0)
};
let payload_length = total_len.saturating_sub(ihl as u32 + header_length as u32);
let event = IPv4Event {
protocol,
src_ip,
dst_ip,
src_port,
dst_port,
packet_length: total_len,
payload_length,
header_length,
timestamp_us,
tcp_flags,
tcp_window_size,
is_forward: false,
protocol: unsafe { std::mem::transmute::<u8, IpProto>(protocol) },
source_ip,
destination_ip,
source_port,
destination_port,
len: total_len,
timestamp,
};
Some(Event::IPv4(event))
}
fn parse_ipv6(packet_data: &[u8], timestamp_us: u64) -> Option<Event> {
fn parse_ipv6(packet_data: &[u8], timestamp: u64) -> Option<Event> {
// Ethernet header (14) + minimum IPv6 header (40) = 54 bytes
if packet_data.len() < 54 {
return None;
}
let ip_header = &packet_data[14..];
let protocol_byte = ip_header[6];
let protocol = unsafe { mem::transmute::<u8, IpProto>(protocol_byte) };
// Parse IPv6 header
let protocol = ip_header[6];
// Source IPv6 address (16 bytes starting at offset 8)
let mut source_ip_bytes = [0u8; 16];
source_ip_bytes.copy_from_slice(&ip_header[8..24]);
let src_ip = u128::from_be_bytes(source_ip_bytes);
let source_ip = u128::from_be_bytes(source_ip_bytes);
// Destination IPv6 address (16 bytes starting at offset 24)
let mut dest_ip_bytes = [0u8; 16];
dest_ip_bytes.copy_from_slice(&ip_header[24..40]);
let dst_ip = u128::from_be_bytes(dest_ip_bytes);
let destination_ip = u128::from_be_bytes(dest_ip_bytes);
// Payload length
let payload_len = u16::from_be_bytes([ip_header[4], ip_header[5]]) as u32;
let total_len = payload_len + 40;
let total_len = payload_len + 40; // IPv6 header is always 40 bytes
if packet_data.len() < 54 + 4 {
return None;
}
let transport_header = &ip_header[40..];
let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]);
let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]);
let (tcp_flags, tcp_window_size, header_length) = if protocol_byte == 6 {
if packet_data.len() < 54 + 20 {
return None;
}
let data_offset = (transport_header[12] >> 4) as u16 * 4;
let flags = TcpFlags::from_byte(transport_header[13]);
let window = u16::from_be_bytes([transport_header[14], transport_header[15]]);
(flags, window, data_offset)
} else if protocol_byte == 17 {
(TcpFlags::default(), 0, 8)
// Parse transport layer (TCP/UDP)
let (source_port, destination_port) = if packet_data.len() >= 54 + 4 {
let transport_header = &ip_header[40..];
let src_port = u16::from_be_bytes([transport_header[0], transport_header[1]]);
let dst_port = u16::from_be_bytes([transport_header[2], transport_header[3]]);
(src_port, dst_port)
} else {
(TcpFlags::default(), 0, 0)
(0, 0)
};
let payload_length = total_len.saturating_sub(40 + header_length as u32);
let event = IPv6Event {
protocol,
src_ip,
dst_ip,
src_port,
dst_port,
packet_length: total_len,
payload_length,
header_length,
timestamp_us,
tcp_flags,
tcp_window_size,
is_forward: false,
protocol: unsafe { std::mem::transmute::<u8, IpProto>(protocol) },
source_ip,
destination_ip,
source_port,
destination_port,
len: total_len,
timestamp,
};
Some(Event::IPv6(event))
}
pub fn format_ipv4(addr: u32) -> String {
let bytes = addr.to_be_bytes();
format!(
"{}.{}.{}.{}",
bytes[0], bytes[1], bytes[2], bytes[3],
)
}
pub fn format_ipv6(addr: u128) -> String {
let bytes = addr.to_be_bytes();
format!(
"{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}:{:02x}{:02x}",
bytes[0], bytes[1], bytes[2], bytes[3],
bytes[4], bytes[5], bytes[6], bytes[7],
bytes[8], bytes[9], bytes[10], bytes[11],
bytes[12], bytes[13], bytes[14], bytes[15]
)
}