PM-pinou
366 lines
16 KiB
Python
366 lines
16 KiB
Python
"""End-to-end integration test for the full TLS analysis pipeline.
|
||
|
||
Pipeline
|
||
1. Create synthetic TLS flow CSV with known entity clusters.
|
||
2. Load via ``load_csv_directory``.
|
||
3. Store in ``SessionStore``.
|
||
4. Auto-detect entity column (``detect_entity_column``).
|
||
5. Aggregate flows → entity profiles (``aggregate_by_entity``).
|
||
6. Run HDBSCAN clustering on aggregated numeric features.
|
||
7. Extract per-cluster distinguishing features.
|
||
|
||
All data is synthetic, deterministic (seed=42), and self-contained —
|
||
no external files or network access required.
|
||
"""
|
||
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
import pytest
|
||
|
||
from analysis.data_loader import load_csv_directory
|
||
from analysis.entity_detector import detect_entity_column
|
||
from analysis.entity_aggregator import aggregate_by_entity
|
||
from analysis.session_store import SessionStore
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Fixtures
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def clean_store():
|
||
"""Reset the SessionStore singleton before and after each test."""
|
||
store = SessionStore()
|
||
store.drop_all()
|
||
yield
|
||
store.drop_all()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Synthetic data generator
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
def create_synthetic_flows(seed: int = 42) -> pl.DataFrame:
|
||
"""Generate deterministic TLS flow data with three known entity groups.
|
||
|
||
Uses the user's exact column naming convention only.
|
||
|
||
Returns
|
||
-------
|
||
pl.DataFrame
|
||
~1000 rows with columns.
|
||
Entities (``:ips``) are partitioned into three behavioural
|
||
groups that HDBSCAN should recover after aggregation.
|
||
"""
|
||
rng = np.random.default_rng(seed)
|
||
|
||
# ── Entity profiles ──────────────────────────────────────────────
|
||
# Group A: 10 high-volume entities (50 flows each)
|
||
group_a = [
|
||
{
|
||
':ips': f'10.0.0.{i+1}',
|
||
'n_flows': 50,
|
||
'8ack_mu': 50_000, '8ack_sigma': 5_000,
|
||
'4dur_mu': 30.0, '4dur_sigma': 5.0,
|
||
'dst_range': (100, 150),
|
||
'0ver': 'TLSv1.3',
|
||
'cipher_suite': 'TLS_AES_256_GCM_SHA384',
|
||
}
|
||
for i in range(10)
|
||
]
|
||
|
||
# Group B: 15 medium-volume entities (20 flows each)
|
||
group_b = [
|
||
{
|
||
':ips': f'10.0.0.{i+11}',
|
||
'n_flows': 20,
|
||
'8ack_mu': 5_000, '8ack_sigma': 1_000,
|
||
'4dur_mu': 15.0, '4dur_sigma': 3.0,
|
||
'dst_range': (150, 200),
|
||
'0ver': 'TLSv1.2',
|
||
'cipher_suite': 'TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256',
|
||
}
|
||
for i in range(15)
|
||
]
|
||
|
||
# Group C: 25 low-volume entities (6 flows each)
|
||
group_c = [
|
||
{
|
||
':ips': f'10.0.0.{i+26}',
|
||
'n_flows': 6,
|
||
'8ack_mu': 500, '8ack_sigma': 100,
|
||
'4dur_mu': 5.0, '4dur_sigma': 1.0,
|
||
'dst_range': (200, 250),
|
||
'0ver': 'TLSv1.2',
|
||
'cipher_suite': 'TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256',
|
||
}
|
||
for i in range(25)
|
||
]
|
||
|
||
entities = group_a + group_b + group_c
|
||
# Total flows: 10×50 + 15×20 + 25×6 = 500 + 300 + 150 = 950
|
||
|
||
# ── Generate flow rows ────────────────────────────────────────────
|
||
rows: list[dict] = []
|
||
dst_ports = [80, 443, 8080, 8443]
|
||
protos = ['TCP', 'UDP']
|
||
proto_weights = [0.9, 0.1]
|
||
|
||
for ent in entities:
|
||
lo, hi = ent['dst_range']
|
||
for _ in range(ent['n_flows']):
|
||
dst_idx = int(rng.integers(lo, hi))
|
||
row = {
|
||
':ips': ent[':ips'],
|
||
':ipd': f"10.0.{dst_idx >> 8}.{dst_idx & 0xff}",
|
||
':prd': int(rng.choice(dst_ports)),
|
||
'8ack': int(max(0, rng.normal(ent['8ack_mu'],
|
||
ent['8ack_sigma']))),
|
||
'4dur': round(max(0.1, abs(rng.normal(ent['4dur_mu'],
|
||
ent['4dur_sigma']))), 2),
|
||
'1ipp': str(rng.choice(protos, p=proto_weights)),
|
||
'0ver': ent['0ver'],
|
||
'cipher_suite': ent['cipher_suite'],
|
||
'snam': f"host{rng.integers(1, 200)}.example.com",
|
||
'timestamp': (
|
||
f"2024-06-15 "
|
||
f"{rng.integers(0, 23):02d}:"
|
||
f"{rng.integers(0, 59):02d}:"
|
||
f"{rng.integers(0, 59):02d}"
|
||
),
|
||
}
|
||
rows.append(row)
|
||
|
||
return pl.DataFrame(rows)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Pipeline tests
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
class TestFullPipeline:
|
||
"""End-to-end: synthetic CSV → load → entity detect → aggregate → cluster."""
|
||
|
||
def _write_csv(self, df: pl.DataFrame, tmpdir: str) -> str:
|
||
"""Write a Polars DataFrame to a temp CSV; return the file path."""
|
||
path = Path(tmpdir) / 'synthetic_flows.csv'
|
||
df.write_csv(path)
|
||
return str(path)
|
||
|
||
def test_full_pipeline(self):
|
||
"""Smoke-test the full pipeline — every step produces valid output."""
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
# ── Step 1: Create synthetic data ─────────────────────────────
|
||
df_flows = create_synthetic_flows(seed=42)
|
||
assert len(df_flows) > 0, 'Synthetic data should not be empty'
|
||
|
||
csv_path = self._write_csv(df_flows, tmpdir)
|
||
|
||
# ── Step 2: Load data ─────────────────────────────────────────
|
||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||
csv_path, encoding='utf-8',
|
||
)
|
||
assert isinstance(lf, pl.LazyFrame), 'Expected LazyFrame'
|
||
assert file_count == 1
|
||
assert row_count > 0
|
||
assert memory_mb >= 0
|
||
assert ':ips' in schema
|
||
|
||
# ── Step 3: Store in session ──────────────────────────────────
|
||
store = SessionStore()
|
||
store.store_dataset(
|
||
'e2e_ds', lf, schema=schema,
|
||
metadata={'row_count': row_count, 'csv_glob': csv_path},
|
||
)
|
||
|
||
# ── Step 4: Entity detection ──────────────────────────────────
|
||
det_result = detect_entity_column('e2e_ds')
|
||
assert 'error' not in det_result, f'detect_entity_column failed: {det_result.get("error")}'
|
||
# :ips is the catch-all entity column (strings with good unique ratio)
|
||
assert det_result['recommended'] in (':ips', ':ipd', 'snam'), (
|
||
f'Expected :ips, :ipd or snam, got {det_result["recommended"]}'
|
||
)
|
||
assert len(det_result['candidates']) > 0
|
||
|
||
# Verify recommended_multi exists
|
||
assert 'recommended_multi' in det_result, 'recommended_multi field missing'
|
||
assert len(det_result['recommended_multi']) > 0, 'recommended_multi is empty'
|
||
|
||
# ── Step 5: Entity aggregation ────────────────────────────────
|
||
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
||
df_agg = agg_lf.collect()
|
||
|
||
# Should be exactly 50 entities (10 + 15 + 25)
|
||
entity_count = len(df_agg)
|
||
assert entity_count == 50, f'Expected 50 entities, got {entity_count}'
|
||
|
||
# Verify per-entity flow counts (sort by flow_count desc so groups
|
||
# are contiguous: A=50, B=20, C=6)
|
||
df_sorted = df_agg.sort('flow_count', descending=True)
|
||
# Group A: 10 entities × 50 flows
|
||
for i in range(10):
|
||
assert df_sorted['flow_count'][i] == 50, (
|
||
f'Row {i}: expected 50, got {df_sorted["flow_count"][i]}'
|
||
)
|
||
# Group B: 15 entities × 20 flows
|
||
for i in range(10, 25):
|
||
assert df_sorted['flow_count'][i] == 20, (
|
||
f'Row {i}: expected 20, got {df_sorted["flow_count"][i]}'
|
||
)
|
||
# Group C: 25 entities × 6 flows
|
||
for i in range(25, 50):
|
||
assert df_sorted['flow_count'][i] == 6, (
|
||
f'Row {i}: expected 6, got {df_sorted["flow_count"][i]}'
|
||
)
|
||
|
||
# Core feature columns are present
|
||
for col in ('flow_count', 'total_bytes_sent',
|
||
'avg_duration', 'unique_dst_ips', 'unique_dst_ports',
|
||
'unique_tls_versions', 'unique_protocols',
|
||
'first_seen', 'last_seen'):
|
||
assert col in feature_cols, f'Missing feature column: {col}'
|
||
|
||
# ── Step 6: Clustering ────────────────────────────────────────
|
||
numeric_cols = [
|
||
'flow_count', 'total_bytes_sent',
|
||
'avg_duration', 'unique_dst_ips', 'unique_dst_ports',
|
||
]
|
||
available_cols = [c for c in numeric_cols if c in df_agg.columns]
|
||
|
||
data = df_agg.select(available_cols).to_numpy()
|
||
assert data.shape[0] == 50, 'All 50 entities should have features'
|
||
|
||
# Standard-scale + HDBSCAN
|
||
from sklearn.preprocessing import StandardScaler
|
||
from sklearn.cluster import HDBSCAN
|
||
|
||
data_scaled = StandardScaler().fit_transform(data)
|
||
clusterer = HDBSCAN(min_cluster_size=5, metric='euclidean')
|
||
labels = clusterer.fit_predict(data_scaled)
|
||
|
||
n_clusters = len(set(labels) - {-1})
|
||
assert n_clusters > 0, (
|
||
f'HDBSCAN should find ≥1 cluster, found {n_clusters} '
|
||
f'(labels={set(labels)})'
|
||
)
|
||
|
||
noise_count = int((labels == -1).sum())
|
||
# At most 40% noise allowed (we have well-separated groups)
|
||
noise_ratio = noise_count / len(labels)
|
||
assert noise_ratio < 0.40, (
|
||
f'Noise ratio too high: {noise_ratio:.2%} '
|
||
f'({noise_count}/{len(labels)})'
|
||
)
|
||
|
||
# ── Step 7: Feature extraction ────────────────────────────────
|
||
unique_labels = sorted(set(l for l in labels if l >= 0))
|
||
assert len(unique_labels) > 0, 'Need at least one valid cluster'
|
||
|
||
# Compute per-cluster distinguishing features using z-score
|
||
features = []
|
||
for label in unique_labels:
|
||
mask = labels == label
|
||
cluster_data = data[mask]
|
||
for col_idx, col_name in enumerate(available_cols):
|
||
global_mean = float(data[:, col_idx].mean())
|
||
global_std = float(data[:, col_idx].std()) or 1.0
|
||
cluster_mean = float(cluster_data[:, col_idx].mean())
|
||
z = (cluster_mean - global_mean) / global_std
|
||
features.append({
|
||
'cluster': int(label),
|
||
'feature': col_name,
|
||
'z_score': round(z, 4),
|
||
})
|
||
|
||
assert len(features) > 0, 'Expected at least one extracted feature'
|
||
# Each cluster should have features for all numeric columns
|
||
assert len(features) >= len(unique_labels) * len(available_cols)
|
||
|
||
# ── Step 8: Subnet aggregation function ────────────────────────
|
||
from analysis.entity_aggregator import ip_to_subnet
|
||
assert callable(ip_to_subnet)
|
||
assert ip_to_subnet('10.0.1.15', 24) == '10.0.1.0/24'
|
||
|
||
# Note: SessionStore is cleaned up by the ``clean_store`` fixture
|
||
|
||
def test_synthetic_data_shape(self):
|
||
"""Verify the synthetic data generator produces expected shapes."""
|
||
df = create_synthetic_flows(seed=42)
|
||
assert len(df.columns) == 10, f'Expected 10 columns, got {len(df.columns)}'
|
||
# Total: 10×50 + 15×20 + 25×6 = 950
|
||
assert len(df) == 950, f'Expected 950 rows, got {len(df)}'
|
||
|
||
# Verify unique entity count
|
||
n_entities = df[':ips'].n_unique()
|
||
assert n_entities == 50, f'Expected 50 unique :ips, got {n_entities}'
|
||
|
||
# Group A should have TLSv1.3
|
||
group_a_ips = [f'10.0.0.{i+1}' for i in range(10)]
|
||
group_a_tls = df.filter(
|
||
pl.col(':ips').is_in(group_a_ips)
|
||
)['0ver'].unique().to_list()
|
||
assert group_a_tls == ['TLSv1.3']
|
||
|
||
# Group C should have CHACHA20 cipher
|
||
group_c_ips = [f'10.0.0.{i+26}' for i in range(25)]
|
||
group_c_cipher = df.filter(
|
||
pl.col(':ips').is_in(group_c_ips)
|
||
)['cipher_suite'].unique().to_list()
|
||
assert any('CHACHA20' in c for c in group_c_cipher)
|
||
|
||
|
||
class TestPipelineEdgeCases:
|
||
"""Edge cases for the analysis pipeline."""
|
||
|
||
def test_empty_dataset_rejected(self):
|
||
"""Empty dataset stored → detect_entity_column returns error."""
|
||
store = SessionStore()
|
||
empty = pl.DataFrame({':ips': pl.Series([], dtype=pl.Utf8)})
|
||
schema = {}
|
||
store.store_dataset(
|
||
'empty_ds', empty.lazy(), schema=schema,
|
||
metadata={'row_count': 0},
|
||
)
|
||
result = detect_entity_column('empty_ds')
|
||
# May fail or return empty candidates; either is acceptable
|
||
# but it should not crash with an unhandled exception.
|
||
assert 'error' in result or result['recommended'] == ''
|
||
|
||
def test_single_entity_clustering(self):
|
||
"""Single entity → aggregation works but clustering may not split."""
|
||
data = pl.DataFrame({
|
||
':ips': ['10.0.0.1'] * 20,
|
||
':ipd': ['10.0.0.2'] * 20,
|
||
'8ack': [100] * 20,
|
||
'4dur': [1.0] * 20,
|
||
'1ipp': ['TCP'] * 20,
|
||
'0ver': ['TLSv1.2'] * 20,
|
||
'cipher_suite': ['AES'] * 20,
|
||
'snam': ['host1.example.com'] * 20,
|
||
'timestamp': ['2024-01-01 00:00:00'] * 20,
|
||
':prd': [443] * 20,
|
||
})
|
||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||
lf = data.lazy()
|
||
agg_lf, _ = aggregate_by_entity(lf, ':ips', schema)
|
||
df_agg = agg_lf.collect()
|
||
|
||
assert len(df_agg) == 1
|
||
assert df_agg['flow_count'][0] == 20
|
||
|
||
def test_aggregate_preserves_schema_types(self):
|
||
"""Aggregated columns should have expected Polars dtypes."""
|
||
data = pl.DataFrame({
|
||
':ips': ['10.0.0.1'] * 10,
|
||
'8ack': [100] * 10,
|
||
'4dur': [1.5] * 10,
|
||
})
|
||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||
lf = data.lazy()
|
||
agg_lf, _ = aggregate_by_entity(lf, ':ips', schema)
|
||
df = agg_lf.collect()
|
||
|
||
assert df['flow_count'].dtype == pl.UInt32 or df['flow_count'].dtype == pl.Int64
|