rustflow_export 0.10.0

IPFIX exporter for captured network traffic (Linux only); the `rustflow export` command
Documentation
use std::collections::HashMap;
use std::net::Ipv4Addr;

use chrono::{DateTime, TimeDelta, Utc};

use crate::ipfix::data::FlowData;

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct FlowKey {
    pub source_ip: Ipv4Addr,
    pub destination_ip: Ipv4Addr,
    pub protocol: u8,
    pub source_port: u16,
    pub destination_port: u16,
}

#[derive(Debug, Clone)]
pub struct Flow {
    pub key: FlowKey,
    pub octet_count: u64,
    pub packet_count: u64,
    pub tcp_flags: u16,
    pub flow_start: DateTime<Utc>,
    pub flow_end: DateTime<Utc>,
}

impl Flow {
    pub fn new(key: FlowKey) -> Self {
        let now = Utc::now();
        Self {
            key,
            octet_count: 0,
            packet_count: 0,
            tcp_flags: 0,
            flow_start: now,
            flow_end: now,
        }
    }

    pub fn update(&mut self, packet_size: u64, tcp_flags: u16) {
        self.octet_count += packet_size;
        self.packet_count += 1;
        self.tcp_flags |= tcp_flags;
        self.flow_end = Utc::now();
    }

    pub fn to_flow_data(&self) -> FlowData {
        FlowData {
            source_ipv4: self.key.source_ip,
            destination_ipv4: self.key.destination_ip,
            protocol: self.key.protocol,
            source_port: self.key.source_port,
            destination_port: self.key.destination_port,
            octet_count: self.octet_count,
            packet_count: self.packet_count,
            tcp_flags: self.tcp_flags,
            flow_start: self.flow_start,
            flow_end: self.flow_end,
        }
    }
}

pub struct FlowCache {
    flows: HashMap<FlowKey, Flow>,
    active_timeout: TimeDelta,
    inactive_timeout: TimeDelta,
}

impl FlowCache {
    pub fn new(active_timeout: u64, inactive_timeout: u64) -> Self {
        Self {
            flows: HashMap::new(),
            active_timeout: TimeDelta::seconds(active_timeout as i64),
            inactive_timeout: TimeDelta::seconds(inactive_timeout as i64),
        }
    }

    pub fn update_flow(&mut self, key: FlowKey, packet_size: u64, tcp_flags: u16) {
        self.flows
            .entry(key.clone())
            .or_insert_with(|| Flow::new(key))
            .update(packet_size, tcp_flags);
    }

    pub fn check_expired_flows(&mut self) -> Vec<Flow> {
        let now = Utc::now();
        let mut expired = Vec::new();

        let active_timeout = self.active_timeout;
        let inactive_timeout = self.inactive_timeout;

        self.flows.retain(|_, flow| {
            let age = now - flow.flow_start;
            let idle_time = now - flow.flow_end;

            if age >= active_timeout || idle_time >= inactive_timeout {
                expired.push(flow.clone());
                false
            } else {
                true
            }
        });

        expired
    }

    pub fn export_all(&mut self) -> Vec<Flow> {
        let flows: Vec<Flow> = self.flows.values().cloned().collect();
        self.flows.clear();
        flows
    }

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