PM-pinou
323 lines
13 KiB
Python
323 lines
13 KiB
Python
"""Tests for entity detection (entity_detector) and aggregation (entity_aggregator).
|
|
|
|
Every test isolates the SessionStore singleton via the ``clean_store`` fixture.
|
|
"""
|
|
|
|
import pytest
|
|
import polars as pl
|
|
|
|
from analysis.session_store import SessionStore
|
|
from analysis.entity_detector import _score_name, detect_entity_column
|
|
from analysis.entity_aggregator import (
|
|
_find_column,
|
|
aggregate_by_entity,
|
|
build_aggregation_params,
|
|
ip_to_subnet,
|
|
)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# 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()
|
|
|
|
|
|
def _store_dataset(
|
|
df: pl.DataFrame,
|
|
dataset_id: str = 'test_ds',
|
|
) -> str:
|
|
"""Helper: store a DataFrame as a dataset in the session store.
|
|
|
|
Returns the dataset_id.
|
|
"""
|
|
store = SessionStore()
|
|
schema = {name: str(dtype) for name, dtype in zip(df.columns, df.dtypes)}
|
|
store.store_dataset(
|
|
dataset_id=dataset_id,
|
|
lazyframe=df.lazy(),
|
|
schema=schema,
|
|
metadata={'row_count': len(df)},
|
|
)
|
|
return dataset_id
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# _score_name (unit-level)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestScoreName:
|
|
"""_score_name() scores column names against entity keywords."""
|
|
|
|
def test_exact_match(self):
|
|
"""Exact keyword match → 3.0."""
|
|
score, kw = _score_name(':ips')
|
|
assert score == 3.0
|
|
assert kw == ':ips'
|
|
|
|
def test_suffix_match(self):
|
|
"""Suffix match via _ separator → 3.0."""
|
|
score, kw = _score_name('my_snam')
|
|
assert score == 3.0
|
|
assert kw == 'snam'
|
|
|
|
def test_prefix_match(self):
|
|
"""Prefix match via _ separator → 3.0."""
|
|
score, kw = _score_name('snam_value')
|
|
assert score == 3.0
|
|
assert kw == 'snam'
|
|
|
|
def test_substring_not_matched(self):
|
|
"""Substring-only match now returns 0.0 (exact matching only)."""
|
|
score, kw = _score_name('myipaddr')
|
|
assert score == 0.0
|
|
assert kw == ''
|
|
|
|
def test_no_match(self):
|
|
score, kw = _score_name('random_column_xyz')
|
|
assert score == 0.0
|
|
assert kw == ''
|
|
|
|
def test_hyphen_normalised(self):
|
|
"""Hyphen normalised to underscore for matching."""
|
|
score, kw = _score_name('cipher-suite')
|
|
assert score == 3.0
|
|
assert kw == 'cipher_suite'
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# detect_entity_column (integration-level)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestDetectEntityColumn:
|
|
"""detect_entity_column() should find entity columns from stored data."""
|
|
|
|
def test_detect_ips(self):
|
|
"""Dataset with ``:ips`` → recommended is ``:ips``."""
|
|
data = pl.DataFrame({
|
|
':ips': [f'10.0.0.{i}' for i in range(1, 51)] * 10,
|
|
':ipd': [f'10.0.1.{i}' for i in range(1, 501)],
|
|
'8ack': [100] * 500,
|
|
})
|
|
ds_id = _store_dataset(data)
|
|
result = detect_entity_column(ds_id)
|
|
|
|
assert result['recommended'] == ':ips'
|
|
assert len(result['candidates']) > 0
|
|
assert result['candidates'][0]['column'] == ':ips'
|
|
assert result['total_rows'] == 500
|
|
assert result['scored_columns'] >= 1
|
|
|
|
def test_detect_snam(self):
|
|
"""Datasets with ``snam`` but no IP column → snam is recommended."""
|
|
data = pl.DataFrame({
|
|
'snam': [f'host{i}.example.com' for i in range(100)],
|
|
'4dur': [float(i) for i in range(100)],
|
|
})
|
|
ds_id = _store_dataset(data)
|
|
result = detect_entity_column(ds_id)
|
|
|
|
assert result['recommended'] == 'snam'
|
|
assert result['candidates'][0]['matched_keyword'] == 'snam'
|
|
|
|
def test_no_entity_column(self):
|
|
"""No entity keywords matched → candidates still found but matched_keyword empty."""
|
|
data = pl.DataFrame({
|
|
'col_a': [str(i) for i in range(100)],
|
|
'col_b': [str(i * 2) for i in range(100)],
|
|
'value_f64': [float(i) for i in range(100)],
|
|
})
|
|
ds_id = _store_dataset(data)
|
|
result = detect_entity_column(ds_id)
|
|
|
|
# String columns are candidates (detector scores on uniqueness, not just keywords)
|
|
assert len(result['candidates']) >= 2
|
|
assert all(c['matched_keyword'] == '' for c in result['candidates'])
|
|
# recommended will be the highest-scoring string column
|
|
assert result['recommended'] in ('col_a', 'col_b')
|
|
|
|
def test_all_numeric_no_strings(self):
|
|
"""No string/categorical columns → no candidates."""
|
|
data = pl.DataFrame({
|
|
'value': list(range(100)),
|
|
'flag': [0, 1] * 50,
|
|
})
|
|
ds_id = _store_dataset(data)
|
|
result = detect_entity_column(ds_id)
|
|
|
|
assert result['recommended'] == ''
|
|
assert len(result['candidates']) == 0
|
|
|
|
def test_dataset_not_found(self):
|
|
"""Unknown dataset_id → error response."""
|
|
result = detect_entity_column('nonexistent_ds')
|
|
assert 'error' in result
|
|
assert result['recommended'] == ''
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# _find_column (helper for aggregation)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestFindColumn:
|
|
"""_find_column() matches keywords against schema column names."""
|
|
|
|
ACK_KEYWORDS = ('8ack',)
|
|
|
|
def test_exact_match(self):
|
|
schema = {'8ack': 'Int64', 'other': 'Utf8'}
|
|
col = _find_column(schema, self.ACK_KEYWORDS)
|
|
assert col == '8ack'
|
|
|
|
def test_no_match(self):
|
|
schema = {'something_else': 'Int64'}
|
|
col = _find_column(schema, self.ACK_KEYWORDS)
|
|
assert col is None
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# aggregate_by_entity
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestAggregateByEntity:
|
|
"""aggregate_by_entity() groups flows and computes per-entity features."""
|
|
|
|
def test_two_entities(self):
|
|
"""50 flows across 2 entities → 2 aggregated rows with correct counts."""
|
|
data = pl.DataFrame({
|
|
':ips': ['10.0.0.1'] * 20 + ['10.0.0.2'] * 30,
|
|
':ipd': ['10.0.1.1'] * 10 + ['10.0.1.2'] * 10
|
|
+ ['10.0.1.3'] * 30,
|
|
':prd': [80] * 20 + [443] * 30,
|
|
'8ack': [100] * 20 + [200] * 30,
|
|
'4dur': [1.0] * 20 + [2.0] * 30,
|
|
'0ver': ['TLSv1.2'] * 20 + ['TLSv1.3'] * 30,
|
|
'cipher_suite': ['AES'] * 20 + ['CHACHA20'] * 30,
|
|
'1ipp': ['TCP'] * 50,
|
|
'timestamp': ['2024-01-01 00:00:00'] * 50,
|
|
})
|
|
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
|
lf = data.lazy()
|
|
|
|
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
|
df = agg_lf.collect()
|
|
|
|
# Two entities
|
|
assert len(df) == 2
|
|
|
|
# Flow counts
|
|
df_sorted = df.sort(':ips')
|
|
assert df_sorted['flow_count'][0] == 20 # 10.0.0.1
|
|
assert df_sorted['flow_count'][1] == 30 # 10.0.0.2
|
|
|
|
# Bytes — each column is now perfectly aligned with entity groups
|
|
assert df_sorted['total_bytes_sent'][0] == 20 * 100
|
|
assert df_sorted['total_bytes_sent'][1] == 30 * 200
|
|
|
|
# Core feature columns present
|
|
assert 'flow_count' in feature_cols
|
|
assert 'total_bytes_sent' in feature_cols
|
|
assert 'avg_duration' in feature_cols
|
|
assert 'unique_dst_ips' in feature_cols
|
|
assert 'unique_dst_ports' in feature_cols
|
|
assert 'unique_protocols' in feature_cols
|
|
assert 'first_seen' in feature_cols
|
|
assert 'last_seen' in feature_cols
|
|
|
|
def test_single_entity(self):
|
|
"""Single entity → single aggregated row."""
|
|
data = pl.DataFrame({
|
|
':ips': ['10.0.0.1'] * 5,
|
|
'8ack': [100] * 5,
|
|
':ipd': ['10.0.0.2'] * 5,
|
|
})
|
|
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 len(df) == 1
|
|
assert df['flow_count'][0] == 5
|
|
|
|
def test_nonexistent_entity_column(self):
|
|
"""Unknown entity column → ValueError."""
|
|
data = pl.DataFrame({'a': [1, 2, 3]})
|
|
schema = {'a': 'Int64'}
|
|
lf = data.lazy()
|
|
with pytest.raises(ValueError, match=r"Entity column\(s\) '\[\'xyz\'\]'"):
|
|
aggregate_by_entity(lf, 'xyz', schema)
|
|
|
|
def test_schema_with_missing_columns(self):
|
|
"""Aggregation still works when some optional columns are missing."""
|
|
data = pl.DataFrame({
|
|
':ips': ['10.0.0.1'] * 3 + ['10.0.0.2'] * 3,
|
|
':ipd': ['10.0.0.3'] * 6,
|
|
# No 8ack, 4dur, etc.
|
|
})
|
|
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
|
lf = data.lazy()
|
|
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
|
|
|
df = agg_lf.collect()
|
|
assert len(df) == 2
|
|
assert df['flow_count'][0] == 3
|
|
# Missing features should have None values
|
|
assert df['total_bytes_sent'][0] is None
|
|
assert 'flow_count' in feature_cols
|
|
assert 'unique_dst_ips' in feature_cols
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# build_aggregation_params
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestBuildAggregationParams:
|
|
"""build_aggregation_params() builds correct expressions per schema."""
|
|
|
|
def test_flow_count_always_present(self):
|
|
params = build_aggregation_params({'a': 'Int64'})
|
|
names = [p[0] for p in params]
|
|
assert 'flow_count' in names
|
|
|
|
def test_ack_detected(self):
|
|
params = build_aggregation_params({'8ack': 'Int64'})
|
|
names = [p[0] for p in params]
|
|
assert 'total_bytes_sent' in names
|
|
|
|
def test_ack_not_detected(self):
|
|
params = build_aggregation_params({'x': 'Int64'})
|
|
param_map = dict(params)
|
|
assert param_map['total_bytes_sent'] is not None # still has a fallback expr
|
|
# The value should be a Polars literal (None)
|
|
names = [p[0] for p in params]
|
|
assert 'total_bytes_sent' in names
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# ip_to_subnet (IP subnet mask conversion)
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class TestIpToSubnet:
|
|
"""ip_to_subnet() converts IP strings to subnet prefixes."""
|
|
|
|
def test_ip_to_subnet_24(self):
|
|
"""ip_to_subnet with /24 mask."""
|
|
assert ip_to_subnet('10.0.1.15', 24) == '10.0.1.0/24'
|
|
assert ip_to_subnet('192.168.1.255', 24) == '192.168.1.0/24'
|
|
|
|
def test_ip_to_subnet_28(self):
|
|
"""ip_to_subnet with /28 mask."""
|
|
assert ip_to_subnet('10.0.1.15', 28) == '10.0.1.0/28'
|
|
assert ip_to_subnet('10.0.1.250', 28) == '10.0.1.240/28'
|
|
|
|
def test_ip_to_subnet_invalid(self):
|
|
"""ip_to_subnet with invalid inputs."""
|
|
assert ip_to_subnet(None, 24) is None
|
|
assert ip_to_subnet('', 24) == ''
|
|
assert ip_to_subnet('not_an_ip', 24) == 'not_an_ip'
|