mirror of
https://github.com/ParrotXray/Mantis.git
synced 2026-08-24 19:50:28 +09:00
Add windowed flow export for active flows (clone + reset)
drain_ready_flows now accepts force_interval_us. Flows alive longer than force_interval_us are cloned, exported, then reset in place so the next window accumulates immediately without removing the tracker entry. Finished and idle flows continue to be removed as before. engine.rs passes flow_timeout_us / 2 (30s default) so active attack flows reach ML and traffic logger within 30 seconds instead of waiting for FIN/RST or 60s idle timeout. Training mode (traffic_logging_mode) and inference mode share the same drain path, keeping feature distributions consistent between training data and inference inputs. https://claude.ai/code/session_01UtTDLtd9MttmxtXac3nSWn
This commit is contained in:
parent
dc20f9364e
commit
15d9efa710
@ -89,7 +89,7 @@ impl Engine {
|
||||
continue;
|
||||
};
|
||||
let total_flows = t.flow_count();
|
||||
let flows = t.drain_ready_flows(self.flow_timeout_us);
|
||||
let flows = t.drain_ready_flows(self.flow_timeout_us, self.flow_timeout_us / 2);
|
||||
t.cleanup_old_flows(self.flow_timeout_us);
|
||||
// active_ips covers both the just-drained flows and flows still in the
|
||||
// tracker (ongoing connections), so their LSTM buffers are preserved.
|
||||
|
||||
@ -232,6 +232,39 @@ impl FlowData {
|
||||
pub fn packet_count(&self) -> usize {
|
||||
self.fwd_packets.len() + self.bwd_packets.len()
|
||||
}
|
||||
|
||||
// Reset per-window counters so this entry can be reused for the next
|
||||
// export window without removing it from the flow table.
|
||||
// Flow key, direction, and last_time_us are preserved.
|
||||
// init_win_bytes are cleared — they are only meaningful at connection start.
|
||||
pub fn reset(&mut self, now_us: u64) {
|
||||
self.fwd_packets.clear();
|
||||
self.bwd_packets.clear();
|
||||
self.fwd_total_bytes = 0;
|
||||
self.bwd_total_bytes = 0;
|
||||
self.fwd_header_bytes = 0;
|
||||
self.bwd_header_bytes = 0;
|
||||
self.fin_count = 0;
|
||||
self.syn_count = 0;
|
||||
self.rst_count = 0;
|
||||
self.psh_count = 0;
|
||||
self.ack_count = 0;
|
||||
self.urg_count = 0;
|
||||
self.cwe_count = 0;
|
||||
self.ece_count = 0;
|
||||
self.init_win_bytes_fwd = 0;
|
||||
self.init_win_bytes_bwd = 0;
|
||||
self.active_periods.clear();
|
||||
self.idle_periods.clear();
|
||||
self.last_packet_time = now_us;
|
||||
self.fwd_bulk_state = BulkState::default();
|
||||
self.bwd_bulk_state = BulkState::default();
|
||||
self.act_data_pkt_fwd = 0;
|
||||
self.start_time_us = now_us;
|
||||
self.active_start_us = now_us;
|
||||
self.fwd_fin_seen = false;
|
||||
self.bwd_fin_seen = false;
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FlowTracker {
|
||||
@ -313,13 +346,18 @@ impl FlowTracker {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn drain_ready_flows(&mut self, timeout_us: u64) -> Vec<FlowData> {
|
||||
// Export flows that are finished or idle.
|
||||
// Active flows alive >= force_interval_us are cloned and exported,
|
||||
// then reset in place so the next window starts accumulating immediately.
|
||||
// This ensures active attack flows reach ML/logger within force_interval_us
|
||||
// and keeps training and inference feature distributions identical.
|
||||
pub fn drain_ready_flows(&mut self, timeout_us: u64, force_interval_us: u64) -> Vec<FlowData> {
|
||||
let now = time::SystemTime::now()
|
||||
.duration_since(time::UNIX_EPOCH)
|
||||
.map(|d| d.as_micros() as u64)
|
||||
.unwrap_or(0);
|
||||
|
||||
let ready_keys: Vec<FlowKey> = self
|
||||
let terminal_keys: Vec<FlowKey> = self
|
||||
.flows
|
||||
.iter()
|
||||
.filter(|(_, flow)| {
|
||||
@ -328,10 +366,19 @@ impl FlowTracker {
|
||||
.map(|(key, _)| key.clone())
|
||||
.collect();
|
||||
|
||||
ready_keys
|
||||
let mut result: Vec<FlowData> = terminal_keys
|
||||
.into_iter()
|
||||
.filter_map(|key| self.flows.remove(&key))
|
||||
.collect()
|
||||
.collect();
|
||||
|
||||
for flow in self.flows.values_mut() {
|
||||
if now.saturating_sub(flow.start_time_us) >= force_interval_us {
|
||||
result.push(flow.clone());
|
||||
flow.reset(now);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub fn active_src_ips(&self) -> impl Iterator<Item = &str> {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user