PM-pinou
588 lines
25 KiB
Python
588 lines
25 KiB
Python
"""Aggregate flow-level data into entity-level profiles.
|
||
|
||
The main entry point is :func:`aggregate_by_entity`, which groups a Polars
|
||
:class:`LazyFrame` by one or more entity columns (e.g. ``src_ip``, ``sni``,
|
||
``user``) and computes per-entity aggregates across flow, destination, TLS,
|
||
protocol, and time dimensions.
|
||
|
||
Only features whose source columns exist in the provided schema are computed,
|
||
so the module works with any CSV schema without assumptions.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from typing import Optional
|
||
|
||
import polars as pl
|
||
import numpy as np
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
|
||
from .type_classifier import DataType
|
||
|
||
_agg_logger = logging.getLogger(__name__)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Keyword sets for schema-independent column discovery
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Flow stats — user exact column names only
|
||
_BYTES_SENT_KEYWORDS = ('8ack', '8byt',)
|
||
_BYTES_REV_KEYWORDS = () # no user column for reverse bytes
|
||
_DURATION_KEYWORDS = ('4dur',)
|
||
_PACKETS_KEYWORDS = ('8ppk', '8pak',)
|
||
_PACKETS_REV_KEYWORDS = () # no user column for reverse packets
|
||
|
||
# Destination diversity — user exact column names only
|
||
_DST_IP_KEYWORDS = (':ips', ':ipd')
|
||
_DST_PORT_KEYWORDS = (':prs', ':prd')
|
||
_DST_NET_KEYWORDS = () # no user column for destination network
|
||
|
||
# TLS features — user exact column names only
|
||
_CIPHER_KEYWORDS = ('cipher_suite', '0cph')
|
||
_TLS_VERSION_KEYWORDS = ('0ver',)
|
||
_SNI_KEYWORDS = ('snam', 'cnam')
|
||
|
||
# Protocol features — user exact column names only
|
||
_PROTO_KEYWORDS = ('1ipp',)
|
||
_TCP_FLAGS_KEYWORDS = () # no user column for TCP flags
|
||
|
||
# Geo-location features — user exact column names only
|
||
_LAT_KEYWORDS = (':ips.latd', ':ipd.latd', 'latd')
|
||
_LON_KEYWORDS = (':ips.lond', ':ipd.lond', 'lond')
|
||
|
||
# Time patterns — user exact column names only
|
||
_TIMESTAMP_KEYWORDS = ('time', 'timestamp', '8dbd')
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# IP subnet helpers for prefix-based aggregation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def ip_to_subnet(ip_str: str, mask: int) -> str:
|
||
"""Convert ``10.0.1.15/24`` → ``10.0.1.0/24``, ``10.0.1.250/28`` →
|
||
``10.0.1.240/28``.
|
||
|
||
Invalid IPs / ``None`` / empty string are returned unchanged.
|
||
Supports ``/24`` (0–255 last octet) and ``/28`` (240–255 in steps of 16)
|
||
masks.
|
||
"""
|
||
if not ip_str or not isinstance(ip_str, str):
|
||
return ip_str
|
||
try:
|
||
parts = [int(o) for o in ip_str.strip().split('.')]
|
||
if len(parts) != 4 or any(p < 0 or p > 255 for p in parts):
|
||
return ip_str
|
||
ip_int = (parts[0] << 24) + (parts[1] << 16) + (parts[2] << 8) + parts[3]
|
||
subnet_int = ip_int & (0xFFFFFFFF << (32 - mask))
|
||
return (
|
||
f"{(subnet_int >> 24) & 255}."
|
||
f"{(subnet_int >> 16) & 255}."
|
||
f"{(subnet_int >> 8) & 255}."
|
||
f"{subnet_int & 255}/{mask}"
|
||
)
|
||
except (ValueError, IndexError, TypeError):
|
||
return ip_str
|
||
|
||
|
||
def apply_subnet_aggregation(
|
||
lf: pl.LazyFrame,
|
||
ip_columns: list[str],
|
||
subnet_masks: list[int],
|
||
) -> pl.LazyFrame:
|
||
"""Apply subnet masking to IP columns in the LazyFrame.
|
||
|
||
Each IP column is transformed to its subnet prefix for each mask.
|
||
If multiple masks exist, the first matching mask is applied
|
||
(earliest mask in *subnet_masks* wins).
|
||
"""
|
||
if not subnet_masks or not ip_columns:
|
||
return lf
|
||
schema = lf.collect_schema().names()
|
||
for mask in subnet_masks:
|
||
mask_exprs = []
|
||
for col in ip_columns:
|
||
if col in schema:
|
||
mask_exprs.append(
|
||
pl.col(col)
|
||
.map_batches(
|
||
lambda s, m=mask: s.cast(pl.Utf8).map_batches(
|
||
lambda x, mm=m: pl.Series(
|
||
[ip_to_subnet(v, mm) for v in x.to_list()]
|
||
),
|
||
return_dtype=pl.Utf8,
|
||
),
|
||
return_dtype=pl.Utf8,
|
||
)
|
||
.alias(col)
|
||
)
|
||
if mask_exprs:
|
||
lf = lf.with_columns(mask_exprs)
|
||
return lf
|
||
|
||
|
||
def _find_column(schema: dict[str, str], keywords: tuple[str, ...]) -> Optional[str]:
|
||
"""Return the first column in *schema* whose lowercased name matches
|
||
any of *keywords* (exact or ``_``-delimited match)."""
|
||
for col in schema:
|
||
normalised = col.lower().lstrip(':+').replace('-', '_').replace('.', '_').replace(' ', '_')
|
||
for kw in sorted(keywords):
|
||
kw_s = kw.lower().lstrip(':+')
|
||
if kw_s == normalised or normalised.endswith(f'_{kw_s}') or normalised.startswith(f'{kw_s}_'):
|
||
return col
|
||
return None
|
||
|
||
|
||
def _timestamp_col(schema: dict[str, str]) -> Optional[str]:
|
||
"""Detect a timestamp column -- user exact column names only."""
|
||
for preference in ('timestamp', 'time', '8dbd'):
|
||
if preference in schema:
|
||
return preference
|
||
return _find_column(schema, _TIMESTAMP_KEYWORDS)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Aggregation plan builder
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def build_aggregation_params(
|
||
schema: dict[str, str],
|
||
type_map: dict[str, DataType] | None = None,
|
||
) -> list[tuple[str, pl.Expr]]:
|
||
"""Build a list of ``(feature_name, pl.Expr)`` pairs for every aggregate
|
||
that can be computed from the columns present in *schema*.
|
||
|
||
Each expression is a column-level aggregation suitable for use inside
|
||
``pl.LazyFrame.group_by(...).agg(...)``.
|
||
"""
|
||
aggs: list[tuple[str, pl.Expr]] = []
|
||
|
||
# ── Flow stats ───────────────────────────────────────────────────────
|
||
# flow_count: always computed via count()
|
||
aggs.append(('flow_count', pl.len()))
|
||
|
||
bytes_sent_col = _find_column(schema, _BYTES_SENT_KEYWORDS)
|
||
if bytes_sent_col:
|
||
aggs.append(('total_bytes_sent', pl.sum(bytes_sent_col).cast(pl.Int64)))
|
||
# Also store raw column name for downstream use
|
||
aggs.append(('_bytes_sent_col', pl.lit(bytes_sent_col)))
|
||
aggs.append(('_bytes_sent_expr', pl.lit(bytes_sent_col)))
|
||
else:
|
||
aggs.append(('total_bytes_sent', pl.lit(None, dtype=pl.Int64)))
|
||
aggs.append(('_bytes_sent_col', pl.lit('')))
|
||
aggs.append(('_bytes_sent_expr', pl.lit('')))
|
||
_agg_logger.info(f'[FIND_COL] keywords=bytes_sent matched={bytes_sent_col} dtype={schema.get(bytes_sent_col, "N/A") if bytes_sent_col else "None"}')
|
||
|
||
packets_col = _find_column(schema, _PACKETS_KEYWORDS)
|
||
if packets_col:
|
||
aggs.append(('total_packets', pl.sum(packets_col).cast(pl.Int64)))
|
||
else:
|
||
aggs.append(('total_packets', pl.lit(None, dtype=pl.Int64)))
|
||
_agg_logger.info(f'[FIND_COL] keywords=packets matched={packets_col} dtype={schema.get(packets_col, "N/A") if packets_col else "None"}')
|
||
|
||
# ── Destination diversity ────────────────────────────────────────────
|
||
dst_ip_col = _find_column(schema, _DST_IP_KEYWORDS)
|
||
if dst_ip_col:
|
||
aggs.append(('unique_dst_ips', pl.n_unique(dst_ip_col)))
|
||
aggs.append(('_dst_ip_col', pl.lit(dst_ip_col)))
|
||
else:
|
||
aggs.append(('unique_dst_ips', pl.lit(None, dtype=pl.UInt32)))
|
||
aggs.append(('_dst_ip_col', pl.lit('')))
|
||
_agg_logger.info(f'[FIND_COL] keywords=dst_ip matched={dst_ip_col} dtype={schema.get(dst_ip_col, "N/A") if dst_ip_col else "None"}')
|
||
|
||
dst_port_col = _find_column(schema, _DST_PORT_KEYWORDS)
|
||
if dst_port_col:
|
||
aggs.append(('unique_dst_ports', pl.n_unique(dst_port_col)))
|
||
aggs.append(('_dst_port_col', pl.lit(dst_port_col)))
|
||
else:
|
||
aggs.append(('unique_dst_ports', pl.lit(None, dtype=pl.UInt32)))
|
||
aggs.append(('_dst_port_col', pl.lit('')))
|
||
_agg_logger.info(f'[FIND_COL] keywords=dst_port matched={dst_port_col} dtype={schema.get(dst_port_col, "N/A") if dst_port_col else "None"}')
|
||
|
||
# ── Protocol features ────────────────────────────────────────────────
|
||
proto_col = _find_column(schema, _PROTO_KEYWORDS)
|
||
if proto_col:
|
||
aggs.append(('unique_protocols', pl.n_unique(proto_col)))
|
||
aggs.append(('_proto_col', pl.lit(proto_col)))
|
||
else:
|
||
aggs.append(('unique_protocols', pl.lit(None, dtype=pl.UInt32)))
|
||
aggs.append(('_proto_col', pl.lit('')))
|
||
_agg_logger.info(f'[FIND_COL] keywords=proto matched={proto_col} dtype={schema.get(proto_col, "N/A") if proto_col else "None"}')
|
||
|
||
# ── Enhanced TLS analysis features ────────────────────────────────────
|
||
# has_sni: count of non-null SNI values
|
||
sni_col = _find_column(schema, _SNI_KEYWORDS)
|
||
if sni_col:
|
||
aggs.append(('has_sni', pl.col(sni_col).is_not_null().sum()))
|
||
aggs.append(('sni_missing_ratio', pl.col(sni_col).is_null().mean()))
|
||
aggs.append(('avg_sni_length', pl.col(sni_col).cast(pl.Utf8).str.len_bytes().mean()))
|
||
aggs.append(('max_sni_length', pl.col(sni_col).cast(pl.Utf8).str.len_bytes().max()))
|
||
|
||
# TLS version risk
|
||
tls_version_col = _find_column(schema, ('0ver',))
|
||
if tls_version_col:
|
||
aggs.append(('tls_modern_ratio',
|
||
pl.when(pl.col(tls_version_col).is_in(['tls_1_3'])).then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('tls_old_ratio',
|
||
pl.when(pl.col(tls_version_col).is_in(['tls_1_0', 'tls_1_1', 'ssl_2_0', 'ssl_3_0'])).then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('unique_tls_versions', pl.n_unique(tls_version_col)))
|
||
|
||
# Modern cipher ratio (AEAD = modern)
|
||
cipher_col = _find_column(schema, ('0cph', 'cipher-suite'))
|
||
if cipher_col:
|
||
aggs.append(('aead_ratio',
|
||
pl.when(pl.col(cipher_col).str.contains('gcm|chacha20|poly1305')).then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('weak_cipher_ratio',
|
||
pl.when(pl.col(cipher_col).str.contains('cbc|rc4|null')).then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('ecdhe_ratio',
|
||
pl.when(pl.col(cipher_col).str.contains('ecdhe')).then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('unique_ciphers', pl.n_unique(cipher_col)))
|
||
|
||
# Non-standard port detection (:prd is often String type)
|
||
prd_col = _find_column(schema, (':prd',))
|
||
if prd_col:
|
||
prd_str = pl.col(prd_col).cast(pl.Utf8)
|
||
aggs.append(('non_standard_port_ratio',
|
||
pl.when(prd_str != '443').then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
aggs.append(('port_443_ratio',
|
||
pl.when(prd_str == '443').then(1)
|
||
.otherwise(0).cast(pl.Float64).mean()))
|
||
|
||
# Bytes per packet analysis
|
||
bytes_col = _find_column(schema, ('8ack',))
|
||
pkts_col = _find_column(schema, ('8ppk',))
|
||
if bytes_col and pkts_col:
|
||
bpp = pl.col(bytes_col).cast(pl.Float64) / pl.col(pkts_col).cast(pl.Float64).clip(1)
|
||
aggs.append(('avg_bytes_per_packet', bpp.mean()))
|
||
aggs.append(('std_bytes_per_packet', bpp.std()))
|
||
aggs.append(('min_bytes_per_packet', bpp.min()))
|
||
aggs.append(('max_bytes_per_packet', bpp.max()))
|
||
|
||
# Duration analysis
|
||
dur_col = _find_column(schema, ('4dur',))
|
||
if dur_col:
|
||
aggs.append(('avg_duration', pl.col(dur_col).cast(pl.Float64).mean()))
|
||
aggs.append(('std_duration', pl.col(dur_col).cast(pl.Float64).std()))
|
||
aggs.append(('max_duration', pl.col(dur_col).cast(pl.Float64).max()))
|
||
|
||
# Recoverable/resumed ratio (columns already converted to boolean)
|
||
cnrs_col = _find_column(schema, ('cnrs',))
|
||
if cnrs_col:
|
||
aggs.append(('recoverable_ratio', pl.col(cnrs_col).cast(pl.Float64).mean()))
|
||
|
||
isrs_col = _find_column(schema, ('isrs',))
|
||
if isrs_col:
|
||
aggs.append(('resumed_ratio', pl.col(isrs_col).cast(pl.Float64).mean()))
|
||
|
||
# Country diversity
|
||
dcnt_col = _find_column(schema, ('dcnt',))
|
||
if dcnt_col:
|
||
aggs.append(('countries_visited', pl.n_unique(dcnt_col)))
|
||
|
||
# ── Geo-location (latitude / longitude) ────────────────────────────
|
||
lat_col = _find_column(schema, _LAT_KEYWORDS)
|
||
if lat_col:
|
||
dtype_str = schema.get(lat_col, '').lower()
|
||
if 'str' in dtype_str or 'utf' in dtype_str:
|
||
lat_expr = pl.mean(pl.col(lat_col).cast(pl.Float64))
|
||
else:
|
||
lat_expr = pl.mean(lat_col)
|
||
aggs.append(('avg_latitude', lat_expr.cast(pl.Float64)))
|
||
else:
|
||
aggs.append(('avg_latitude', pl.lit(None, dtype=pl.Float64)))
|
||
_agg_logger.info(f'[FIND_COL] keywords=lat matched={lat_col} dtype={schema.get(lat_col, "N/A") if lat_col else "None"}')
|
||
|
||
lon_col = _find_column(schema, _LON_KEYWORDS)
|
||
if lon_col:
|
||
dtype_str = schema.get(lon_col, '').lower()
|
||
if 'str' in dtype_str or 'utf' in dtype_str:
|
||
lon_expr = pl.mean(pl.col(lon_col).cast(pl.Float64))
|
||
else:
|
||
lon_expr = pl.mean(lon_col)
|
||
aggs.append(('avg_longitude', lon_expr.cast(pl.Float64)))
|
||
else:
|
||
aggs.append(('avg_longitude', pl.lit(None, dtype=pl.Float64)))
|
||
_agg_logger.info(f'[FIND_COL] keywords=lon matched={lon_col} dtype={schema.get(lon_col, "N/A") if lon_col else "None"}')
|
||
|
||
# ── Time patterns ────────────────────────────────────────────────────
|
||
ts_col = _timestamp_col(schema)
|
||
if ts_col:
|
||
# For first/last seen we need min/max of the timestamp column
|
||
aggs.append(('first_seen', pl.min(ts_col)))
|
||
aggs.append(('last_seen', pl.max(ts_col)))
|
||
aggs.append(('_ts_col', pl.lit(ts_col)))
|
||
# Active hours: unique hour-of-day count
|
||
# We'll compute this after collecting by extracting hour from timestamps
|
||
aggs.append(('_has_timestamp', pl.lit(True)))
|
||
else:
|
||
aggs.append(('first_seen', pl.lit(None, dtype=pl.Utf8)))
|
||
aggs.append(('last_seen', pl.lit(None, dtype=pl.Utf8)))
|
||
aggs.append(('_ts_col', pl.lit('')))
|
||
aggs.append(('_has_timestamp', pl.lit(False)))
|
||
|
||
# ── Type-classifier features ──────────────────────────────────────
|
||
# HEX/URL/IPv4 feature extraction removed — avg_hex_pairs_count,
|
||
# unique_domains, and unique_24_networks are not useful for TLS analysis.
|
||
|
||
return aggs
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main aggregation entry point
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Columns whose ``_``-prefixed names are internal bookkeeping, not user-facing
|
||
_INTERNAL_COLUMNS = tuple([
|
||
'_bytes_sent_col', '_dst_ip_col', '_dst_port_col',
|
||
'_proto_col', '_ts_col', '_has_timestamp',
|
||
'_bytes_sent_expr',
|
||
])
|
||
|
||
|
||
def aggregate_by_entity(
|
||
lf: pl.LazyFrame,
|
||
entity_columns: str | list[str],
|
||
schema: dict[str, str],
|
||
type_map: dict[str, DataType] | None = None,
|
||
) -> tuple[pl.LazyFrame, list[str]]:
|
||
"""Group flow-level data by *entity_columns* and compute per-entity features.
|
||
|
||
Parameters
|
||
----------
|
||
lf:
|
||
A Polars :class:`LazyFrame` containing flow-level data.
|
||
entity_columns:
|
||
Column name(s) to group by (e.g. ``'src_ip'``, or ``['src_ip', 'dst_ip']``).
|
||
schema:
|
||
Dictionary mapping column names to dtype strings (e.g. from
|
||
``SessionStore.get_dataset(...)['schema']``).
|
||
|
||
Returns
|
||
-------
|
||
tuple[pl.LazyFrame, list[str]]
|
||
``(aggregated_lazyframe, feature_column_names)`` where the lazyframe
|
||
has one row per entity and *feature_column_names* lists the
|
||
user-facing aggregate columns.
|
||
"""
|
||
# Normalise to list
|
||
if isinstance(entity_columns, str):
|
||
entity_columns = [entity_columns]
|
||
|
||
# ── Boolean indicator column conversion ──────────────────────────────
|
||
# cnrs/isrs have values "+" (True) or blank (False). Override schema
|
||
# type and convert to boolean before aggregation.
|
||
for col in ('cnrs', 'isrs'):
|
||
if col in schema:
|
||
schema[col] = 'Boolean'
|
||
lf = lf.with_columns(
|
||
pl.when(pl.col(col).cast(pl.Utf8).str.contains(r'\+|True|yes', literal=False))
|
||
.then(pl.lit(True))
|
||
.otherwise(pl.lit(False))
|
||
.alias(col)
|
||
)
|
||
|
||
# ── Proxy / Anomaly Detection Scores ──────────────────────────────
|
||
# 1. Non-standard port ratio (ports other than 443, 80, 53, 8080, 8443)
|
||
port_col = _find_column(schema, (':prd', 'dst_port', '8prd', 'port'))
|
||
if port_col:
|
||
std_ports = [443, 80, 53, 8080, 8443]
|
||
lf = lf.with_columns(
|
||
(~pl.col(port_col).cast(pl.Utf8).str.replace(r'\.0$', '').cast(pl.Int64, strict=False)
|
||
.is_in(std_ports)).cast(pl.Float64)
|
||
.alias('_non_std_port_ratio')
|
||
)
|
||
|
||
# 2. Modern TLS version ratio (TLS 1.2 / 1.3 = 0303 / 0304)
|
||
tls_col = _find_column(schema, ('0ver', 'tls_version', '_tlsver'))
|
||
if tls_col:
|
||
lf = lf.with_columns(
|
||
pl.col(tls_col).cast(pl.Utf8).str.replace(' ', '')
|
||
.is_in(['0303', '0304']).cast(pl.Float64)
|
||
.alias('_modern_tls_ratio')
|
||
)
|
||
|
||
# 3. SNI missing ratio (null or empty string)
|
||
sni_col = _find_column(schema, _SNI_KEYWORDS)
|
||
if sni_col:
|
||
lf = lf.with_columns(
|
||
pl.col(sni_col).cast(pl.Utf8)
|
||
.map_batches(lambda s: (s == '') | s.is_null(), return_dtype=pl.Boolean)
|
||
.cast(pl.Float64)
|
||
.alias('_sni_missing_ratio')
|
||
)
|
||
|
||
# Build aggregation expressions based on the schema (and optional type_map)
|
||
agg_params = build_aggregation_params(schema, type_map=type_map)
|
||
|
||
# Separate feature names and expressions
|
||
feature_names: list[str] = []
|
||
exprs: list[pl.Expr] = []
|
||
for name, expr in agg_params:
|
||
exprs.append(expr.alias(name))
|
||
if not name.startswith('_'):
|
||
feature_names.append(name)
|
||
|
||
# Add mean aggregation for scoring pre-computed columns
|
||
_lf_cols = lf.collect_schema().names()
|
||
for col_name, alias_name in [
|
||
('_non_std_port_ratio', 'non_std_port_rate'),
|
||
('_modern_tls_ratio', 'modern_tls_rate'),
|
||
]:
|
||
if col_name in _lf_cols:
|
||
exprs.append(pl.col(col_name).mean().alias(alias_name))
|
||
feature_names.append(alias_name)
|
||
|
||
# Check entity columns exist
|
||
missing = [c for c in entity_columns if c not in schema]
|
||
if missing:
|
||
raise ValueError(
|
||
f"Entity column(s) '{missing}' not found in schema. "
|
||
f"Available columns: {list(schema.keys())}"
|
||
)
|
||
|
||
# Perform aggregation
|
||
aggregated_lf = lf.group_by(entity_columns).agg(exprs)
|
||
|
||
# ── Adaptive anomaly scoring (data-driven, no hardcoded weights) ──
|
||
aggregated_lf = _add_adaptive_scores(aggregated_lf, feature_names)
|
||
_agg_cols = aggregated_lf.collect_schema().names()
|
||
if 'proxy_score' in _agg_cols:
|
||
feature_names.append('proxy_score')
|
||
if 'threat_score' in _agg_cols:
|
||
feature_names.append('threat_score')
|
||
if 'risk_level' in _agg_cols:
|
||
feature_names.append('risk_level')
|
||
|
||
# Post-process: compute active_hours_count if a timestamp column exists
|
||
# (We compute this as a second pass since it requires extraction logic)
|
||
any(name == '_has_timestamp' for name, _ in agg_params)
|
||
|
||
# ── Post-aggregation: log number of entities for adaptive sampling info ──
|
||
try:
|
||
n_entities = aggregated_lf.select(pl.len()).collect(streaming=True).item()
|
||
except Exception:
|
||
n_entities = 0
|
||
if n_entities > 50000:
|
||
logger.info('[SCORE] %s entities — adaptive scoring will use sampled PCA weights', n_entities)
|
||
|
||
# Filter to only user-facing feature columns names (excluding _internal)
|
||
user_features = [f for f in feature_names if f not in _INTERNAL_COLUMNS]
|
||
|
||
return aggregated_lf, user_features
|
||
|
||
|
||
# ── Adaptive anomaly scoring (data-driven, no hardcoded weights) ──────
|
||
|
||
_ANOMALY_FEATURE_COLS = [
|
||
'non_std_port_rate', 'modern_tls_rate', 'sni_missing_ratio',
|
||
'recoverable_ratio', 'flow_count', 'total_bytes_sent',
|
||
'total_packets', 'unique_dst_ips', 'unique_dst_ports',
|
||
'unique_protocols', 'has_sni', 'avg_sni_length',
|
||
'countries_visited',
|
||
]
|
||
|
||
|
||
def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
|
||
"""Add anomaly scores using data-driven weights (PCA + IQR + percentile).
|
||
|
||
Steps:
|
||
1. Find which anomaly feature columns exist in the aggregated data
|
||
2. Collect a sample (max 10K rows) for PCA weight learning
|
||
3. Learn feature weights from PCA first component loadings
|
||
4. Compute weighted proxy_score across ALL rows (vectorized, no per-row ops)
|
||
5. Add percentile rank calibration
|
||
6. Infer risk thresholds from data distribution
|
||
"""
|
||
_agg_cols = lf.collect_schema().names()
|
||
avail = [c for c in _ANOMALY_FEATURE_COLS if c in _agg_cols]
|
||
if not avail:
|
||
return lf
|
||
|
||
try:
|
||
# Step 1-2: Sample up to 10K rows for weight learning
|
||
n_rows = lf.select(pl.len()).collect(streaming=True).item()
|
||
sample_lf = lf if n_rows <= 10000 else lf.sample(n=10000, seed=42)
|
||
df_sample = sample_lf.select(avail).collect(streaming=True)
|
||
mat = df_sample.to_numpy()
|
||
mat = np.nan_to_num(mat, nan=0.0)
|
||
|
||
# Step 3: PCA weight learning
|
||
from sklearn.decomposition import PCA
|
||
n_comp = min(1, min(mat.shape) - 1)
|
||
if n_comp < 1 or mat.shape[1] < 2:
|
||
return lf
|
||
pca = PCA(n_components=1, random_state=42).fit(mat)
|
||
raw_weights = np.abs(pca.components_[0])
|
||
# Normalise to sum = 1
|
||
weights = dict(zip(avail, raw_weights / raw_weights.sum()))
|
||
|
||
# Step 4: Weighted sum expression (all vectorised, no loop)
|
||
expr_parts = []
|
||
for col, w in weights.items():
|
||
if col == 'modern_tls_rate':
|
||
# modern_tls_rate is good — invert for anomaly scoring
|
||
expr_parts.append((1.0 - pl.col(col).fill_null(0.5)) * w)
|
||
else:
|
||
expr_parts.append(pl.col(col).fill_null(0.0) * w)
|
||
if expr_parts:
|
||
proxy_expr = expr_parts[0]
|
||
for p in expr_parts[1:]:
|
||
proxy_expr = proxy_expr + p
|
||
else:
|
||
return lf
|
||
|
||
lf = lf.with_columns(proxy_expr.clip(0, 1).alias('_raw_score'))
|
||
|
||
# Step 5: Percentile rank calibration (0-1, relative position)
|
||
# Polars doesn't have native rank%, so we approximate with
|
||
# quantile-based bucketing on the collected sample thresholds
|
||
sample_scores = df_sample.select(
|
||
pl.sum_horizontal([
|
||
(1.0 - pl.col('modern_tls_rate').fill_null(0.5)) * weights.get('modern_tls_rate', 0)
|
||
if c == 'modern_tls_rate'
|
||
else pl.col(c).fill_null(0.0) * weights.get(c, 0)
|
||
for c in avail
|
||
]).clip(0, 1)
|
||
).to_series().to_numpy()
|
||
|
||
# Compute IQR thresholds from sample
|
||
q25, q50, q75 = np.percentile(sample_scores, [25, 50, 75])
|
||
iqr = q75 - q25
|
||
upper_fence = float(q75 + 1.5 * iqr) if iqr > 0 else 0.9
|
||
extreme_fence = float(q75 + 3.0 * iqr) if iqr > 0 else 0.99
|
||
|
||
# Step 6: Risk level labels (0.0-1.0)
|
||
lf = lf.with_columns([
|
||
pl.col('_raw_score').alias('proxy_score'),
|
||
pl.when(pl.col('_raw_score') >= extreme_fence)
|
||
.then(pl.lit('critical'))
|
||
.when(pl.col('_raw_score') >= upper_fence)
|
||
.then(pl.lit('suspicious'))
|
||
.when(pl.col('_raw_score') >= q50)
|
||
.then(pl.lit('watch'))
|
||
.otherwise(pl.lit('normal'))
|
||
.alias('risk_level'),
|
||
(pl.col('_raw_score') * 0.6 + 0.4 * (pl.col('_raw_score') > q75).cast(pl.Float64))
|
||
.alias('threat_score'),
|
||
]).drop('_raw_score')
|
||
|
||
valid_counts = {c: int((mat[:, i] != 0).sum()) for i, c in enumerate(avail)}
|
||
logger.info(
|
||
'[SCORE] adaptive weights from PCA: %s | thresholds: p50=%.3f upper=%.3f extreme=%.3f | sample=%d/%d',
|
||
{c: f'{w:.3f}' for c, w in sorted(weights.items(), key=lambda x: -x[1])},
|
||
float(q50), upper_fence, extreme_fence,
|
||
min(n_rows, 10000), n_rows,
|
||
)
|
||
|
||
except Exception as exc:
|
||
logger.warning('[SCORE] adaptive scoring skipped: %s', exc)
|
||
|
||
return lf
|