Initial commit: code, paper, small artifacts

tests/common/test_data_contract.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+import sys
+from pathlib import Path
+import numpy as np
+import pytest
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+from common.data_contract import BENIGN_TOKEN, CANONICAL_FLOW_FEATURE_NAMES, FLOW_COUNT_FEATURE_NAMES, FLOW_COUNT_IDX, FLOW_CONTINUOUS_IDX, FLOW_D, IDLE_THRESHOLD_MS, PACKET_BINARY_CHANNEL_IDX, PACKET_CONTINUOUS_CHANNEL_IDX, PACKET_D, PACKET_FEATURE_NAMES, UNKNOWN_LABEL_TOKEN, apply_mixed_dequant, canonical_5tuple, compute_flow_features_from_packets, fit_packet_stats, normalize_label, zscore
+
+def test_packet_schema_invariants():
+    assert PACKET_D == 9
+    assert len(PACKET_FEATURE_NAMES) == PACKET_D
+    assert len(set(PACKET_FEATURE_NAMES)) == PACKET_D
+    cont = set(PACKET_CONTINUOUS_CHANNEL_IDX)
+    binary = set(PACKET_BINARY_CHANNEL_IDX)
+    assert cont.isdisjoint(binary)
+    assert cont | binary == set(range(PACKET_D))
+
+def test_flow_schema_invariants():
+    assert FLOW_D == 20
+    assert len(CANONICAL_FLOW_FEATURE_NAMES) == FLOW_D
+    assert len(set(CANONICAL_FLOW_FEATURE_NAMES)) == FLOW_D
+    assert set(FLOW_COUNT_IDX).isdisjoint(FLOW_CONTINUOUS_IDX)
+    assert set(FLOW_COUNT_IDX) | set(FLOW_CONTINUOUS_IDX) == set(range(FLOW_D))
+    for name in FLOW_COUNT_FEATURE_NAMES:
+        assert name in CANONICAL_FLOW_FEATURE_NAMES
+
+@pytest.mark.parametrize('raw,expected', [('BENIGN', BENIGN_TOKEN), ('Benign', BENIGN_TOKEN), ('benign', BENIGN_TOKEN), ('normal', BENIGN_TOKEN), ('  BENIGN  ', BENIGN_TOKEN), ('DrDoS_DNS', 'DrDoS_DNS'), ('', UNKNOWN_LABEL_TOKEN), ('   ', UNKNOWN_LABEL_TOKEN)])
+def test_normalize_label(raw, expected):
+    assert normalize_label(raw) == expected
+
+def test_canonical_5tuple_direction_agnostic():
+    a = canonical_5tuple('10.0.0.1', 1234, '10.0.0.2', 80, 6)
+    b = canonical_5tuple('10.0.0.2', 80, '10.0.0.1', 1234, 6)
+    assert a == b
+
+def test_canonical_5tuple_distinct_flows():
+    a = canonical_5tuple('10.0.0.1', 1234, '10.0.0.2', 80, 6)
+    c = canonical_5tuple('10.0.0.1', 1234, '10.0.0.3', 80, 6)
+    assert a != c
+
+def test_canonical_5tuple_string_inputs():
+    a = canonical_5tuple('10.0.0.1', '1234', '10.0.0.2', '80', '6')
+    b = canonical_5tuple('10.0.0.2', 80, '10.0.0.1', 1234, 6)
+    assert a == b
+
+def _make_packets(n_flows: int, T: int, n_real: int, seed: int=0) -> tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    tokens = np.zeros((n_flows, T, PACKET_D), dtype=np.float32)
+    tokens[:, :n_real, 0] = rng.uniform(0, 5, (n_flows, n_real))
+    tokens[:, :n_real, 1] = rng.uniform(0, 3, (n_flows, n_real))
+    tokens[:, :n_real, 2] = rng.integers(0, 2, (n_flows, n_real))
+    tokens[:, :n_real, 3] = rng.integers(0, 2, (n_flows, n_real))
+    tokens[:, :n_real, 7] = rng.integers(0, 2, (n_flows, n_real))
+    tokens[:, :n_real, 8] = rng.uniform(0, 4, (n_flows, n_real))
+    lens = np.full(n_flows, n_real, dtype=np.int32)
+    return (tokens, lens)
+
+def test_fit_packet_stats_ignores_padding():
+    (tokens, lens) = _make_packets(8, 16, 4, seed=1)
+    tokens[:, 4:, :] = 999.0
+    (mean, std) = fit_packet_stats(tokens, lens)
+    assert np.all(np.abs(mean) < 50)
+    assert np.all(std < 50)
+
+def test_zscore_basic():
+    x = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+    mean = np.array([2.0, 3.0], dtype=np.float32)
+    std = np.array([1.0, 1.0], dtype=np.float32)
+    z = zscore(x, mean, std)
+    assert np.allclose(z, [[-1.0, -1.0], [1.0, 1.0]])
+
+def test_zscore_handles_zero_std():
+    x = np.ones((3, 2), dtype=np.float32)
+    mean = np.ones(2, dtype=np.float32)
+    std = np.zeros(2, dtype=np.float32)
+    z = zscore(x, mean, std)
+    assert np.all(np.isfinite(z))
+
+def test_apply_mixed_dequant_zeros_padding():
+    (tokens, lens) = _make_packets(4, 8, 3, seed=2)
+    (mean, std) = fit_packet_stats(tokens, lens)
+    z = apply_mixed_dequant(tokens, lens, mean, std, split_tag='train', seed=0)
+    assert np.all(z[:, 3:, :] == 0.0)
+
+def test_apply_mixed_dequant_stable_under_same_seed():
+    (tokens, lens) = _make_packets(4, 8, 3, seed=3)
+    (mean, std) = fit_packet_stats(tokens, lens)
+    z1 = apply_mixed_dequant(tokens, lens, mean, std, split_tag='train', seed=42)
+    z2 = apply_mixed_dequant(tokens, lens, mean, std, split_tag='train', seed=42)
+    assert np.allclose(z1, z2)
+
+def test_apply_mixed_dequant_split_noise_differs():
+    (tokens, lens) = _make_packets(4, 8, 3, seed=4)
+    (mean, std) = fit_packet_stats(tokens, lens)
+    z_train = apply_mixed_dequant(tokens, lens, mean, std, split_tag='train', seed=42)
+    z_val = apply_mixed_dequant(tokens, lens, mean, std, split_tag='val', seed=42)
+    b_idx = list(PACKET_BINARY_CHANNEL_IDX)
+    assert not np.allclose(z_train[..., b_idx], z_val[..., b_idx])
+
+def test_flow_features_shape():
+    (tokens, lens) = _make_packets(5, 16, 8, seed=5)
+    feats = compute_flow_features_from_packets(tokens, lens)
+    assert feats.shape == (5, FLOW_D)
+    assert feats.dtype == np.float32
+
+def test_flow_features_rejects_wrong_shape():
+    tokens_bad = np.zeros((3, 4, 7), dtype=np.float32)
+    lens = np.full(3, 2, dtype=np.int32)
+    with pytest.raises(ValueError):
+        compute_flow_features_from_packets(tokens_bad, lens)
+
+def test_flow_features_zero_for_empty_flow():
+    tokens = np.zeros((2, 4, PACKET_D), dtype=np.float32)
+    lens = np.array([0, 3], dtype=np.int32)
+    tokens[1, :3, 0] = 2.0
+    feats = compute_flow_features_from_packets(tokens, lens)
+    assert np.all(feats[0] == 0.0)
+    assert np.any(feats[1] != 0.0)
+
+def test_flow_features_ack_syn_counts():
+    tokens = np.zeros((1, 5, PACKET_D), dtype=np.float32)
+    tokens[0, :3, 7] = 1.0
+    tokens[0, 0, 3] = 1.0
+    tokens[0, 3:, 7] = 1.0
+    tokens[0, 3:, 3] = 1.0
+    lens = np.array([3], dtype=np.int32)
+    feats = compute_flow_features_from_packets(tokens, lens)
+    ack_idx = CANONICAL_FLOW_FEATURE_NAMES.index('ack_cnt')
+    syn_idx = CANONICAL_FLOW_FEATURE_NAMES.index('syn_cnt')
+    assert feats[0, ack_idx] == pytest.approx(3.0)
+    assert feats[0, syn_idx] == pytest.approx(1.0)
+
+def test_flow_features_fwd_bwd_counts():
+    tokens = np.zeros((1, 6, PACKET_D), dtype=np.float32)
+    tokens[0, :4, 2] = np.array([0, 1, 0, 1])
+    lens = np.array([4], dtype=np.int32)
+    feats = compute_flow_features_from_packets(tokens, lens)
+    fwd_idx = CANONICAL_FLOW_FEATURE_NAMES.index('fwd_count')
+    bwd_idx = CANONICAL_FLOW_FEATURE_NAMES.index('bwd_count')
+    assert feats[0, fwd_idx] == 2.0
+    assert feats[0, bwd_idx] == 2.0
+
+def test_flow_features_active_vs_idle():
+    tokens = np.zeros((1, 4, PACKET_D), dtype=np.float32)
+    lens = np.array([4], dtype=np.int32)
+    tokens[0, 0, 1] = 0.0
+    tokens[0, 1, 1] = np.log1p(100.0)
+    tokens[0, 2, 1] = np.log1p(5000.0)
+    tokens[0, 3, 1] = np.log1p(200.0)
+    feats = compute_flow_features_from_packets(tokens, lens, idle_threshold_ms=1000.0)
+    active_mean_idx = CANONICAL_FLOW_FEATURE_NAMES.index('active_mean')
+    idle_mean_idx = CANONICAL_FLOW_FEATURE_NAMES.index('idle_mean')
+    assert feats[0, active_mean_idx] == pytest.approx(np.log1p(150.0), rel=0.0001)
+    assert feats[0, idle_mean_idx] == pytest.approx(np.log1p(5000.0), rel=0.0001)
+
+def test_flow_features_padding_does_not_leak():
+    T = 12
+    (tokens_clean, lens) = _make_packets(3, T, 5, seed=6)
+    feats_clean = compute_flow_features_from_packets(tokens_clean, lens)
+    tokens_poisoned = tokens_clean.copy()
+    tokens_poisoned[:, 5:, :] = 9999.0
+    feats_poisoned = compute_flow_features_from_packets(tokens_poisoned, lens)
+    assert np.allclose(feats_clean, feats_poisoned, atol=0.0001)
+
+def test_flow_features_single_packet_graceful():
+    tokens = np.zeros((1, 4, PACKET_D), dtype=np.float32)
+    tokens[0, 0, 0] = 3.0
+    tokens[0, 0, 2] = 0
+    lens = np.array([1], dtype=np.int32)
+    feats = compute_flow_features_from_packets(tokens, lens)
+    assert np.all(np.isfinite(feats))
+    iat_idx = CANONICAL_FLOW_FEATURE_NAMES.index('iat_mean')
+    assert feats[0, iat_idx] == 0.0