PM-pinou
301 lines
10 KiB
Python
301 lines
10 KiB
Python
"""天璇类型系统 — 仅8种固定类型，面向对象接口。
|
|
|
|
每个类型继承 :class:`ColumnType`，实现:
|
|
- ``distance()`` — 为该类型列计算距离矩阵（追加到 LazyFrame）
|
|
- ``display()`` — 将值格式化为可读字符串
|
|
- ``polars_dtype`` — 映射到 Polars 数据类型
|
|
- ``name`` — 中文名称
|
|
|
|
:data:`TYPE_REGISTRY` 是唯一的类型注册表，按名称索引。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Base type — all 8 types inherit from this
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class ColumnType(ABC):
|
|
"""Base class for all column types.
|
|
|
|
Subclasses MUST define:
|
|
- ``name`` (class attribute, Chinese)
|
|
- ``polars_dtype`` (class attribute)
|
|
- ``distance()``
|
|
- ``display()``
|
|
"""
|
|
|
|
name: str = ""
|
|
polars_dtype: pl.DataType = pl.Utf8
|
|
|
|
@abstractmethod
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
"""Compute per-row distances for *columns* of this type.
|
|
|
|
Returns a new LazyFrame with distance columns appended
|
|
(prefixed ``_dist_``).
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def display(self, value: Any) -> str:
|
|
"""Format *value* for human-readable display."""
|
|
...
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# 8 concrete types
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
class IPv4Type(ColumnType):
|
|
name = "IPv4"
|
|
polars_dtype = pl.Utf8
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
if len(columns) < 2:
|
|
return lf
|
|
from analysis.distance import ip_distance
|
|
return ip_distance(lf, columns)
|
|
|
|
def display(self, value: Any) -> str:
|
|
return str(value) if value else "-"
|
|
|
|
|
|
class StringType(ColumnType):
|
|
name = "字符串"
|
|
polars_dtype = pl.Utf8
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
from analysis.distance import string_distance
|
|
return string_distance(lf, columns)
|
|
|
|
def display(self, value: Any) -> str:
|
|
return str(value)[:80] if value else "-"
|
|
|
|
|
|
class EnumType(ColumnType):
|
|
name = "枚举"
|
|
polars_dtype = pl.Utf8
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
# 0-1 match distance: compare to mode
|
|
for col in columns:
|
|
mode_val = lf.select(pl.col(col).mode().first()).collect(streaming=True).item()
|
|
lf = lf.with_columns(
|
|
(pl.col(col) != mode_val).cast(pl.Float64).alias(f'_enum_dist_{col}')
|
|
)
|
|
return lf
|
|
|
|
def display(self, value: Any) -> str:
|
|
return str(value) if value else "-"
|
|
|
|
|
|
class FloatType(ColumnType):
|
|
name = "浮点"
|
|
polars_dtype = pl.Float64
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
# Normalised euclidean distance — handled in preprocessing pipeline
|
|
for col in columns:
|
|
lf = lf.with_columns(
|
|
pl.col(col).cast(pl.Float64).alias(f'_raw_{col}')
|
|
)
|
|
return lf
|
|
|
|
def display(self, value: Any) -> str:
|
|
try:
|
|
return f"{float(value):.2f}"
|
|
except (ValueError, TypeError):
|
|
return str(value)
|
|
|
|
|
|
class IntType(ColumnType):
|
|
name = "整数"
|
|
polars_dtype = pl.Int64
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
# Normalised manhattan distance — handled in preprocessing pipeline
|
|
for col in columns:
|
|
lf = lf.with_columns(
|
|
pl.col(col).cast(pl.Int64).alias(f'_raw_{col}')
|
|
)
|
|
return lf
|
|
|
|
def display(self, value: Any) -> str:
|
|
try:
|
|
return str(int(value))
|
|
except (ValueError, TypeError):
|
|
return str(value)
|
|
|
|
|
|
class BoolType(ColumnType):
|
|
name = "布尔"
|
|
polars_dtype = pl.Boolean
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
# 0-1 match distance
|
|
for col in columns:
|
|
lf = lf.with_columns(
|
|
pl.col(col).cast(pl.Float64).alias(f'_bool_dist_{col}')
|
|
)
|
|
return lf
|
|
|
|
def display(self, value: Any) -> str:
|
|
if value is None:
|
|
return "-"
|
|
return "是" if str(value).lower() in ('true', '1', '+', 't') else "否"
|
|
|
|
|
|
class BytesType(ColumnType):
|
|
name = "字节"
|
|
polars_dtype = pl.Utf8
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
from analysis.distance import bytes_distance
|
|
return bytes_distance(lf, columns)
|
|
|
|
def display(self, value: Any) -> str:
|
|
s = str(value).strip() if value else ""
|
|
if len(s) > 20:
|
|
return s[:20] + "..."
|
|
return s or "-"
|
|
|
|
|
|
class TimestampType(ColumnType):
|
|
name = "时间戳"
|
|
polars_dtype = pl.Float64
|
|
|
|
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
|
from analysis.distance import timestamp_fft_distance
|
|
return timestamp_fft_distance(lf, columns)
|
|
|
|
def display(self, value: Any) -> str:
|
|
try:
|
|
t = float(value)
|
|
if t < 1e8 or t > 2e12:
|
|
return str(value)
|
|
from datetime import datetime
|
|
return datetime.fromtimestamp(t).strftime("%Y-%m-%d %H:%M")
|
|
except (ValueError, TypeError, OSError):
|
|
return str(value) if value else "-"
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
# Registry — single source of truth
|
|
# ═══════════════════════════════════════════════════════════════════════
|
|
|
|
TYPE_REGISTRY: dict[str, ColumnType] = {
|
|
"枚举": EnumType(),
|
|
"字节": BytesType(),
|
|
"浮点": FloatType(),
|
|
"整数": IntType(),
|
|
"布尔": BoolType(),
|
|
"字符串": StringType(),
|
|
"IPv4": IPv4Type(),
|
|
"时间戳": TimestampType(),
|
|
}
|
|
|
|
# Alias map for config.yaml type strings
|
|
TYPE_ALIASES: dict[str, str] = {
|
|
'enum': '枚举', '枚举': '枚举',
|
|
'bytes': '字节', '字节': '字节', 'hex': '字节',
|
|
'float': '浮点', '浮点': '浮点',
|
|
'int': '整数', '整数': '整数',
|
|
'bool': '布尔', '布尔': '布尔', 'boolean': '布尔',
|
|
'string': '字符串', '字符串': '字符串', 'str': '字符串',
|
|
'ipv4': 'IPv4', 'IPv4': 'IPv4',
|
|
'timestamp': '时间戳', '时间戳': '时间戳',
|
|
}
|
|
|
|
|
|
def get_type(name: str) -> ColumnType:
|
|
"""Resolve a type name (Chinese or English alias) to a ColumnType."""
|
|
resolved = TYPE_ALIASES.get(name.strip().lower() if name.isascii() else name.strip())
|
|
if resolved is None:
|
|
return StringType() # fallback
|
|
return TYPE_REGISTRY[resolved]
|
|
|
|
|
|
def classify_column(name: str, series: pl.Series, config_type: str | None = None) -> ColumnType:
|
|
"""Classify a single column to a ColumnType.
|
|
|
|
Priority: config override → TlsDB lookup → name heuristics → value sampling.
|
|
"""
|
|
from analysis.tls_ref import get_tls_ref_dict
|
|
import re
|
|
|
|
# 1. Config override
|
|
if config_type is not None:
|
|
t = get_type(config_type)
|
|
if t.name != "字符串" or config_type.lower().strip() == 'string':
|
|
return t
|
|
|
|
# 2. TlsDB lookup
|
|
tls_ref = get_tls_ref_dict()
|
|
if name in tls_ref:
|
|
t = get_type(tls_ref[name]['data_type'])
|
|
return t
|
|
|
|
# 3. Name heuristics
|
|
_patterns: list[tuple[re.Pattern, str]] = [
|
|
(re.compile(r'^(?:src_ip|dst_ip|:ips|:ipd|server.ip|client.ip)$', re.I), 'IPv4'),
|
|
(re.compile(r'^(?::prs|:prd|scnt|dcnt|1ipp|4dbn|4srs|eiph|name|tabl|0ver|cnrs|isrs)$', re.I), '枚举'),
|
|
(re.compile(r'^(?:0cph|0crv|cipher|ecdhe)', re.I), '字节'),
|
|
(re.compile(r'^(?:4dur|8ses|2tmo|0rnt)$', re.I), '浮点'),
|
|
(re.compile(r'^(?:4ksz|8ack|8ppk|8pak|8seq|8did|8byt|row)$', re.I), '整数'),
|
|
(re.compile(r'^(?:timestamp|time|8dbd|created|updated)', re.I), '时间戳'),
|
|
(re.compile(r'\.(ispn|orgn|city|anon|doma|latd|lond)$', re.I), '字符串'),
|
|
(re.compile(r'^(?:cnam|snam|0rnd|source.node|crcc|orga|orgu|@)', re.I), '字符串'),
|
|
]
|
|
for pat, type_name in _patterns:
|
|
if pat.search(name):
|
|
return TYPE_REGISTRY[type_name]
|
|
|
|
# 4. Value-based sampling
|
|
non_null = series.drop_nulls()
|
|
if len(non_null) == 0:
|
|
return TYPE_REGISTRY["字符串"]
|
|
|
|
strs = [str(v).strip().lower() for v in non_null.head(100)]
|
|
ip4 = re.compile(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$')
|
|
hex_p = re.compile(r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$')
|
|
|
|
if all(ip4.match(v) for v in strs):
|
|
return TYPE_REGISTRY["IPv4"]
|
|
if all(hex_p.match(v) for v in strs):
|
|
return TYPE_REGISTRY["字节"]
|
|
if series.dtype.is_integer():
|
|
return TYPE_REGISTRY["整数"]
|
|
if series.dtype == pl.Float64:
|
|
return TYPE_REGISTRY["浮点"]
|
|
if series.dtype == pl.Boolean:
|
|
return TYPE_REGISTRY["布尔"]
|
|
if non_null.n_unique() < 50:
|
|
return TYPE_REGISTRY["枚举"]
|
|
|
|
return TYPE_REGISTRY["字符串"]
|
|
|
|
|
|
def classify_schema(
|
|
lf: pl.LazyFrame,
|
|
config_overrides: dict[str, str] | None = None,
|
|
) -> dict[str, ColumnType]:
|
|
"""Classify all columns in a LazyFrame."""
|
|
config_overrides = config_overrides or {}
|
|
df_sample = lf.collect(streaming=True)
|
|
result = {}
|
|
for col_name in df_sample.columns:
|
|
config_type = config_overrides.get(col_name)
|
|
result[col_name] = classify_column(
|
|
name=col_name,
|
|
series=df_sample[col_name],
|
|
config_type=config_type,
|
|
)
|
|
return result
|