flowscope 0.22.0

//! CICFlowMeter-compatible feature vector subset (issue #15).

#![cfg(feature = "ml-features")]

use flowscope::ml_features::{CicFlowFeatures, count_tcp_flags};
use flowscope::{
    EndReason, FlowExtractor, FlowRecord, FlowTracker, PacketView, Timestamp,
    extract::{FiveTuple, parse::test_frames::ipv4_tcp},
};

fn end_to_end_record(payload_len: usize, end_reason: EndReason) -> FlowRecord {
    let mac = [0u8; 6];
    let ip_a = [10, 0, 0, 1];
    let ip_b = [10, 0, 0, 2];
    let mut t: FlowTracker<FiveTuple, ()> = FlowTracker::new(FiveTuple::bidirectional());
    let body: Vec<u8> = vec![b'A'; payload_len];
    let frames = [
        // SYN
        ipv4_tcp(mac, mac, ip_a, ip_b, 1234, 80, 1000, 0, 0x02, b""),
        // SYN-ACK
        ipv4_tcp(mac, mac, ip_b, ip_a, 80, 1234, 5000, 1001, 0x12, b""),
        // ACK
        ipv4_tcp(mac, mac, ip_a, ip_b, 1234, 80, 1001, 5001, 0x10, b""),
        // PSH+ACK with data
        ipv4_tcp(mac, mac, ip_a, ip_b, 1234, 80, 1001, 5001, 0x18, &body),
        // Responder ACK
        ipv4_tcp(
            mac,
            mac,
            ip_b,
            ip_a,
            80,
            1234,
            5001,
            1001 + payload_len as u32,
            0x10,
            b"",
        ),
    ];
    let mut key = None;
    for (i, f) in frames.iter().enumerate() {
        let pv = PacketView::new(f, Timestamp::new(i as u32, 0));
        if key.is_none() {
            key = Some(FiveTuple::bidirectional().extract(pv).expect("ext").key);
        }
        t.track(pv);
    }
    let key = key.unwrap();
    let stats = t.snapshot_stats(&key).expect("stats");
    FlowRecord::from_parts(&stats, &key, Some(end_reason))
}

#[test]
fn end_to_end_flow_produces_features() {
    let rec = end_to_end_record(100, EndReason::Fin);
    let feats = CicFlowFeatures::from_flow_record(&rec);
    assert!(
        feats.total_fwd_packets >= 3,
        "expected >=3 fwd packets, got {}",
        feats.total_fwd_packets
    );
    assert!(feats.total_fwd_bytes >= 100);
    assert!(feats.flow_duration_us > 0);
    assert!(feats.flow_bytes_per_sec > 0.0);
    assert_eq!(feats.conn_state, Some("SF"));
}

#[test]
fn flag_count_helper_compiles_through_public_api() {
    // Note: `FlowRecord::from_parts` does not auto-populate
    // tcp_control_bits — the caller must invoke
    // `observe_tcp_flags` per packet (typically wired from
    // the per-packet TCP-flags emit hook). Driving that
    // explicitly here keeps the test in the consumer's
    // shape.
    let mut rec = end_to_end_record(10, EndReason::Fin);
    rec.observe_tcp_flags(true, false, true, false, false, false, false, false, false);
    rec.observe_tcp_flags(false, false, true, false, false, true, false, false, false);
    let counts = count_tcp_flags(&rec);
    assert_eq!(counts.fwd_syn, 1);
    assert_eq!(counts.bwd_syn, 1);
}

#[test]
fn fin_end_reason_maps_to_sf() {
    let rec = end_to_end_record(10, EndReason::Fin);
    let feats = CicFlowFeatures::from_flow_record(&rec);
    assert_eq!(feats.conn_state, Some("SF"));
}

#[test]
fn idle_timeout_end_reason_maps_to_oth() {
    let rec = end_to_end_record(10, EndReason::IdleTimeout);
    let feats = CicFlowFeatures::from_flow_record(&rec);
    assert_eq!(feats.conn_state, Some("OTH"));
}

#[test]
fn missing_end_reason_yields_none_conn_state() {
    let mut rec = end_to_end_record(10, EndReason::Fin);
    rec.flow_end_reason = None;
    let feats = CicFlowFeatures::from_flow_record(&rec);
    assert_eq!(feats.conn_state, None);
}

#[test]
fn down_up_ratio_matches_responder_over_initiator() {
    let rec = end_to_end_record(0, EndReason::Fin);
    let feats = CicFlowFeatures::from_flow_record(&rec);
    // Both directions sent only handshake-control bytes; bwd
    // and fwd are similar but not zero. Ratio is safe (no inf).
    assert!(!feats.down_up_ratio.is_nan());
    assert!(!feats.down_up_ratio.is_infinite());
}

#[test]
fn five_tuple_app_label_populates_application_name() {
    // The 3WHS+data flow targets dport 80 → app_label "http".
    // from_parts copies that into application_name; the
    // feature vector does not surface application_name
    // directly, but verifying the FlowRecord's wiring is
    // useful sanity.
    let rec = end_to_end_record(10, EndReason::Fin);
    assert_eq!(rec.application_name.as_deref(), Some("http"));
    let _feats = CicFlowFeatures::from_flow_record(&rec);
}

#[test]
fn iat_stats_populated_by_tracker() {
    // Drive a 4-packet flow with explicit 1ms / 2ms / 3ms
    // inter-arrival gaps to verify the per-packet IAT
    // observation actually fires inside the tracker and
    // ends up in CicFlowFeatures via `with_iat`.
    use flowscope::extract::FiveTuple;
    use flowscope::extract::parse::test_frames::ipv4_tcp;
    use flowscope::{FlowExtractor, FlowTracker, PacketView};
    let mac = [0u8; 6];
    let mut t: FlowTracker<FiveTuple, ()> = FlowTracker::new(FiveTuple::bidirectional());
    let frames = [
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            1,
            0,
            0x02,
            b"",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 2],
            [10, 0, 0, 1],
            80,
            1234,
            5000,
            1,
            0x12,
            b"",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            1,
            5001,
            0x10,
            b"",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            1,
            5001,
            0x18,
            b"hi",
        ),
    ];
    // Timestamps in 1ms / 3ms / 6ms (1ms / 2ms / 3ms gaps).
    let ts = [
        Timestamp::new(0, 0),
        Timestamp::new(0, 1_000_000),
        Timestamp::new(0, 3_000_000),
        Timestamp::new(0, 6_000_000),
    ];
    let mut key_opt = None;
    for (i, f) in frames.iter().enumerate() {
        let pv = PacketView::new(f, ts[i]);
        if key_opt.is_none() {
            key_opt = Some(FiveTuple::bidirectional().extract(pv).expect("ext").key);
        }
        t.track(pv);
    }
    let key = key_opt.unwrap();
    let stats = t.snapshot_stats(&key).expect("stats");
    let rec = FlowRecord::from_parts(&stats, &key, Some(EndReason::Fin));
    let feats = CicFlowFeatures::from_flow_record(&rec).with_iat(&stats);

    // 4 packets → 3 flow-IAT observations.
    assert_eq!(stats.iat_flow.count(), 3);
    // IAT mean = (1000 + 2000 + 3000) / 3 = 2000 µs.
    assert!((feats.flow_iat_mean_us - 2000.0).abs() < 1e-9);
    assert!((feats.flow_iat_min_us - 1000.0).abs() < 1e-9);
    assert!((feats.flow_iat_max_us - 3000.0).abs() < 1e-9);
    // Forward direction: 3 packets at ts 0 / 3ms / 6ms → 2 IAT (3000, 3000).
    assert_eq!(stats.iat_initiator.count(), 2);
    assert!((feats.fwd_iat_mean_us - 3000.0).abs() < 1e-9);
    // Backward direction: 1 packet → 0 IAT (need 2+).
    assert_eq!(stats.iat_responder.count(), 0);
    assert_eq!(feats.bwd_iat_mean_us, 0.0);
}

#[test]
fn active_idle_periods_split_on_threshold_gap() {
    // Two bursts of two packets each, 2 seconds apart.
    // CICFlowMeter default threshold is 1 second, so the
    // 2-second gap should close the first active period
    // and observe one idle gap. End-of-flow doesn't close
    // the final active period (no closing IAT) — that's
    // CICFlowMeter behaviour too.
    use flowscope::extract::FiveTuple;
    use flowscope::extract::parse::test_frames::ipv4_tcp;
    use flowscope::{FlowExtractor, FlowTracker, PacketView};
    let mac = [0u8; 6];
    let mut t: FlowTracker<FiveTuple, ()> = FlowTracker::new(FiveTuple::bidirectional());
    // Two same-direction packets at ts=0/10ms (1 active period),
    // then two more at ts=2_010ms/2_020ms (next active period
    // after a 2-second gap).
    let frames = [
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            1,
            0,
            0x02,
            b"",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            2,
            0,
            0x18,
            b"a",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            3,
            0,
            0x18,
            b"b",
        ),
        ipv4_tcp(
            mac,
            mac,
            [10, 0, 0, 1],
            [10, 0, 0, 2],
            1234,
            80,
            4,
            0,
            0x18,
            b"c",
        ),
    ];
    let ts = [
        Timestamp::new(0, 0),
        Timestamp::new(0, 10_000_000),
        Timestamp::new(2, 10_000_000),
        Timestamp::new(2, 20_000_000),
    ];
    let mut key_opt = None;
    for (i, f) in frames.iter().enumerate() {
        let pv = PacketView::new(f, ts[i]);
        if key_opt.is_none() {
            key_opt = Some(FiveTuple::bidirectional().extract(pv).expect("ext").key);
        }
        t.track(pv);
    }
    let key = key_opt.unwrap();
    let stats = t.snapshot_stats(&key).expect("stats");
    let rec = FlowRecord::from_parts(&stats, &key, Some(EndReason::Fin));
    let feats = CicFlowFeatures::from_flow_record(&rec).with_iat(&stats);

    // Active periods: 1 completed (10ms first burst). Second
    // burst is still "open" at flow-end — not observed.
    assert_eq!(stats.active_periods.count(), 1);
    assert!((feats.active_mean_us - 10_000.0).abs() < 1e-6);
    // Idle periods: 1 (the 2 second gap = 2_000_000 µs).
    assert_eq!(stats.idle_periods.count(), 1);
    assert!((feats.idle_mean_us - 2_000_000.0).abs() < 1e-6);
}

#[cfg(feature = "serde")]
#[test]
fn features_serialize_to_json() {
    // CicFlowFeatures is serialize-only via serde — the
    // `conn_state: Option<&'static str>` field can be
    // emitted but not deserialized back (the borrowed `'static`
    // lifetime doesn't survive a round-trip). For producer
    // pipelines (the common ML-features case) this is the
    // right tradeoff: zero allocations on the hot path.
    let rec = end_to_end_record(50, EndReason::Fin);
    let feats = CicFlowFeatures::from_flow_record(&rec);
    let json = serde_json::to_string(&feats).expect("serialize");
    assert!(json.contains("flow_duration_us"));
    assert!(json.contains("total_fwd_packets"));
    assert!(json.contains("down_up_ratio"));
}