PM-pinou
9508cc8d58
- globe.py: add streaming=True to all lf.head().collect() calls (3 places) - entities.py: chunked aggregation (100K batches) for datasets > 1M rows - constants.py: reduce MAX_SAMPLE_ROWS 50000->30000 for 8GB RAM - tools/__init__.py: add __all__ for '*' re-export (fix tool_registry import) - clustering.py: fix misplaced logger/import traceback syntax errors - analysis.py: fix indentation of pass statement - anomalies.py: fix misplaced import traceback syntax error
295 lines
13 KiB
Python
295 lines
13 KiB
Python
"""Entity profile building + adaptive scoring handlers."""
|
|
import traceback
|
|
from analysis.constants import MAX_DISTRIBUTION_SAMPLE, RANDOM_SEED
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from ..session_store import SessionStore, _generate_id
|
|
from ..profile_util import profile
|
|
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_DTYPES
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Constants
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
_ENTITY_KEYWORDS = [
|
|
'src_ip', 'dst_ip', 'source_ip', 'destination_ip', 'src_addr', 'dst_addr',
|
|
'sip', 'dip', 'ip_src', 'ip_dst', 'saddr', 'daddr',
|
|
'sni', 'host', 'domain', 'uri', 'url', 'fqdn', 'server_name',
|
|
'mac', 'mac_addr', 'src_mac', 'dst_mac',
|
|
'user_agent', 'ua', 'asn', 'organization', 'org',
|
|
]
|
|
|
|
_NUMERIC_TYPES_FOR_AGG = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
|
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
|
pl.Float32, pl.Float64)
|
|
|
|
_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',
|
|
]
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# _handle_build_entity_profiles
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
async def _handle_build_entity_profiles(
|
|
dataset_id: str,
|
|
entity_column: Optional[str] = None,
|
|
entity_columns: Optional[list[str]] = None,
|
|
auto_detect: bool = True,
|
|
) -> dict:
|
|
"""Group raw flow data by entity column(s) and compute per-entity aggregate features."""
|
|
entry, err = _resolve_dataset(dataset_id)
|
|
if err:
|
|
return err
|
|
|
|
lf: pl.LazyFrame = entry['lazyframe']
|
|
schema = entry.get('schema', {})
|
|
col_names = list(schema.keys())
|
|
|
|
# ── Resolve entity columns ──────────────────────────────────────────
|
|
if entity_column and entity_column in col_names:
|
|
group_cols = [entity_column]
|
|
elif entity_columns:
|
|
group_cols = [c for c in entity_columns if c in col_names]
|
|
if not group_cols:
|
|
return {'error': f'None of entity_columns {entity_columns} exist in dataset. Available: {col_names[:20]}',
|
|
'truncated': False}
|
|
elif auto_detect:
|
|
matched = []
|
|
for kw in _ENTITY_KEYWORDS:
|
|
found = [c for c in col_names if kw.lower() in c.lower()]
|
|
matched.extend(found)
|
|
seen = set()
|
|
group_cols = [c for c in matched if not (c in seen or seen.add(c))]
|
|
if not group_cols:
|
|
for c in col_names:
|
|
if c.startswith('_'):
|
|
continue
|
|
group_cols = [c]
|
|
break
|
|
group_cols = group_cols[:3]
|
|
else:
|
|
return {'error': 'No entity column specified and auto_detect is disabled',
|
|
'truncated': False}
|
|
|
|
if not group_cols:
|
|
return {'error': 'Could not determine entity column(s) for grouping',
|
|
'truncated': False}
|
|
|
|
# ── Build aggregation expressions ──────────────────────────────────
|
|
agg_exprs: list[pl.Expr] = [
|
|
pl.len().alias('flow_count'),
|
|
pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
|
|
]
|
|
|
|
numeric_cols = [c for c in col_names if c not in group_cols
|
|
and schema[c].startswith(('Int', 'UInt', 'Float'))]
|
|
for c in numeric_cols[:20]:
|
|
agg_exprs.append(pl.sum(c).alias(f'total_{c}'))
|
|
agg_exprs.append(pl.mean(c).alias(f'avg_{c}'))
|
|
|
|
str_cols = [c for c in col_names if c not in group_cols
|
|
and schema[c].startswith(('Utf8', 'String', 'Cat'))]
|
|
for c in str_cols[:10]:
|
|
agg_exprs.append(pl.col(c).n_unique().alias(f'unique_{c}'))
|
|
|
|
for kw in ['dst_ip', 'dip', 'dst_addr', 'daddr', 'ip_dst', 'dest_ip',
|
|
'dst_port', 'dport', 'destination_port', 'dest_port',
|
|
'port', 'protocol', 'proto', 'sni', 'server_name',
|
|
'cipher_suite', 'cipher', 'tls_version', 'version', '0ver',
|
|
'country', 'scnt', 'dcnt', 'asn', 'organization']:
|
|
found = [c for c in col_names if c not in group_cols and kw.lower() in c.lower()]
|
|
for c in found:
|
|
alias_name = f'unique_{c}'
|
|
if alias_name not in [str(a.meta.output_name()) if hasattr(a, 'meta') else '' for a in agg_exprs]:
|
|
agg_exprs.append(pl.col(c).n_unique().alias(alias_name))
|
|
|
|
# ── Execute grouping ───────────────────────────────────────────────
|
|
row_count = _count_rows(lf)
|
|
ENTITY_CHUNK_THRESHOLD = 1_000_000 # 1M rows
|
|
ENTITY_CHUNK_SIZE = 100_000 # 100K rows per chunk
|
|
|
|
if row_count >= ENTITY_CHUNK_THRESHOLD:
|
|
# Chunked path: process in 100K batches → group → merge across chunks
|
|
all_chunk_dfs = []
|
|
for offset in range(0, row_count, ENTITY_CHUNK_SIZE):
|
|
chunk_df = lf.slice(offset, ENTITY_CHUNK_SIZE).collect(streaming=True)
|
|
chunk_grouped = chunk_df.group_by(group_cols).agg(agg_exprs)
|
|
all_chunk_dfs.append(chunk_grouped)
|
|
merged_df = pl.concat(all_chunk_dfs)
|
|
|
|
# Re-aggregate across chunks (same entity may appear in multiple chunks)
|
|
re_agg_exprs: list[pl.Expr] = []
|
|
for col_name in merged_df.columns:
|
|
if col_name in group_cols:
|
|
continue
|
|
if col_name == 'flow_count':
|
|
re_agg_exprs.append(pl.sum('flow_count').alias('flow_count'))
|
|
elif col_name.startswith('total_'):
|
|
re_agg_exprs.append(pl.sum(col_name).alias(col_name))
|
|
elif col_name.startswith('avg_'):
|
|
# Weighted average: sum(total) / sum(flow_count)
|
|
base = col_name[4:]
|
|
total_col = f'total_{base}'
|
|
if total_col in merged_df.columns:
|
|
re_agg_exprs.append(
|
|
(pl.sum(total_col) / pl.sum('flow_count')).alias(col_name))
|
|
elif col_name.startswith('unique_'):
|
|
# Approximate: max of per-chunk distinct counts
|
|
re_agg_exprs.append(pl.max(col_name).alias(col_name))
|
|
|
|
entity_lf = merged_df.lazy().group_by(group_cols).agg(re_agg_exprs)
|
|
logger.info('[ENTITY] chunked aggregation: %d rows → %d chunks → %d merged rows',
|
|
row_count, len(all_chunk_dfs), len(merged_df))
|
|
else:
|
|
entity_lf = lf.group_by(group_cols).agg(agg_exprs)
|
|
|
|
n_entities = _count_rows(entity_lf)
|
|
|
|
# ── Store result ───────────────────────────────────────────────────
|
|
new_schema = {name: str(dtype) for name, dtype in
|
|
zip(entity_lf.collect_schema().names(), entity_lf.collect_schema().dtypes())}
|
|
new_id = f'entity_{_generate_id("ds")}'
|
|
store = SessionStore()
|
|
store.store_dataset(
|
|
dataset_id=new_id,
|
|
lazyframe=entity_lf,
|
|
schema=new_schema,
|
|
metadata={
|
|
'row_count': n_entities,
|
|
'entity_columns': group_cols,
|
|
'parent_dataset_id': dataset_id,
|
|
'n_agg_features': len(agg_exprs),
|
|
},
|
|
parent_id=dataset_id,
|
|
)
|
|
|
|
return {
|
|
'dataset_id': new_id,
|
|
'entity_columns': group_cols,
|
|
'n_entities': n_entities,
|
|
'n_features': len(agg_exprs),
|
|
'aggregated_columns': list(new_schema.keys())[:20],
|
|
}
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Adaptive scoring helpers
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
@profile
|
|
def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
|
|
"""Add anomaly scores using data-driven weights (SVD + IQR + percentile)."""
|
|
_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:
|
|
n_rows = _count_rows(lf)
|
|
sample_lf = lf if n_rows <= MAX_DISTRIBUTION_SAMPLE else lf.sample(n=MAX_DISTRIBUTION_SAMPLE, seed=RANDOM_SEED)
|
|
df_sample = sample_lf.select(avail).collect(streaming=True)
|
|
mat = df_sample.to_numpy()
|
|
mat = np.nan_to_num(mat, nan=0.0)
|
|
|
|
from sklearn.decomposition import TruncatedSVD
|
|
n_comp = min(1, min(mat.shape) - 1)
|
|
if n_comp < 1 or mat.shape[1] < 2:
|
|
return lf
|
|
svd = TruncatedSVD(n_components=1, random_state=RANDOM_SEED).fit(mat)
|
|
raw_weights = np.abs(svd.components_[0])
|
|
weights = dict(zip(avail, raw_weights / raw_weights.sum()))
|
|
|
|
expr_parts = []
|
|
for col, w in weights.items():
|
|
if col == 'modern_tls_rate':
|
|
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'))
|
|
|
|
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()
|
|
|
|
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
|
|
|
|
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 = logging.getLogger(__name__)
|
|
logger.info(
|
|
'[SCORE] adaptive weights from SVD: %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, MAX_DISTRIBUTION_SAMPLE), n_rows,
|
|
)
|
|
|
|
except Exception as exc:
|
|
logging.getLogger(__name__).warning('[SCORE] adaptive scoring skipped: %s', exc)
|
|
|
|
return lf
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# _handle_compute_scores
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
async def _handle_compute_scores(dataset_id: str) -> dict:
|
|
"""Compute adaptive proxy/anomaly/threat scores on entity data."""
|
|
entry, err = _resolve_dataset(dataset_id)
|
|
if err:
|
|
return err
|
|
|
|
lf = entry['lazyframe']
|
|
lf = _add_adaptive_scores(lf, [])
|
|
try:
|
|
updated_schema = lf.collect_schema()
|
|
new_schema = {name: str(dtype) for name, dtype in
|
|
zip(updated_schema.names(), updated_schema.dtypes())}
|
|
except Exception:
|
|
logger.error("_handle_compute_scores failed: {}".format(traceback.format_exc()))
|
|
new_schema = entry.get('schema', {})
|
|
store = SessionStore()
|
|
store.store_dataset(dataset_id, lf, schema=new_schema,
|
|
metadata=entry.get('metadata', {}))
|
|
return {'status': 'scored', 'dataset_id': dataset_id}
|