tianxuan/analysis/entity_aggregator.py

"""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

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',)
_BYTES_REV_KEYWORDS = ()  # no user column for reverse bytes
_DURATION_KEYWORDS = ('4dur',)
_PACKETS_KEYWORDS = ('8ppk',)
_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 = ('latd',)
_LON_KEYWORDS = ('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: look for _norm column
    tls_norm = _find_column(schema, ('0ver_norm',))
    if not tls_norm:
        tls_norm = _find_column(schema, ('0ver',))
    if tls_norm:
        aggs.append(('tls_modern_ratio',
            pl.when(pl.col(tls_norm).is_in(['tls_1_3'])).then(1)
            .otherwise(0).cast(pl.Float64).mean()))
        aggs.append(('tls_old_ratio',
            pl.when(pl.col(tls_norm).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_norm)))

    # Modern cipher ratio (AEAD = modern)
    cph_norm = _find_column(schema, ('0cph_norm',))
    if not cph_norm:
        cph_norm = _find_column(schema, ('cipher_suite_norm',))
    if cph_norm:
        aggs.append(('aead_ratio',
            pl.when(pl.col(cph_norm).str.contains('gcm|chacha20|poly1305')).then(1)
            .otherwise(0).cast(pl.Float64).mean()))
        aggs.append(('weak_cipher_ratio',
            pl.when(pl.col(cph_norm).str.contains('cbc|rc4|null')).then(1)
            .otherwise(0).cast(pl.Float64).mean()))
        aggs.append(('ecdhe_ratio',
            pl.when(pl.col(cph_norm).str.contains('ecdhe')).then(1)
            .otherwise(0).cast(pl.Float64).mean()))
        aggs.append(('unique_ciphers', pl.n_unique(cph_norm)))

    # 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
    cnrs_norm = _find_column(schema, ('cnrs_norm',))
    if cnrs_norm:
        aggs.append(('recoverable_ratio',
            pl.when(pl.col(cnrs_norm) == 'yes').then(1)
            .otherwise(0).cast(pl.Float64).mean()))

    isrs_norm = _find_column(schema, ('isrs_norm',))
    if isrs_norm:
        aggs.append(('resumed_ratio',
            pl.when(pl.col(isrs_norm) == 'yes').then(1)
            .otherwise(0).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]

    # 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)

    # 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)

    # 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)

    # 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