Mantis/mantis/src/detection/ml/flow_tracker.rs

use std::collections::HashMap;
use std::time;

use common::model::event::Event;

use crate::model::direction::Direction;
use crate::model::ml_detection::{BulkState, FlowKey, PacketData};

#[derive(Debug, Clone)]
pub struct FlowData {
    pub flow_key: FlowKey,
    pub direction: Direction,
    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,
    pub act_data_pkt_fwd: u32,
    active_start_us: u64,
    fwd_fin_seen: bool,
    bwd_fin_seen: bool,
}

impl FlowData {
    pub fn new(flow_key: FlowKey, first_packet: &Event, direction: Direction) -> Self {
        Self {
            flow_key,
            direction,
            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(),
            act_data_pkt_fwd: 0,
            active_start_us: first_packet.timestamp_us(),
            fwd_fin_seen: false,
            bwd_fin_seen: false,
        }
    }

    pub fn add_packet(&mut self, packet: &Event) {
        const MAX_PACKETS_PER_DIRECTION: usize = 1000;
        const MAX_PERIODS: usize = 10000;

        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;
            // Track per-direction FIN for proper 4-way teardown detection
            if packet.is_forward() {
                self.fwd_fin_seen = true;
            } else {
                self.bwd_fin_seen = true;
            }
        }
        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 {
            // Record cumulative active period before this idle gap
            let active_dur = self.last_packet_time.saturating_sub(self.active_start_us);
            if active_dur > 0 && self.active_periods.len() < MAX_PERIODS {
                self.active_periods.push(active_dur);
            }
            if self.idle_periods.len() < MAX_PERIODS {
                self.idle_periods.push(iat);
            }
            self.active_start_us = packet.timestamp_us();
        }

        self.last_packet_time = packet.timestamp_us();
        self.last_time_us = packet.timestamp_us();

        if packet.is_forward() && packet.payload_length() > 0 {
            self.act_data_pkt_fwd += 1;
        }

        if packet.is_forward() {
            if self.fwd_packets.len() < MAX_PACKETS_PER_DIRECTION {
                self.fwd_packets.push(packet_data.clone());
            }
            self.fwd_total_bytes += packet.payload_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 {
            if self.bwd_packets.len() < MAX_PACKETS_PER_DIRECTION {
                self.bwd_packets.push(packet_data.clone());
            }
            self.bwd_total_bytes += packet.payload_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);
        }
    }

    /// TCP flow is finished when both sides have sent FIN, or either side sent RST.
    pub fn is_finished(&self) -> bool {
        (self.fwd_fin_seen && self.bwd_fin_seen) || self.rst_count > 0
    }

    fn update_bulk_state(bulk_state: &mut BulkState, packet: &PacketData) {
        const BULK_IDLE_US: u64 = 1_000_000;

        if packet.payload_length > 0 {
            // Idle break: discard helper, start fresh
            if bulk_state.in_bulk && packet.timestamp_us.saturating_sub(bulk_state.last_bulk_packet_us) > BULK_IDLE_US {
                bulk_state.in_bulk = false;
                bulk_state.last_bulk_bytes = 0;
                bulk_state.last_bulk_packets = 0;
                bulk_state.last_bulk_start_us = 0;
                bulk_state.last_bulk_packet_us = 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;
                bulk_state.last_bulk_start_us = packet.timestamp_us;
                bulk_state.last_bulk_packet_us = packet.timestamp_us;
            } else {
                let prev_us = bulk_state.last_bulk_packet_us;
                bulk_state.last_bulk_bytes += packet.length as u64;
                bulk_state.last_bulk_packets += 1;
                bulk_state.last_bulk_packet_us = packet.timestamp_us;

                if bulk_state.last_bulk_packets == 4 {
                    // Threshold first reached: commit initial 4 packets
                    bulk_state.bulk_count += 1;
                    bulk_state.total_packets += 4;
                    bulk_state.total_bytes += bulk_state.last_bulk_bytes;
                    bulk_state.total_duration_us += packet.timestamp_us.saturating_sub(bulk_state.last_bulk_start_us);
                } else if bulk_state.last_bulk_packets > 4 {
                    // Each subsequent packet adds incrementally
                    bulk_state.total_packets += 1;
                    bulk_state.total_bytes += packet.length as u64;
                    bulk_state.total_duration_us += packet.timestamp_us.saturating_sub(prev_us);
                }
            }
        } else {
            // Zero-payload: end bulk sequence (no commit — incomplete bulks discarded)
            bulk_state.in_bulk = false;
            bulk_state.last_bulk_bytes = 0;
            bulk_state.last_bulk_packets = 0;
            bulk_state.last_bulk_start_us = 0;
            bulk_state.last_bulk_packet_us = 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: HashMap<FlowKey, FlowData>,
    max_flows: usize,
}

impl FlowTracker {
    pub fn new(max_flows: usize) -> Self {
        Self {
            flows: HashMap::new(),
            max_flows,
        }
    }

    pub fn process_packet(&mut self, mut packet: Event, is_ingress: bool, payload: &[u8]) {
        let direction = if is_ingress {
            Direction::Ingress
        } else {
            Direction::Egress
        };

        let packet_key = FlowKey::from_packet(&packet);
        let reversed_key = packet_key.clone().reverse();

        // Try-both: canonical key is whichever orientation already exists in the flow table.
        // For new flows, identify the initiator using (in priority order):
        //   1. TCP SYN / SYN+ACK flags
        //   2. DPI: TLS ClientHello/ServerHello, HTTP request/response, DNS QR bit
        //   3. Best effort: use packet as-is
        let (actual_key, is_forward) = if self.flows.contains_key(&packet_key) {
            (packet_key, true)
        } else if self.flows.contains_key(&reversed_key) {
            (reversed_key, false)
        } else {
            let flags = packet.tcp_flags();
            if flags.syn && flags.ack {
                // Normal: Server (egress side) sends SYN+ACK, packet arrives on ingress → reverse
                // Bot attack: Client (egress side) sends SYN+ACK, packet arrives on egress → keep as-is
                if is_ingress {
                    (reversed_key, false)
                } else {
                    (packet_key, true)
                }
            } else if flags.syn {
                (packet_key, true)
            } else {
                match detect_initiator(payload, packet_key.protocol, packet_key.src_port, packet_key.dst_port) {
                    Some(true) => (packet_key, true),
                    Some(false) => (reversed_key, false),
                    None => (packet_key, true),
                }
            }
        };

        packet.set_is_forward(is_forward);

        let initiator_direction = if is_forward { direction } else { direction.flip() };

        let flow = self
            .flows
            .entry(actual_key.clone())
            .or_insert_with(|| FlowData::new(actual_key, &packet, initiator_direction));

        flow.add_packet(&packet);

        if self.flows.len() > self.max_flows {
            if let Some(key) = self.flows.keys().next().cloned() {
                self.flows.remove(&key);
            }
        }
    }

    pub fn drain_flows(&mut self) -> Vec<FlowData> {
        self.flows.drain().map(|(_, flow)| flow).collect()
    }

    pub fn get_flows_snapshot(&self) -> Vec<FlowData> {
        self.flows.values().cloned().collect()
    }

    pub fn cleanup_old_flows(&mut 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);
        self.flows
            .retain(|_, flow| !flow.is_finished() && now.saturating_sub(flow.last_time_us) < max_age_us);
    }

    pub fn flow_count(&self) -> usize {
        self.flows.len()
    }
}

fn detect_initiator(payload: &[u8], protocol: u8, src_port: u16, dst_port: u16) -> Option<bool> {
    if payload.is_empty() {
        return None;
    }

    // TLS: record type 0x16 (Handshake), byte 5 = handshake type
    //   0x01 = ClientHello → this side is the initiator
    //   0x02 = ServerHello → this side is the responder
    if payload.len() >= 6 && payload[0] == 0x16 {
        return match payload[5] {
            0x01 => Some(true),
            0x02 => Some(false),
            _ => None,
        };
    }

    // HTTP: request line starts with a method verb (initiator),
    //       response starts with "HTTP/" (responder)
    if payload.len() >= 5 {
        if payload.starts_with(b"GET ")
            || payload.starts_with(b"POST ")
            || payload.starts_with(b"PUT ")
            || payload.starts_with(b"HEAD ")
            || payload.starts_with(b"DELETE ")
            || payload.starts_with(b"OPTIONS ")
            || payload.starts_with(b"PATCH ")
        {
            return Some(true);
        }
        if payload.starts_with(b"HTTP/") {
            return Some(false);
        }
    }

    // DNS over UDP (port 53): flags byte 2, MSB = QR bit
    //   0 = query (initiator), 1 = response (responder)
    if protocol == 17 && (src_port == 53 || dst_port == 53) && payload.len() >= 3 {
        return Some((payload[2] >> 7) == 0);
    }

    None
}