v7: 19 issues fixed — SQLite storage, new clustering pipeline, display_id, globe rewrite, Chinese tools
Key changes: - New: data_loader.py SQLite persistence with drop_sqlite_table - New: db_utils.py retry_on_lock decorator (3 retries, exponential backoff) - New: tool_registry.py with 27 MCP tools (filter_and_cluster, compute_scores, etc.) - New: tls_ref.py for TLS cipher/reference data - New: import_tlsdb.py management command - New: scripts/start_server.py for portable runtime - New: migrations 0003-0007 for SQLite table, display_id, llm fields - Changed: views.py unified pipeline worker, retry_run, display_id everywhere - Changed: models.py with display_id auto-assignment, run_type, sqlite_table, llm_thinking, tool_calls_json - Changed: urls.py added retry_run route - Changed: session_store.py robust JSON persistence - Changed: AGENTS.md v7 fix summary added - Changed: templates — globe rewrite (inertia/polar flip), auto.html (thinking/tool accordions), base.html (toast/config), all pages use display_id - Changed: run.bat PYTHONUTF8=1 - Deleted: entity_detector.py, entity_aggregator.py (replaced by filter_and_cluster clustering pipeline) - Test: 92/92 unit tests passing
This commit is contained in:
PM-pinou
@@ -7,6 +7,7 @@ validates schema consistency, and returns a merged :class:`pl.LazyFrame`.
|
||||
import glob
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
@@ -169,6 +170,113 @@ def load_config(config_path: Optional[str]) -> dict:
|
||||
return yaml.safe_load(f) or {}
|
||||
|
||||
|
||||
# ── SQLite persistence ──────────────────────────────────────────────────
|
||||
|
||||
def _dtype_to_sqlite(dtype: pl.DataType) -> str:
|
||||
"""Map a Polars dtype to a SQLite column type affinity."""
|
||||
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
|
||||
return 'INTEGER'
|
||||
if dtype in (pl.Float32, pl.Float64):
|
||||
return 'REAL'
|
||||
if dtype == pl.Boolean:
|
||||
return 'INTEGER'
|
||||
return 'TEXT' # String, Date, Datetime, etc. all map to TEXT
|
||||
|
||||
|
||||
def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str:
|
||||
"""Write a LazyFrame to a SQLite table ``_data_{run_id}``.
|
||||
|
||||
Collects the LazyFrame (streaming if possible), creates or replaces the
|
||||
table, bulk-inserts rows in chunks, and returns the table name.
|
||||
|
||||
Args:
|
||||
run_id: Analysis run ID (used as ``_data_{run_id}``).
|
||||
lf: LazyFrame to persist.
|
||||
mode: ``'replace'`` (default) → drop+create+insert.
|
||||
``'append'`` → skip DDL, only INSERT (table must exist).
|
||||
|
||||
This is a *persistence copy* — the original CSV files remain untouched.
|
||||
"""
|
||||
from django.conf import settings
|
||||
|
||||
table_name = f"_data_{run_id}"
|
||||
db_path = settings.DATABASES['default']['NAME']
|
||||
|
||||
df = lf.collect(streaming=True)
|
||||
conn = sqlite3.connect(db_path)
|
||||
|
||||
try:
|
||||
if mode == 'replace':
|
||||
# Drop existing table for retry/idempotency safety
|
||||
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
|
||||
|
||||
# Build CREATE TABLE from Polars schema
|
||||
col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}'
|
||||
for name, dtype in zip(df.columns, df.dtypes)]
|
||||
conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})')
|
||||
|
||||
# Bulk insert in chunks
|
||||
col_names = [f'"{c}"' for c in df.columns]
|
||||
placeholders = ','.join(['?' for _ in df.columns])
|
||||
insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})'
|
||||
|
||||
total_rows = len(df)
|
||||
chunk_size = 10000
|
||||
for i in range(0, total_rows, chunk_size):
|
||||
chunk = df.slice(i, chunk_size)
|
||||
conn.executemany(insert_sql, chunk.rows())
|
||||
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode)
|
||||
return table_name
|
||||
|
||||
|
||||
def load_from_db(table_name: str) -> Optional[pl.LazyFrame]:
|
||||
"""Load a SQLite table into a LazyFrame.
|
||||
|
||||
Returns ``None`` if the table does not exist.
|
||||
"""
|
||||
from django.conf import settings
|
||||
|
||||
db_path = settings.DATABASES['default']['NAME']
|
||||
conn = sqlite3.connect(db_path)
|
||||
|
||||
try:
|
||||
# Check table exists
|
||||
cursor = conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
|
||||
(table_name,),
|
||||
)
|
||||
if cursor.fetchone() is None:
|
||||
return None
|
||||
|
||||
conn.row_factory = sqlite3.Row
|
||||
rows = conn.execute(f'SELECT * FROM "{table_name}"').fetchall()
|
||||
if not rows:
|
||||
return pl.DataFrame().lazy()
|
||||
df = pl.from_dicts([dict(r) for r in rows])
|
||||
return df.lazy()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def drop_sqlite_table(table_name: str) -> None:
|
||||
"""Drop a SQLite data table silently if it exists."""
|
||||
from django.conf import settings
|
||||
|
||||
db_path = settings.DATABASES['default']['NAME']
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ── helpers ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _file_count(paths: list[str]) -> int:
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Database utilities for concurrent access."""
|
||||
import functools
|
||||
import time
|
||||
import logging
|
||||
from django.db import OperationalError, close_old_connections
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def retry_on_lock(max_retries=3, base_delay=0.5):
|
||||
"""Decorator: retry on OperationalError('database is locked')."""
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
for attempt in range(max_retries):
|
||||
close_old_connections()
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except OperationalError as e:
|
||||
if 'database is locked' in str(e) and attempt < max_retries - 1:
|
||||
delay = base_delay * (2 ** attempt)
|
||||
logger.warning(f'[DB_LOCKED] retry {attempt+1}/{max_retries} in {delay:.1f}s')
|
||||
time.sleep(delay)
|
||||
close_old_connections()
|
||||
else:
|
||||
raise
|
||||
return wrapper
|
||||
return decorator
|
||||
@@ -1,587 +0,0 @@
|
||||
"""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
|
||||
@@ -1,177 +0,0 @@
|
||||
"""Auto-detect entity column from CSV schema using heuristic scoring.
|
||||
|
||||
Provides :func:`detect_entity_column` which scores all string/categorical
|
||||
columns by keyword-name match, unique-value ratio, and dtype, returning
|
||||
a ranked list of candidates and a recommended column.
|
||||
|
||||
Usage::
|
||||
|
||||
from analysis.entity_detector import detect_entity_column
|
||||
|
||||
result = detect_entity_column(dataset_id)
|
||||
result['recommended'] # e.g. 'src_ip'
|
||||
result['candidates'] # sorted list of scored columns
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import polars as pl
|
||||
|
||||
from .session_store import SessionStore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entity-related keywords (lowercased)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ENTITY_KEYWORDS: tuple[str, ...] = (
|
||||
# User exact column names only — no generic/fuzzy keywords
|
||||
':ips', ':ipd', ':prs', ':prd',
|
||||
'server_ip', 'client_ip', 'scnt', 'dcnt',
|
||||
'cnam', 'snam', '0ver', '4dur', '8ses', '2tmo', '4ksz',
|
||||
'cnrs', 'isrs', '8ack', '8ppk', '8dbd',
|
||||
'1ipp', '4dbn', 'tabl', 'name', 'source_node',
|
||||
'cipher_suite', 'ecdhe_named_curve',
|
||||
'0cph', '0crv', '0rnd', '0rnt', 'row', 'time', 'timestamp',
|
||||
'latd', 'lond', 'ispn', 'orgn', 'city',
|
||||
)
|
||||
|
||||
|
||||
def _score_name(column: str) -> tuple[float, str]:
|
||||
"""Score a column name against *ENTITY_KEYWORDS*.
|
||||
|
||||
Returns ``(score, matched_keyword_or_empty)`` where *score* is 3.0
|
||||
for every matched keyword. Partial matches (keyword is a substring
|
||||
of the lowercased column name) also count.
|
||||
"""
|
||||
col_lower = column.lower().replace('-', '_').replace('.', '_').replace(' ', '_')
|
||||
for kw in ENTITY_KEYWORDS:
|
||||
if kw == col_lower or col_lower.endswith(f'_{kw}') or col_lower.startswith(f'{kw}_'):
|
||||
return 3.0, kw
|
||||
return 0.0, ''
|
||||
|
||||
|
||||
def detect_entity_column(dataset_id: str, entity_columns: Optional[list[str]] = None) -> dict:
|
||||
"""Score all string/categorical columns and return candidates.
|
||||
|
||||
The scoring algorithm:
|
||||
|
||||
* **Name match** (``+3``): column name (normalised) is a keyword, or
|
||||
ends/starts with a keyword via ``_`` separator.
|
||||
* **Unique ratio** (0–10): columns whose ratio of unique values to
|
||||
non-null values falls in the range ``[0.01, 0.50]`` receive up to
|
||||
10 points (``min(ratio * 100, 10)``).
|
||||
* **String dtype** (``+2``): ``Utf8`` / ``String`` / ``Categorical``
|
||||
columns get a bonus.
|
||||
* **Null penalty** (``*0.1``): if null ratio > 50%, score severely reduced.
|
||||
* **Penalty** (``-5``): unique ratio > 0.50 (likely random IDs or
|
||||
transaction-level identifiers).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dataset_id:
|
||||
Dataset identifier in the session store.
|
||||
entity_columns:
|
||||
Optional list of column names to restrict detection to. When given,
|
||||
only these columns are scored (must exist in the schema).
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
``candidates``: list of ``{column, score, unique_ratio, matched_keyword}``
|
||||
sorted descending by score.
|
||||
|
||||
``recommended``: highest-scoring column name (or ``''`` when none found).
|
||||
|
||||
``recommended_multi``: list of the top 3 candidate column names (may be
|
||||
shorter than 3 when fewer candidates exist).
|
||||
|
||||
``dataset_id``: the input identifier (echoed).
|
||||
"""
|
||||
store = SessionStore()
|
||||
entry = store.get_dataset(dataset_id)
|
||||
if entry is None:
|
||||
return {
|
||||
'dataset_id': dataset_id,
|
||||
'candidates': [],
|
||||
'recommended': '',
|
||||
'error': f'Dataset not found: {dataset_id}',
|
||||
}
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema: dict[str, str] = entry.get('schema', {})
|
||||
|
||||
# Identify string / categorical columns
|
||||
string_types = frozenset({'Utf8', 'String', 'Categorical', 'str', 'cat'})
|
||||
candidate_scores: list[dict] = []
|
||||
|
||||
try:
|
||||
# Collect a sample to compute unique ratios efficiently
|
||||
df_sample = lf.collect(streaming=True)
|
||||
except Exception as exc:
|
||||
return {
|
||||
'dataset_id': dataset_id,
|
||||
'candidates': [],
|
||||
'recommended': '',
|
||||
'error': f'Failed to collect dataset: {exc}',
|
||||
}
|
||||
|
||||
total_rows = len(df_sample)
|
||||
|
||||
for col_name, dtype_str in schema.items():
|
||||
# If entity_columns specified, skip columns not in the list
|
||||
if entity_columns is not None and col_name not in entity_columns:
|
||||
continue
|
||||
|
||||
# Only score string-like columns
|
||||
dtype_base = dtype_str.split('(')[0].strip()
|
||||
if dtype_base not in string_types:
|
||||
continue
|
||||
|
||||
series_all = df_sample[col_name]
|
||||
non_null_series = series_all.drop_nulls()
|
||||
non_null_count = len(non_null_series)
|
||||
null_count = total_rows - non_null_count
|
||||
|
||||
if non_null_count == 0:
|
||||
continue
|
||||
|
||||
unique_count = int(non_null_series.n_unique())
|
||||
unique_ratio = unique_count / non_null_count
|
||||
|
||||
# --- scoring ---
|
||||
name_score, matched_kw = _score_name(col_name)
|
||||
dtype_score = 2.0 # String/categorical bonus
|
||||
ratio_score = min(unique_ratio * 100.0, 10.0)
|
||||
penalty = -5.0 if unique_ratio > 0.50 else 0.0
|
||||
|
||||
total_score = name_score + dtype_score + ratio_score + penalty
|
||||
|
||||
# Null ratio penalty: if more than 50% null, severely reduce score
|
||||
if null_count / max(1, non_null_count) > 0.5:
|
||||
total_score *= 0.1
|
||||
|
||||
candidate_scores.append({
|
||||
'column': col_name,
|
||||
'score': round(total_score, 2),
|
||||
'unique_ratio': round(unique_ratio, 4),
|
||||
'matched_keyword': matched_kw,
|
||||
'non_null_count': non_null_count,
|
||||
'unique_count': unique_count,
|
||||
'null_count': null_count,
|
||||
})
|
||||
|
||||
# Sort descending by score, then by unique_count desc (prefer richer columns)
|
||||
candidate_scores.sort(key=lambda c: (c['score'], c['unique_count']), reverse=True)
|
||||
|
||||
recommended = candidate_scores[0]['column'] if candidate_scores else ''
|
||||
recommended_multi = [c['column'] for c in candidate_scores[:3]]
|
||||
|
||||
return {
|
||||
'dataset_id': dataset_id,
|
||||
'candidates': candidate_scores,
|
||||
'recommended': recommended,
|
||||
'recommended_multi': recommended_multi, # NEW
|
||||
'total_rows': total_rows,
|
||||
'scored_columns': len(candidate_scores),
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Django management command: ``python manage.py import_tlsdb``
|
||||
|
||||
Reads ``TlsDB.csv`` from the project root and populates the
|
||||
``tls_ref_data`` SQLite table, replacing the CSV as the authoritative
|
||||
source of TLS field-code reference metadata.
|
||||
|
||||
Usage:
|
||||
python manage.py import_tlsdb
|
||||
python manage.py import_tlsdb --csv /path/to/TlsDB.csv
|
||||
python manage.py import_tlsdb --clear # re-import from scratch
|
||||
"""
|
||||
|
||||
import csv
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
from django.db import connection, transaction
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
"""Import ``TlsDB.csv`` field definitions into ``tls_ref_data``."""
|
||||
|
||||
help = 'Import TlsDB.csv reference data into the tls_ref_data SQLite table'
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
'--csv', type=str, default=None,
|
||||
help='Path to TlsDB.csv (default: project-root/TlsDB.csv)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--clear', action='store_true', default=False,
|
||||
help='Clear existing data before importing (default: upsert)',
|
||||
)
|
||||
|
||||
def _resolve_csv_path(self, provided: str | None) -> Path:
|
||||
"""Locate the TlsDB.csv file.
|
||||
|
||||
Priority:
|
||||
1. Explicit ``--csv`` path.
|
||||
2. Project root (where ``manage.py`` lives).
|
||||
"""
|
||||
if provided:
|
||||
path = Path(provided).resolve()
|
||||
else:
|
||||
# manage.py lives in the project root — use that as anchor
|
||||
manage_dir = Path(__file__).resolve().parent.parent.parent.parent
|
||||
path = manage_dir / 'TlsDB.csv'
|
||||
|
||||
if not path.exists():
|
||||
raise CommandError(f'TlsDB.csv not found at: {path}')
|
||||
return path
|
||||
|
||||
def handle(self, *args, **options):
|
||||
csv_path = self._resolve_csv_path(options['csv'])
|
||||
clear = options['clear']
|
||||
|
||||
# ── Parse CSV ────────────────────────────────────────────────────────
|
||||
rows: list[dict[str, str]] = []
|
||||
with open(csv_path, 'r', encoding='utf-8-sig') as f:
|
||||
reader = csv.DictReader(f)
|
||||
# Expected columns: 标题,含义,数据类型
|
||||
for line_no, row in enumerate(reader, start=2):
|
||||
field_code = (row.get('标题') or '').strip()
|
||||
meaning = (row.get('含义') or '').strip()
|
||||
data_type = (row.get('数据类型') or '').strip()
|
||||
if not field_code:
|
||||
self.stderr.write(self.style.WARNING(
|
||||
f' [SKIP] Line {line_no}: empty field_code, skipping'
|
||||
))
|
||||
continue
|
||||
rows.append({
|
||||
'field_code': field_code,
|
||||
'meaning': meaning,
|
||||
'data_type': data_type,
|
||||
})
|
||||
|
||||
if not rows:
|
||||
self.stderr.write(self.style.ERROR('No valid rows found in TlsDB.csv'))
|
||||
return
|
||||
|
||||
# ── Insert into tls_ref_data ─────────────────────────────────────────
|
||||
with connection.cursor() as cursor:
|
||||
with transaction.atomic():
|
||||
if clear:
|
||||
cursor.execute('DELETE FROM tls_ref_data')
|
||||
self.stdout.write(self.style.WARNING(f' Cleared existing data ({len(rows)} rows to insert)'))
|
||||
|
||||
inserted = 0
|
||||
updated = 0
|
||||
for r in rows:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO tls_ref_data (field_code, meaning, data_type)
|
||||
VALUES (%s, %s, %s)
|
||||
ON CONFLICT(field_code) DO UPDATE SET
|
||||
meaning = EXCLUDED.meaning,
|
||||
data_type = EXCLUDED.data_type
|
||||
""",
|
||||
[r['field_code'], r['meaning'], r['data_type']],
|
||||
)
|
||||
if cursor.rowcount == 1:
|
||||
inserted += 1
|
||||
else:
|
||||
updated += 1
|
||||
|
||||
self.stdout.write(self.style.SUCCESS(
|
||||
f'\n[OK] Imported {len(rows)} TLS field definitions from: {csv_path}'
|
||||
))
|
||||
self.stdout.write(self.style.SUCCESS(
|
||||
f' New: {inserted} | Updated: {updated} | Total: {len(rows)}'
|
||||
))
|
||||
@@ -1,19 +1,17 @@
|
||||
"""Django management command: ``python manage.py run_pipeline``
|
||||
|
||||
Runs the complete TLS analysis pipeline from CLI, no LLM required.
|
||||
Usage: uv run python manage.py run_pipeline <csv_glob> [--entity-col NAME] [--algo hdbscan|kmeans]
|
||||
Runs the analysis pipeline from CLI: load → cluster → extract features.
|
||||
Usage: uv run python manage.py run_pipeline <csv_glob> [--algo hdbscan|kmeans]
|
||||
"""
|
||||
import os, sys, json, asyncio
|
||||
from django.core.management.base import BaseCommand, CommandError
|
||||
from django.utils import timezone
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = 'Run the full TLS analysis pipeline from CLI (no LLM needed)'
|
||||
help = 'Run the analysis pipeline from CLI (load → cluster → extract)'
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument('csv_glob', type=str, help='Glob pattern for CSV files')
|
||||
parser.add_argument('--entity-col', type=str, default=None,
|
||||
help='Entity column (auto-detect if omitted)')
|
||||
parser.add_argument('--algo', type=str, default='hdbscan', choices=['hdbscan', 'kmeans'],
|
||||
help='Clustering algorithm (default: hdbscan)')
|
||||
parser.add_argument('--output', type=str, default=None,
|
||||
@@ -25,7 +23,6 @@ class Command(BaseCommand):
|
||||
from analysis.views import _run_pipeline_worker
|
||||
|
||||
csv_glob = options['csv_glob']
|
||||
entity_col = options['entity_col']
|
||||
algo = options['algo']
|
||||
output = options['output']
|
||||
|
||||
@@ -42,14 +39,12 @@ class Command(BaseCommand):
|
||||
warnings.filterwarnings('ignore', category=RuntimeWarning, module='sklearn')
|
||||
from analysis.session_store import SessionStore
|
||||
from analysis.data_loader import load_csv_directory
|
||||
from analysis.entity_detector import detect_entity_column
|
||||
from analysis.entity_aggregator import aggregate_by_entity
|
||||
|
||||
store = SessionStore()
|
||||
store.drop_all()
|
||||
|
||||
# ── Step 1: Load CSV ─────────────────────────────────────────────
|
||||
_self.stdout.write(f'[1/5] 加载 CSV: {csv_glob}')
|
||||
_self.stdout.write(f'[1/3] 加载 CSV: {csv_glob}')
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
|
||||
ds_id = 'manual_ds'
|
||||
store.store_dataset(ds_id, lf, schema=schema, metadata={
|
||||
@@ -57,47 +52,25 @@ class Command(BaseCommand):
|
||||
})
|
||||
_self.stdout.write(_self.style.SUCCESS(f' → {row_count} rows, {file_count} files, {memory_mb:.1f} MB'))
|
||||
|
||||
# ── Step 2: Detect entity column ────────────────────────────────
|
||||
_self.stdout.write('[2/5] 检测实体列')
|
||||
if entity_col:
|
||||
used_col = entity_col
|
||||
_self.stdout.write(f' → 使用指定列: {used_col}')
|
||||
else:
|
||||
result = detect_entity_column(ds_id)
|
||||
candidates = result.get('candidates', [])
|
||||
used_col = result.get('recommended')
|
||||
if not used_col:
|
||||
_self.stderr.write(_self.style.ERROR(' ✗ 未检测到实体列，请用 --entity-col 手动指定'))
|
||||
run.status = 'failed'
|
||||
run.error_message = '未检测到实体列'
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
_self.stdout.write(f' → 自动检测: {used_col} (共 {len(candidates)} 候选)')
|
||||
|
||||
# ── Step 3: Aggregate by entity ──────────────────────────────────
|
||||
_self.stdout.write('[3/5] 实体聚合')
|
||||
agg_lf, feature_columns = aggregate_by_entity(lf, used_col, schema)
|
||||
df_agg = agg_lf.collect(streaming=True)
|
||||
entity_ds_id = 'manual_entity'
|
||||
store.store_dataset(entity_ds_id, agg_lf, schema=dict(zip(feature_columns, ['']*len(feature_columns))),
|
||||
metadata={'row_count': len(df_agg), 'entity_column': used_col, 'csv_glob': csv_glob})
|
||||
_self.stdout.write(_self.style.SUCCESS(f' → {len(df_agg)} 个实体, {len(feature_columns)} 维特征'))
|
||||
|
||||
# ── Step 4: Clustering + Extraction + PCA (shared pipeline) ─────
|
||||
_self.stdout.write(f'[4/5] 聚类 ({algo})')
|
||||
from analysis.views import _run_clustering_pipeline
|
||||
user_features = [
|
||||
c for c in feature_columns
|
||||
if not c.startswith('_') and c not in ('first_seen', 'last_seen', 'src_ip')
|
||||
# Auto-select numeric columns for clustering
|
||||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64',
|
||||
'UInt8', 'UInt16', 'UInt32', 'UInt64',
|
||||
'Float32', 'Float64'}
|
||||
feature_cols = [
|
||||
c for c in schema.keys()
|
||||
if not c.startswith('_')
|
||||
and schema[c].split('(')[0].strip() in numeric_types
|
||||
][:10]
|
||||
_self.stdout.write(f' → 聚类特征: {user_features}')
|
||||
run.entity_column = used_col
|
||||
|
||||
# ── Step 2: Clustering + Extraction + PCA (shared pipeline) ─────
|
||||
_self.stdout.write(f'[2/3] 聚类 ({algo})')
|
||||
_self.stdout.write(f' → 聚类特征: {feature_cols}')
|
||||
run.total_flows = row_count
|
||||
run.entity_count = len(df_agg)
|
||||
run.save(update_fields=['entity_column', 'total_flows', 'entity_count'])
|
||||
run.save(update_fields=['total_flows'])
|
||||
from analysis.views import _run_clustering_pipeline
|
||||
_run_clustering_pipeline(
|
||||
run=run, store=store, entity_ds_id=entity_ds_id,
|
||||
feature_columns=user_features,
|
||||
run=run, store=store, entity_ds_id=ds_id,
|
||||
feature_columns=feature_cols,
|
||||
algorithm=algo, min_cluster_size=5,
|
||||
run_pca=True,
|
||||
)
|
||||
@@ -112,26 +85,20 @@ class Command(BaseCommand):
|
||||
if output:
|
||||
from analysis.tool_registry import _handle_export_results
|
||||
os.makedirs(output, exist_ok=True)
|
||||
asyncio.run(_handle_export_results(result_id=entity_ds_id, output_path=output, format='csv', overwrite=True))
|
||||
asyncio.run(_handle_export_results(result_id=ds_id, output_path=output, format='csv', overwrite=True))
|
||||
_self.stdout.write(_self.style.SUCCESS(f' → 导出到 {output}'))
|
||||
|
||||
# ── Update run record with final stats ──────────────────────────
|
||||
run.total_flows = row_count
|
||||
run.entity_count = len(df_agg)
|
||||
run.entity_column = used_col
|
||||
run.save(update_fields=['total_flows', 'entity_count', 'entity_column'])
|
||||
|
||||
_run_pipeline_worker(run_id, _analysis_fn)
|
||||
|
||||
# ── Read final status and print summary ─────────────────────────────
|
||||
run = AnalysisRun.objects.get(id=run_id)
|
||||
if run.status == 'failed':
|
||||
err_msg = (run.error_message or '')[:500]
|
||||
self.stderr.write(self.style.ERROR(f'\n*** 分析失败! Run ID: #{run.id}'))
|
||||
self.stderr.write(self.style.ERROR(f'\n*** 分析失败! Run ID: #{run.display_id}'))
|
||||
if err_msg:
|
||||
self.stderr.write(self.style.ERROR(f' {err_msg}'))
|
||||
return
|
||||
self.stdout.write(self.style.SUCCESS(
|
||||
f'\n*** 分析完成! Run ID: #{run.id}'
|
||||
f' 查看 Django: http://127.0.0.1:8000/runs/{run.id}/'
|
||||
f'\n*** 分析完成! Run ID: #{run.display_id}'
|
||||
f' 查看 Django: http://127.0.0.1:8000/runs/{run.display_id}/'
|
||||
))
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# Generated by Django 4.2.30 on 2026-07-20 05:11
|
||||
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('tianxuan_analysis', '0002_analysisrun_progress_msg_analysisrun_progress_pct_and_more'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AddField(
|
||||
model_name='analysisrun',
|
||||
name='sqlite_table',
|
||||
field=models.CharField(blank=True, default='', help_text='SQLite table name for persisted raw data', max_length=256),
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Migration 0004 — create raw SQLite table ``tls_ref_data``.
|
||||
|
||||
Replaces the ``TlsDB.csv`` reference file with a proper SQLite table.
|
||||
The table stores TLS field-code metadata: human-readable meaning and
|
||||
data type for each field code used in TLS flow data.
|
||||
"""
|
||||
|
||||
from django.db import migrations
|
||||
|
||||
# Raw SQL for creating the tls_ref_data table
|
||||
CREATE_TLS_REF_DATA = """
|
||||
CREATE TABLE IF NOT EXISTS tls_ref_data (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
field_code TEXT NOT NULL UNIQUE,
|
||||
meaning TEXT NOT NULL DEFAULT '',
|
||||
data_type TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
"""
|
||||
|
||||
DROP_TLS_REF_DATA = """
|
||||
DROP TABLE IF EXISTS tls_ref_data;
|
||||
"""
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
"""Create ``tls_ref_data`` table via raw SQL — no Django model needed."""
|
||||
|
||||
dependencies = [
|
||||
('tianxuan_analysis', '0003_analysisrun_sqlite_table'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.RunSQL(
|
||||
sql=CREATE_TLS_REF_DATA,
|
||||
reverse_sql=DROP_TLS_REF_DATA,
|
||||
state_operations=None,
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
# Generated by Django 4.2.30 on 2026-07-20 05:20
|
||||
|
||||
from django.db import migrations
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('tianxuan_analysis', '0004_tls_ref_data'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.RemoveField(
|
||||
model_name='analysisrun',
|
||||
name='entity_column',
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
# Generated by Django 4.2.30 on 2026-07-20 05:25
|
||||
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('tianxuan_analysis', '0005_remove_entity_column'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AddField(
|
||||
model_name='analysisrun',
|
||||
name='display_id',
|
||||
field=models.IntegerField(blank=True, help_text='User-facing numeric ID (recycles gaps)', null=True, unique=True),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name='analysisrun',
|
||||
name='llm_thinking',
|
||||
field=models.TextField(blank=True, default='', help_text='Accumulated LLM reasoning / thinking text'),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name='analysisrun',
|
||||
name='run_type',
|
||||
field=models.CharField(choices=[('upload', '手动上传'), ('manual', '手动分析'), ('auto', 'LLM自动')], default='upload', max_length=16),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name='analysisrun',
|
||||
name='tool_calls_json',
|
||||
field=models.JSONField(blank=True, default=list, help_text='Structured tool call list: [{step, name, input, output}, ...]'),
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,20 @@
|
||||
# Generated by Django 4.2.30 on 2026-07-20 05:26
|
||||
|
||||
from django.db import migrations
|
||||
from django.db.models import F
|
||||
|
||||
|
||||
def populate_display_id(apps, schema_editor):
|
||||
AnalysisRun = apps.get_model('tianxuan_analysis', 'AnalysisRun')
|
||||
AnalysisRun.objects.filter(display_id__isnull=True).update(display_id=F('id'))
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('tianxuan_analysis', '0006_analysisrun_display_id_analysisrun_llm_thinking_and_more'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.RunPython(populate_display_id, migrations.RunPython.noop),
|
||||
]
|
||||
+35
-4
@@ -1,4 +1,4 @@
|
||||
from django.db import models
|
||||
from django.db import models, connection
|
||||
|
||||
|
||||
class AnalysisRun(models.Model):
|
||||
@@ -14,9 +14,18 @@ class AnalysisRun(models.Model):
|
||||
('failed', 'Failed'),
|
||||
]
|
||||
|
||||
RUN_TYPE_CHOICES = [
|
||||
('upload', '手动上传'),
|
||||
('manual', '手动分析'),
|
||||
('auto', 'LLM自动'),
|
||||
]
|
||||
|
||||
display_id = models.IntegerField(unique=True, null=True, blank=True,
|
||||
help_text="User-facing numeric ID (recycles gaps)")
|
||||
csv_glob = models.CharField(max_length=1024, help_text="Glob pattern for CSV files")
|
||||
config = models.JSONField(default=dict, blank=True, help_text="YAML config snapshot")
|
||||
status = models.CharField(max_length=32, choices=STATUS_CHOICES, default='pending')
|
||||
run_type = models.CharField(max_length=16, choices=RUN_TYPE_CHOICES, default='upload')
|
||||
created_at = models.DateTimeField(auto_now_add=True)
|
||||
updated_at = models.DateTimeField(auto_now=True)
|
||||
error_message = models.TextField(blank=True, default='')
|
||||
@@ -28,14 +37,36 @@ class AnalysisRun(models.Model):
|
||||
total_flows = models.IntegerField(null=True, blank=True)
|
||||
entity_count = models.IntegerField(null=True, blank=True)
|
||||
cluster_count = models.IntegerField(null=True, blank=True)
|
||||
entity_column = models.CharField(max_length=256, blank=True, default='',
|
||||
help_text="Auto-detected entity column name")
|
||||
sqlite_table = models.CharField(max_length=256, blank=True, default='',
|
||||
help_text="SQLite table name for persisted raw data")
|
||||
|
||||
# LLM auto-analysis fields
|
||||
llm_thinking = models.TextField(blank=True, default='',
|
||||
help_text="Accumulated LLM reasoning / thinking text")
|
||||
tool_calls_json = models.JSONField(default=list, blank=True,
|
||||
help_text="Structured tool call list: [{step, name, input, output}, ...]")
|
||||
|
||||
class Meta:
|
||||
ordering = ['-created_at']
|
||||
|
||||
def __str__(self):
|
||||
return f"Run #{self.id} ({self.created_at:%Y-%m-%d %H:%M}) - {self.status}"
|
||||
return f"Run #{self.display_id or self.id} ({self.created_at:%Y-%m-%d %H:%M}) - {self.status}"
|
||||
|
||||
def save(self, *args, **kwargs):
|
||||
if self._state.adding and self.display_id is None:
|
||||
table = self._meta.db_table
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(f"""
|
||||
SELECT COALESCE(MIN(t1.display_id) + 1, 1)
|
||||
FROM {table} t1
|
||||
WHERE t1.display_id IS NOT NULL
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} t2
|
||||
WHERE t2.display_id = t1.display_id + 1
|
||||
)
|
||||
""")
|
||||
self.display_id = cursor.fetchone()[0]
|
||||
super().save(*args, **kwargs)
|
||||
|
||||
|
||||
class ClusterResult(models.Model):
|
||||
|
||||
@@ -106,10 +106,14 @@ class SessionStore:
|
||||
)
|
||||
|
||||
def restore_from_disk(self) -> int:
|
||||
"""Re-register datasets whose CSV files still exist on disk.
|
||||
"""Re-register datasets from persisted metadata.
|
||||
|
||||
Uses ``csv_glob`` saved in metadata to re-load each dataset via
|
||||
:func:`analysis.data_loader.load_csv_directory`.
|
||||
Attempts two strategies in order:
|
||||
1. **SQLite table** — when ``metadata.sqlite_table`` exists (fast,
|
||||
always available, no CSV dependency).
|
||||
2. **CSV glob** — falls back to ``load_csv_directory`` for datasets
|
||||
that were stored before the SQLite migration or whose SQLite table
|
||||
was dropped.
|
||||
|
||||
Returns the number of successfully restored datasets.
|
||||
"""
|
||||
@@ -121,18 +125,38 @@ class SessionStore:
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return 0
|
||||
|
||||
# Lazy import to avoid circular dependency at module level
|
||||
from analysis.data_loader import load_csv_directory # fmt: skip
|
||||
# Lazy imports to avoid circular dependency at module level
|
||||
from analysis.data_loader import load_csv_directory, load_from_db # fmt: skip
|
||||
|
||||
restored = 0
|
||||
for ds_id, info in data.items():
|
||||
meta = dict(info.get('metadata', {}))
|
||||
schema = dict(info.get('schema', {}))
|
||||
restored_ok = False
|
||||
|
||||
# Strategy 1: SQLite table (fast path)
|
||||
sqlite_table = meta.get('sqlite_table')
|
||||
if sqlite_table:
|
||||
try:
|
||||
lf = load_from_db(sqlite_table)
|
||||
if lf is not None:
|
||||
self.store_dataset(ds_id, lf, schema=schema, metadata=meta)
|
||||
restored += 1
|
||||
restored_ok = True
|
||||
except Exception:
|
||||
pass # fall through to CSV restore
|
||||
|
||||
if restored_ok:
|
||||
continue
|
||||
|
||||
# Strategy 2: CSV glob (legacy / fallback)
|
||||
csv_glob = info.get('csv_glob')
|
||||
if not csv_glob:
|
||||
continue
|
||||
try:
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
|
||||
meta = dict(info.get('metadata', {}))
|
||||
self.store_dataset(ds_id, lf, schema=schema, metadata=meta)
|
||||
lf, schema_csv, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
|
||||
# Merge schema from CSV (may be more accurate than persisted schema)
|
||||
self.store_dataset(ds_id, lf, schema=schema_csv, metadata=meta)
|
||||
restored += 1
|
||||
except Exception:
|
||||
# File(s) no longer exist or schema changed — skip gracefully
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""SQLite-backed TLS reference data access — replaces ``TlsDB.csv``.
|
||||
|
||||
Provides functions to query the ``tls_ref_data`` table (populated by
|
||||
the ``import_tlsdb`` management command) instead of reading the
|
||||
now-deprecated ``TlsDB.csv`` file.
|
||||
|
||||
Usage::
|
||||
|
||||
from analysis.tls_ref import get_tls_ref_data, get_tls_field, get_tls_ref_dict
|
||||
|
||||
# All rows as list of dicts
|
||||
all_fields = get_tls_ref_data()
|
||||
|
||||
# Single field
|
||||
field = get_tls_field('0ver')
|
||||
# -> {'field_code': '0ver', 'meaning': 'TLS版本', 'data_type': '2字节'}
|
||||
|
||||
# Dict lookup (cached in memory after first call)
|
||||
lookup = get_tls_ref_dict()
|
||||
# -> {'0ver': {'meaning': 'TLS版本', 'data_type': '2字节'}, ...}
|
||||
"""
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from django.db import connection
|
||||
|
||||
|
||||
def get_tls_ref_data() -> list[dict[str, str]]:
|
||||
"""Return all rows from ``tls_ref_data`` as a list of dicts.
|
||||
|
||||
Returns:
|
||||
Each dict has keys ``field_code``, ``meaning``, ``data_type``.
|
||||
Returns an empty list if the table does not exist (e.g. before
|
||||
migration).
|
||||
"""
|
||||
if not _table_exists():
|
||||
return []
|
||||
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
'SELECT field_code, meaning, data_type FROM tls_ref_data ORDER BY id'
|
||||
)
|
||||
columns = [col[0] for col in cursor.description]
|
||||
return [dict(zip(columns, row)) for row in cursor.fetchall()]
|
||||
|
||||
|
||||
def get_tls_field(field_code: str) -> Optional[dict[str, str]]:
|
||||
"""Return a single TLS field definition by its code.
|
||||
|
||||
Args:
|
||||
field_code: e.g. ``'0ver'``, ``':ips'``, ``'cipher-suite'``.
|
||||
|
||||
Returns:
|
||||
Dict with keys ``field_code``, ``meaning``, ``data_type``, or
|
||||
``None`` if not found.
|
||||
"""
|
||||
if not _table_exists():
|
||||
return None
|
||||
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
'SELECT field_code, meaning, data_type FROM tls_ref_data WHERE field_code = %s',
|
||||
[field_code],
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return {
|
||||
'field_code': row[0],
|
||||
'meaning': row[1],
|
||||
'data_type': row[2],
|
||||
}
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_tls_ref_dict() -> dict[str, dict[str, str]]:
|
||||
"""Return a dict mapping ``field_code -> {meaning, data_type}``.
|
||||
|
||||
The result is cached in memory after the first call. Call
|
||||
``get_tls_ref_dict.cache_clear()`` to invalidate if the
|
||||
underlying table is modified at runtime.
|
||||
"""
|
||||
rows = get_tls_ref_data()
|
||||
return {r['field_code']: {'meaning': r['meaning'], 'data_type': r['data_type']} for r in rows}
|
||||
|
||||
|
||||
def _table_exists() -> bool:
|
||||
"""Check if the ``tls_ref_data`` table exists in the database."""
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='tls_ref_data'"
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
+680
-336
File diff suppressed because it is too large
Load Diff
+6
-6
@@ -11,15 +11,15 @@ urlpatterns = [
|
||||
path('analyze/auto/', views.auto_page, name='auto'),
|
||||
path('analyze/llm/', views.run_llm_analysis_view, name='run_llm_analysis'),
|
||||
path('runs/', views.run_list, name='run_list'),
|
||||
path('runs/<int:run_id>/', views.run_detail, name='run_detail'),
|
||||
path('runs/<int:run_id>/retry/', views.retry_analysis, name='retry_analysis'),
|
||||
path('runs/<int:run_id>/status/', views.run_status_api, name='run_status'),
|
||||
path('clusters/<int:run_id>/', views.cluster_overview, name='cluster_overview'),
|
||||
path('clusters/<int:run_id>/<int:cluster_label>/', views.cluster_detail, name='cluster_detail'),
|
||||
path('runs/<int:display_id>/', views.run_detail, name='run_detail'),
|
||||
path('runs/<int:display_id>/status/', views.run_status_api, name='run_status'),
|
||||
path('clusters/<int:display_id>/', views.cluster_overview, name='cluster_overview'),
|
||||
path('clusters/<int:display_id>/<int:cluster_label>/', views.cluster_detail, name='cluster_detail'),
|
||||
path('entities/<int:entity_id>/', views.entity_profile, name='entity_profile'),
|
||||
path('config/', views.config_view, name='config'),
|
||||
path('api/llm/test/', views.llm_test, name='llm_test'),
|
||||
path('runs/<int:run_id>/delete/', views.delete_upload, name='delete_upload'),
|
||||
path('runs/<int:display_id>/retry/', views.retry_run, name='retry_run'),
|
||||
path('runs/<int:display_id>/delete/', views.delete_upload, name='delete_upload'),
|
||||
path('logs/', views.log_viewer, name='log_viewer'),
|
||||
path('globe/', views.globe_view, name='globe'),
|
||||
path('tools/run/', views.tool_lab_run, name='tool_lab_run'),
|
||||
|
||||
+274
-207
@@ -21,21 +21,28 @@ def dashboard(request):
|
||||
|
||||
|
||||
def run_list(request):
|
||||
"""List all analysis runs with pagination."""
|
||||
"""List all analysis runs with pagination and optional ?type= filter."""
|
||||
run_type_filter = request.GET.get('type')
|
||||
page = int(request.GET.get('page', 1))
|
||||
per_page = 20
|
||||
total = AnalysisRun.objects.count()
|
||||
runs = AnalysisRun.objects.all()[(page - 1) * per_page: page * per_page]
|
||||
|
||||
qs = AnalysisRun.objects.all()
|
||||
if run_type_filter in dict(AnalysisRun.RUN_TYPE_CHOICES):
|
||||
qs = qs.filter(run_type=run_type_filter)
|
||||
|
||||
total = qs.count()
|
||||
runs = qs[(page - 1) * per_page: page * per_page]
|
||||
return render(request, 'analysis/run_list.html', {
|
||||
'runs': runs,
|
||||
'page': page,
|
||||
'total_pages': (total + per_page - 1) // per_page if total else 1,
|
||||
'current_type': run_type_filter or '',
|
||||
})
|
||||
|
||||
|
||||
def run_detail(request, run_id):
|
||||
def run_detail(request, display_id):
|
||||
"""Details for a single analysis run."""
|
||||
run = get_object_or_404(AnalysisRun, id=run_id)
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
clusters = run.clusters.all().order_by('-size')
|
||||
return render(request, 'analysis/run_detail.html', {
|
||||
'run': run,
|
||||
@@ -43,74 +50,6 @@ def run_detail(request, run_id):
|
||||
})
|
||||
|
||||
|
||||
@csrf_exempt
|
||||
def retry_analysis(request, run_id):
|
||||
"""Retry a failed analysis run. Creates a new run with same parameters."""
|
||||
if request.method != 'POST':
|
||||
return JsonResponse({'error': 'POST required'}, status=405)
|
||||
|
||||
original = get_object_or_404(AnalysisRun, id=run_id)
|
||||
if original.status != 'failed':
|
||||
return JsonResponse({'error': '只能重试失败的分析'}, status=400)
|
||||
|
||||
# Create new run copying original parameters
|
||||
new_run = AnalysisRun.objects.create(
|
||||
csv_glob=original.csv_glob,
|
||||
entity_column=original.entity_column,
|
||||
status='pending',
|
||||
total_flows=original.total_flows,
|
||||
)
|
||||
|
||||
def _retry_fn(run, ctx):
|
||||
from analysis.data_loader import load_csv_directory
|
||||
from analysis.session_store import SessionStore
|
||||
from analysis.entity_detector import detect_entity_column
|
||||
from analysis.entity_aggregator import aggregate_by_entity
|
||||
|
||||
store = SessionStore()
|
||||
try:
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(run.csv_glob)
|
||||
ds_id = f'retry_{run.id}'
|
||||
store.store_dataset(ds_id, lf, schema=schema, metadata={
|
||||
'row_count': row_count, 'file_count': file_count, 'csv_glob': run.csv_glob,
|
||||
})
|
||||
run.total_flows = row_count
|
||||
run.save(update_fields=['total_flows'])
|
||||
|
||||
# Detect entity column
|
||||
result = detect_entity_column(ds_id)
|
||||
entity_col = result.get('recommended')
|
||||
if not entity_col:
|
||||
raise Exception('Could not detect entity column')
|
||||
|
||||
run.entity_column = entity_col
|
||||
run.save(update_fields=['entity_column'])
|
||||
|
||||
# Aggregate
|
||||
agg_lf, feature_columns = aggregate_by_entity(lf, entity_col, schema)
|
||||
df_agg = agg_lf.collect(streaming=True)
|
||||
entity_ds_id = f'entity_{run.id}'
|
||||
store.store_dataset(entity_ds_id, agg_lf, schema=dict(zip(feature_columns, ['']*len(feature_columns))))
|
||||
run.entity_count = len(df_agg)
|
||||
run.save(update_fields=['entity_count'])
|
||||
|
||||
# Run clustering pipeline (same defaults as manual analysis)
|
||||
_run_clustering_pipeline(
|
||||
run=run, store=store, entity_ds_id=entity_ds_id,
|
||||
feature_columns=None, algorithm='hdbscan',
|
||||
min_cluster_size=5, run_pca=True,
|
||||
)
|
||||
except Exception as e:
|
||||
run.status = 'failed'
|
||||
run.error_message = str(e)
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
|
||||
t = threading.Thread(target=_run_pipeline_worker, args=(new_run.id, _retry_fn), daemon=True)
|
||||
t.start()
|
||||
|
||||
return JsonResponse({'status': 'started', 'run_id': new_run.id, 'redirect': f'/runs/{new_run.id}/'})
|
||||
|
||||
|
||||
def _extract_lat(val):
|
||||
"""Try to extract a latitude from *val* (number or str)."""
|
||||
if val is None:
|
||||
@@ -133,9 +72,9 @@ def _extract_lon(val):
|
||||
return None
|
||||
|
||||
|
||||
def cluster_overview(request, run_id):
|
||||
def cluster_overview(request, display_id):
|
||||
"""Overview of all clusters for a run, with PCA scatter data and geo scatter data."""
|
||||
run = get_object_or_404(AnalysisRun, id=run_id)
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
clusters = run.clusters.exclude(cluster_label=-1).order_by('-size').prefetch_related('features')
|
||||
noise = run.clusters.filter(cluster_label=-1).first()
|
||||
|
||||
@@ -191,9 +130,9 @@ def cluster_overview(request, run_id):
|
||||
})
|
||||
|
||||
|
||||
def cluster_detail(request, run_id, cluster_label):
|
||||
def cluster_detail(request, display_id, cluster_label):
|
||||
"""Detailed view of a single cluster with entity list."""
|
||||
run = get_object_or_404(AnalysisRun, id=run_id)
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
cluster = get_object_or_404(ClusterResult, run=run, cluster_label=cluster_label)
|
||||
features = cluster.features.all().order_by('-distinguishing_score')[:50]
|
||||
entities = cluster.entities.all()[:50]
|
||||
@@ -370,8 +309,6 @@ def upload_csv(request):
|
||||
files = request.FILES.getlist('files')
|
||||
if not files:
|
||||
return JsonResponse({'error': 'No files uploaded'})
|
||||
if len(files) > 10000:
|
||||
return JsonResponse({'error': '一次最多上传 10000 个文件'})
|
||||
for f in files:
|
||||
if f.size > 500 * 1024 * 1024:
|
||||
return JsonResponse({'warning': f'文件 {f.name} 超过 500MB，建议分割后上传'})
|
||||
@@ -398,21 +335,40 @@ def upload_csv(request):
|
||||
run = AnalysisRun.objects.create(
|
||||
csv_glob=str(upload_dir / '*.csv'),
|
||||
status='loading',
|
||||
run_type='upload',
|
||||
)
|
||||
t = threading.Thread(target=_background_process, args=(run.id, upload_dir), daemon=True)
|
||||
t.start()
|
||||
return JsonResponse({'run_id': run.id, 'files': saved, 'status': 'loading'})
|
||||
return JsonResponse({'run_id': run.display_id, 'files': saved, 'status': 'loading'})
|
||||
|
||||
|
||||
def delete_upload(request, run_id):
|
||||
"""Delete uploaded files and AnalysisRun."""
|
||||
run = get_object_or_404(AnalysisRun, id=run_id)
|
||||
def delete_upload(request, display_id):
|
||||
"""Delete uploaded files, SQLite data table, and AnalysisRun.
|
||||
|
||||
Handles all run statuses (loading, profiling, clustering, completed, failed).
|
||||
Gracefully tolerates null csv_glob and missing/errored SQLite tables.
|
||||
"""
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
import logging
|
||||
import shutil
|
||||
upload_dir = None
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
# 1. Drop the SQLite data table if one exists
|
||||
if run.sqlite_table:
|
||||
try:
|
||||
from analysis.data_loader import drop_sqlite_table # fmt: skip
|
||||
drop_sqlite_table(run.sqlite_table)
|
||||
except Exception as exc:
|
||||
_logger.warning('[DELETE] failed to drop SQLite table %s: %s', run.sqlite_table, exc)
|
||||
|
||||
# 2. Remove the upload directory if csv_glob is set
|
||||
if run.csv_glob:
|
||||
p = Path(run.csv_glob).parent
|
||||
if p.exists():
|
||||
shutil.rmtree(p, ignore_errors=True)
|
||||
|
||||
# 3. Delete the ORM record
|
||||
run.delete()
|
||||
return JsonResponse({'deleted': True})
|
||||
|
||||
@@ -422,14 +378,11 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _background_process(run_id, upload_dir):
|
||||
"""Background: load → detect → aggregate → ready."""
|
||||
"""Background: stream-load CSVs → save to SQLite → ready for analysis."""
|
||||
import django
|
||||
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
||||
django.setup()
|
||||
from analysis.models import AnalysisRun
|
||||
from analysis.data_loader import load_csv_directory
|
||||
from analysis.entity_detector import detect_entity_column
|
||||
from analysis.entity_aggregator import aggregate_by_entity
|
||||
from analysis.session_store import SessionStore
|
||||
|
||||
run = AnalysisRun.objects.get(id=run_id)
|
||||
@@ -444,56 +397,88 @@ def _background_process(run_id, upload_dir):
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
|
||||
run.status = 'loading'; run.progress_pct = 10; run.progress_msg = '正在加载数据...'
|
||||
import glob
|
||||
import polars as pl
|
||||
csv_paths = sorted(glob.glob(str(upload_dir / '*.csv')))
|
||||
if not csv_paths:
|
||||
msg = '没有找到 CSV 文件'
|
||||
logger.error('[BACKGROUND] %s (run_id=%s)', msg, run_id)
|
||||
run.status = 'failed'
|
||||
run.error_message = msg
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
|
||||
total_files = len(csv_paths)
|
||||
# Adaptive batch size: aim for ~20 batches total, cap at 500 files per batch
|
||||
batch_size = max(1, min(500, total_files // 20 or 1))
|
||||
total_batches = (total_files + batch_size - 1) // batch_size
|
||||
|
||||
run.status = 'loading'; run.progress_pct = 10
|
||||
run.progress_msg = f'正在流式加载 {total_files} 个文件（共 {total_batches} 批）...'
|
||||
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
||||
lf, schema, row_count, file_count, mem = load_csv_directory(
|
||||
str(upload_dir / '*.csv'), schema_strict=False
|
||||
)
|
||||
|
||||
from analysis.data_loader import save_to_db, drop_sqlite_table, load_from_db
|
||||
|
||||
table_name = f'_data_{run_id}'
|
||||
completed_batches = 0
|
||||
schema = {}
|
||||
total_rows = 0
|
||||
|
||||
# ── Phase 1: Stream CSVs into SQLite in batches ─────────────────
|
||||
for i in range(0, total_files, batch_size):
|
||||
batch_files = csv_paths[i:i + batch_size]
|
||||
mode = 'replace' if i == 0 else 'append'
|
||||
|
||||
batch_lf = pl.scan_csv(batch_files, low_memory=True)
|
||||
|
||||
# Capture schema from first batch (efficient — header-only scan)
|
||||
if not schema:
|
||||
raw_schema = batch_lf.collect_schema()
|
||||
schema = {name: str(dtype) for name, dtype in raw_schema.items()}
|
||||
|
||||
save_to_db(run.id, batch_lf, mode=mode)
|
||||
|
||||
completed_batches += 1
|
||||
run.progress_pct = int((completed_batches / total_batches) * 80 + 10)
|
||||
run.progress_msg = f'正在流式加载... ({completed_batches}/{total_batches} 批)'
|
||||
run.save(update_fields=['progress_pct', 'progress_msg'])
|
||||
|
||||
# ── Phase 2: Load from SQLite for analysis ─────────────────────
|
||||
lf = load_from_db(table_name)
|
||||
if lf is None:
|
||||
raise RuntimeError(f'SQLite 表 {table_name} 加载失败')
|
||||
|
||||
# Get row count
|
||||
df_count = lf.select(pl.len()).collect(streaming=True)
|
||||
total_rows = df_count[0, 0] if df_count.height > 0 else 0
|
||||
|
||||
run.sqlite_table = table_name
|
||||
run.total_flows = total_rows
|
||||
run.save(update_fields=['sqlite_table', 'total_flows'])
|
||||
|
||||
ds_id = f'upload_{run_id}'
|
||||
store.store_dataset(ds_id, lf, schema=schema, metadata={
|
||||
'row_count': row_count, 'file_count': file_count, 'upload_dir': str(upload_dir),
|
||||
'row_count': total_rows, 'file_count': total_files, 'upload_dir': str(upload_dir),
|
||||
'csv_glob': str(upload_dir / '*.csv'),
|
||||
'sqlite_table': table_name,
|
||||
})
|
||||
run.total_flows = row_count; run.save(update_fields=['total_flows'])
|
||||
|
||||
run.status = 'profiling'; run.progress_pct = 30; run.progress_msg = '正在分析数据特征...'
|
||||
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
||||
entity_result = detect_entity_column(ds_id)
|
||||
entity_col = entity_result.get('recommended', '')
|
||||
entity_cols_multi = entity_result.get('recommended_multi', [])
|
||||
# Store as comma-separated (multi) or single-column
|
||||
if entity_cols_multi:
|
||||
entity_cols = entity_cols_multi[:2] # top 2
|
||||
run.entity_column = ','.join(entity_cols)
|
||||
else:
|
||||
entity_cols = [entity_col] if entity_col else []
|
||||
run.entity_column = entity_col
|
||||
run.save(update_fields=['entity_column'])
|
||||
|
||||
if entity_cols:
|
||||
run.status = 'aggregating'; run.progress_pct = 50; run.progress_msg = '正在构建实体画像...'
|
||||
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
||||
agg_lf, features = aggregate_by_entity(lf, entity_cols, schema)
|
||||
df_agg = agg_lf.collect(streaming=True)
|
||||
entity_ds_id = f'entity_{run_id}'
|
||||
store.store_dataset(entity_ds_id, agg_lf,
|
||||
schema=dict(zip(features, [''] * len(features))),
|
||||
metadata={'row_count': len(df_agg), 'entity_column': ','.join(entity_cols),
|
||||
'csv_glob': str(upload_dir / '*.csv')})
|
||||
run.entity_count = len(df_agg)
|
||||
run.entity_column = ','.join(entity_cols)
|
||||
run.save(update_fields=['entity_column', 'entity_count'])
|
||||
|
||||
# ── Phase 3: Data ready for filtering/clustering via MCP tools ──
|
||||
run.status = 'ready'; run.progress_pct = 60; run.progress_msg = '数据准备就绪'
|
||||
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
||||
except Exception:
|
||||
tb = traceback.format_exc()
|
||||
logger.error(tb)
|
||||
# Drop partial SQLite table on failure (idempotent)
|
||||
try:
|
||||
from analysis.data_loader import drop_sqlite_table # fmt: skip
|
||||
drop_sqlite_table(f'_data_{run_id}')
|
||||
except Exception:
|
||||
pass
|
||||
run.status = 'failed'
|
||||
run.run_log += f'\n[ERROR] {tb}'
|
||||
run.save(update_fields=['run_log'])
|
||||
run.error_message = tb
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
run.run_log += f'\n[ERROR] {tb}'
|
||||
run.save(update_fields=['status', 'error_message', 'run_log'])
|
||||
|
||||
|
||||
# ── Manual analysis page ───────────────────────────────────────────────
|
||||
@@ -547,7 +532,16 @@ def _run_pipeline_worker(run_id, analysis_fn, **ctx):
|
||||
os.environ['DJANGO_ALLOW_ASYNC_UNSAFE'] = 'true'
|
||||
django.setup()
|
||||
|
||||
run = AnalysisRun.objects.get(id=run_id)
|
||||
from analysis.db_utils import retry_on_lock
|
||||
from django.db import close_old_connections
|
||||
|
||||
close_old_connections()
|
||||
|
||||
@retry_on_lock()
|
||||
def _get_run():
|
||||
return AnalysisRun.objects.get(id=run_id)
|
||||
|
||||
run = _get_run()
|
||||
try:
|
||||
analysis_fn(run, ctx)
|
||||
except Exception:
|
||||
@@ -557,7 +551,7 @@ def _run_pipeline_worker(run_id, analysis_fn, **ctx):
|
||||
run.status = 'failed'
|
||||
run.error_message = tb
|
||||
run.run_log += f'\n[FATAL] {tb}'
|
||||
run.save(update_fields=['status', 'error_message', 'run_log'])
|
||||
retry_on_lock()(lambda: run.save(update_fields=['status', 'error_message', 'run_log']))()
|
||||
|
||||
|
||||
# ── Shared clustering pipeline ─────────────────────────────────────────
|
||||
@@ -619,16 +613,13 @@ def _run_clustering_pipeline(run, store, entity_ds_id, feature_columns, algorith
|
||||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64',
|
||||
'UInt8', 'UInt16', 'UInt32', 'UInt64',
|
||||
'Float32', 'Float64'}
|
||||
entity_cols_set = set()
|
||||
if hasattr(run, 'entity_column') and run.entity_column:
|
||||
entity_cols_set = {c.strip() for c in run.entity_column.split(',')}
|
||||
|
||||
if feature_columns:
|
||||
feature_cols = [c for c in feature_columns if c in schema and c not in entity_cols_set]
|
||||
feature_cols = [c for c in feature_columns if c in schema]
|
||||
else:
|
||||
feature_cols = [
|
||||
c for c in schema.keys()
|
||||
if not c.startswith('_') and c not in entity_cols_set
|
||||
if not c.startswith('_')
|
||||
and schema[c].split('(')[0].strip() in numeric_types
|
||||
][:10]
|
||||
|
||||
@@ -753,11 +744,10 @@ def _run_clustering_pipeline(run, store, entity_ds_id, feature_columns, algorith
|
||||
logger.error(tb)
|
||||
try:
|
||||
run.status = 'failed'
|
||||
run.error_message = tb
|
||||
run.progress_msg = f'失败: {tb}'
|
||||
run.run_log += f'\n[ERROR] {tb}'
|
||||
run.save(update_fields=['run_log'])
|
||||
run.error_message = tb
|
||||
run.save(update_fields=['status', 'progress_msg', 'error_message'])
|
||||
run.save(update_fields=['status', 'error_message', 'progress_msg', 'run_log'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -766,14 +756,16 @@ def _run_clustering_pipeline(run, store, entity_ds_id, feature_columns, algorith
|
||||
|
||||
@csrf_exempt
|
||||
def run_analysis(request):
|
||||
"""AJAX POST: run analysis pipeline and return result page URL."""
|
||||
"""AJAX POST: run clustering pipeline on the upload dataset."""
|
||||
if request.method != 'POST':
|
||||
return JsonResponse({'error': 'POST required'}, status=405)
|
||||
import json
|
||||
body = json.loads(request.body)
|
||||
run_id = body.get('run_id')
|
||||
entity_col = body.get('entity_column', '')
|
||||
entity_columns = body.get('entity_columns')
|
||||
display_id = body.get('run_id')
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
pk = run.id
|
||||
run.run_type = 'manual'
|
||||
run.save(update_fields=['run_type'])
|
||||
feature_columns = body.get('feature_columns')
|
||||
algorithm = body.get('algorithm', 'hdbscan')
|
||||
min_cluster_size = int(body.get('min_cluster_size', 5))
|
||||
@@ -781,53 +773,32 @@ def run_analysis(request):
|
||||
|
||||
def _analysis_fn(run, ctx):
|
||||
from analysis.session_store import SessionStore
|
||||
from analysis.entity_aggregator import aggregate_by_entity
|
||||
|
||||
entity_columns = ctx['entity_columns']
|
||||
feature_columns = ctx['feature_columns']
|
||||
algorithm = ctx['algorithm']
|
||||
min_cluster_size = ctx['min_cluster_size']
|
||||
head = ctx.get('head')
|
||||
pk = ctx['pk']
|
||||
|
||||
store = SessionStore()
|
||||
try:
|
||||
if entity_columns:
|
||||
upload_ds_id = f'upload_{run_id}'
|
||||
upload_entry = store.get_dataset(upload_ds_id)
|
||||
if upload_entry is None:
|
||||
run.error_message = '请先上传并等待预处理完成'
|
||||
run.status = 'failed'
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
upload_lf = upload_entry['lazyframe']
|
||||
upload_schema = upload_entry.get('schema', {})
|
||||
agg_lf, agg_features = aggregate_by_entity(upload_lf, entity_columns, upload_schema)
|
||||
df_agg = agg_lf.collect(streaming=True)
|
||||
agg_schema = {n: str(d) for n, d in zip(df_agg.columns, [df_agg[c].dtype for c in df_agg.columns])}
|
||||
entity_ds_id = f'entity_{run_id}_custom'
|
||||
store.store_dataset(
|
||||
dataset_id=entity_ds_id,
|
||||
lazyframe=df_agg.lazy(),
|
||||
schema=agg_schema,
|
||||
metadata={'row_count': len(df_agg), 'entity_columns': entity_columns,
|
||||
'csv_glob': run.csv_glob},
|
||||
)
|
||||
run.entity_column = ','.join(entity_columns)
|
||||
run.entity_count = len(df_agg)
|
||||
run.save(update_fields=['entity_column', 'entity_count'])
|
||||
else:
|
||||
entity_ds_id = f'entity_{run_id}'
|
||||
entry = store.get_dataset(entity_ds_id)
|
||||
if entry is None:
|
||||
run.error_message = '请先上传并等待预处理完成'
|
||||
run.status = 'failed'
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
upload_ds_id = f'upload_{pk}'
|
||||
entry = store.get_dataset(upload_ds_id)
|
||||
if entry is None:
|
||||
# Fallback: try entity dataset (from old pipeline)
|
||||
entry = store.get_dataset(f'entity_{pk}')
|
||||
if entry is None:
|
||||
run.error_message = '请先上传并等待预处理完成'
|
||||
run.status = 'failed'
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
|
||||
ds_id = upload_ds_id if store.get_dataset(upload_ds_id) else f'entity_{pk}'
|
||||
|
||||
# Use the unified clustering pipeline (clustering → extraction → PCA)
|
||||
from analysis.views import _run_clustering_pipeline
|
||||
_run_clustering_pipeline(
|
||||
run=run, store=store, entity_ds_id=entity_ds_id,
|
||||
run=run, store=store, entity_ds_id=ds_id,
|
||||
feature_columns=feature_columns,
|
||||
algorithm=algorithm, min_cluster_size=min_cluster_size,
|
||||
run_pca=True, head=head,
|
||||
@@ -836,22 +807,20 @@ def run_analysis(request):
|
||||
tb = traceback.format_exc()
|
||||
logger.error(tb)
|
||||
run.status = 'failed'
|
||||
run.error_message = tb
|
||||
run.progress_msg = f'失败: {tb}'
|
||||
run.run_log += f'\n[ERROR] {tb}'
|
||||
run.save(update_fields=['run_log'])
|
||||
run.error_message = tb
|
||||
run.save(update_fields=['status', 'progress_msg', 'error_message'])
|
||||
run.save(update_fields=['status', 'error_message', 'progress_msg', 'run_log'])
|
||||
|
||||
t = threading.Thread(
|
||||
target=_run_pipeline_worker,
|
||||
args=(run_id, _analysis_fn),
|
||||
kwargs=dict(entity_col=entity_col, entity_columns=entity_columns,
|
||||
feature_columns=feature_columns, algorithm=algorithm,
|
||||
min_cluster_size=min_cluster_size, head=head),
|
||||
kwargs=dict(feature_columns=feature_columns, algorithm=algorithm,
|
||||
min_cluster_size=min_cluster_size, head=head, pk=pk),
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
return JsonResponse({'status': 'started', 'redirect': f'/runs/{run_id}/'})
|
||||
return JsonResponse({'status': 'started', 'redirect': f'/runs/{run.display_id}/'})
|
||||
|
||||
|
||||
|
||||
@@ -866,42 +835,68 @@ def run_llm_analysis_view(request):
|
||||
return JsonResponse({'error': 'POST required'}, status=405)
|
||||
import json
|
||||
body = json.loads(request.body)
|
||||
run_id = body.get('run_id')
|
||||
entity_col = body.get('entity_column', '')
|
||||
display_id = body.get('run_id')
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
pk = run.id
|
||||
run.run_type = 'auto'
|
||||
run.save(update_fields=['run_type'])
|
||||
head = body.get('head')
|
||||
|
||||
from analysis.session_store import SessionStore
|
||||
store = SessionStore()
|
||||
upload_ds_id = f'upload_{run_id}'
|
||||
upload_ds_id = f'upload_{pk}'
|
||||
entry = store.get_dataset(upload_ds_id)
|
||||
|
||||
if entry is None:
|
||||
# Try entity dataset as fallback
|
||||
entity_ds_id = f'entity_{run_id}'
|
||||
entity_ds_id = f'entity_{pk}'
|
||||
entry = store.get_dataset(entity_ds_id)
|
||||
if entry is None:
|
||||
return JsonResponse({'error': '数据集不存在，请先上传并等待预处理完成'})
|
||||
|
||||
dataset_id = upload_ds_id if store.get_dataset(upload_ds_id) else f'entity_{run_id}'
|
||||
dataset_id = upload_ds_id if store.get_dataset(upload_ds_id) else f'entity_{pk}'
|
||||
|
||||
def _analysis_fn(run, ctx):
|
||||
from tianxuan.llm_orchestrator import run_llm_pipeline, LLMConfig
|
||||
dataset_id = ctx['dataset_id']
|
||||
entity_col = ctx['entity_col']
|
||||
pk = ctx['pk']
|
||||
|
||||
def progress_callback(step, tool_name, result):
|
||||
def progress_callback(step, tool_name, result=None, meta=None):
|
||||
nonlocal run
|
||||
try:
|
||||
run = AnalysisRun.objects.get(id=run_id)
|
||||
run = AnalysisRun.objects.get(id=pk)
|
||||
pct = min(10 + step * 6, 95)
|
||||
run.progress_pct = pct
|
||||
run.progress_msg = f'步骤 {step}: {tool_name}'
|
||||
if result:
|
||||
summary = str(result)
|
||||
run.run_log += f'[{step}] {tool_name}: {summary}\n'
|
||||
|
||||
if tool_name == 'thinking':
|
||||
# Accumulate LLM thinking text
|
||||
if result:
|
||||
run.llm_thinking = (run.llm_thinking or '') + f'\n[{step}] {result}'
|
||||
save_fields = ['progress_pct', 'progress_msg', 'llm_thinking']
|
||||
elif tool_name == 'complete':
|
||||
pass # No extra tracking needed
|
||||
else:
|
||||
run.run_log += f'[{step}] {tool_name}\n'
|
||||
run.save(update_fields=['progress_pct', 'progress_msg', 'run_log'])
|
||||
# Real tool call — log it
|
||||
if result:
|
||||
summary = str(result)
|
||||
run.run_log += f'[{step}] {tool_name}: {summary}\n'
|
||||
else:
|
||||
run.run_log += f'[{step}] {tool_name}\n'
|
||||
|
||||
# Capture structured tool call data
|
||||
if meta and isinstance(meta, dict) and 'args' in meta:
|
||||
tool_calls = list(run.tool_calls_json or [])
|
||||
tool_calls.append({
|
||||
'step': step,
|
||||
'name': tool_name,
|
||||
'input': meta['args'],
|
||||
'output': result,
|
||||
})
|
||||
run.tool_calls_json = tool_calls
|
||||
save_fields = ['progress_pct', 'progress_msg', 'run_log', 'tool_calls_json']
|
||||
|
||||
run.save(update_fields=save_fields)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -915,7 +910,7 @@ def run_llm_analysis_view(request):
|
||||
run.progress_msg = 'LLM 初始化中...'
|
||||
run.save(update_fields=['status', 'progress_pct', 'progress_msg', 'run_log'])
|
||||
|
||||
result = run_llm_pipeline(dataset_id, entity_col, config=llm_cfg, callback=progress_callback)
|
||||
result = run_llm_pipeline(dataset_id, config=llm_cfg, callback=progress_callback)
|
||||
|
||||
if 'error' in result:
|
||||
run.status = 'failed'
|
||||
@@ -931,7 +926,7 @@ def run_llm_analysis_view(request):
|
||||
tb = traceback.format_exc()
|
||||
logger.error(tb)
|
||||
try:
|
||||
run = AnalysisRun.objects.get(id=run_id)
|
||||
run = AnalysisRun.objects.get(id=pk)
|
||||
run.status = 'failed'
|
||||
run.progress_msg = f'失败: {tb}'
|
||||
run.run_log += f'\n[ERROR] {tb}'
|
||||
@@ -942,12 +937,12 @@ def run_llm_analysis_view(request):
|
||||
|
||||
t = threading.Thread(
|
||||
target=_run_pipeline_worker,
|
||||
args=(run_id, _analysis_fn),
|
||||
kwargs=dict(dataset_id=dataset_id, entity_col=entity_col, head=head),
|
||||
args=(pk, _analysis_fn),
|
||||
kwargs=dict(dataset_id=dataset_id, head=head, pk=pk),
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
return JsonResponse({'status': 'started', 'run_id': run_id})
|
||||
return JsonResponse({'status': 'started', 'run_id': run.display_id})
|
||||
|
||||
|
||||
def auto_page(request):
|
||||
@@ -962,11 +957,11 @@ def auto_page(request):
|
||||
|
||||
# ── Progress API ───────────────────────────────────────────────────────
|
||||
|
||||
def run_status_api(request, run_id):
|
||||
def run_status_api(request, display_id):
|
||||
"""Return current run status as JSON."""
|
||||
run = get_object_or_404(AnalysisRun, id=run_id)
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
return JsonResponse({
|
||||
'id': run.id,
|
||||
'id': run.display_id,
|
||||
'status': run.status,
|
||||
'entity_count': run.entity_count,
|
||||
'cluster_count': run.cluster_count,
|
||||
@@ -974,6 +969,8 @@ def run_status_api(request, run_id):
|
||||
'progress_pct': run.progress_pct,
|
||||
'progress_msg': run.progress_msg,
|
||||
'run_log': run.run_log[-2000:], # last 2000 chars
|
||||
'llm_thinking': run.llm_thinking[-5000:], # last 5K chars
|
||||
'tool_calls': run.tool_calls_json,
|
||||
})
|
||||
|
||||
|
||||
@@ -1029,6 +1026,9 @@ def _extract_flows_from_df(df, MAX_ROWS):
|
||||
dst_ip_col = next((c for c in df.columns if 'dst_ip' in c.lower() or ':ipd' in c.lower()), None)
|
||||
tls_col = next((c for c in df.columns if c.lower() in ('0ver', 'tls_version', '_tlsver', 'version')), None)
|
||||
bytes_col = next((c for c in df.columns if c.lower() in ('8ack', '8pak', '8byt', 'bytes_sent', 'bytes')), None)
|
||||
time_col = next((c for c in df.columns if c.lower() in (
|
||||
'timestamp', 'time', '_time', 'ts', 'datetime', 'epoch',
|
||||
'created_at', 'created', 'event_time')), None)
|
||||
|
||||
flows = []
|
||||
for row in df.iter_rows(named=True):
|
||||
@@ -1066,12 +1066,46 @@ def _extract_flows_from_df(df, MAX_ROWS):
|
||||
tls_label = 'other'
|
||||
|
||||
bs = float(row.get(bytes_col or '', 0) or 0)
|
||||
|
||||
# Extract timestamp for animation ordering
|
||||
t_raw = row.get(time_col) if time_col else None
|
||||
t_val = None
|
||||
if t_raw is not None:
|
||||
if isinstance(t_raw, (int, float)):
|
||||
t_val = float(t_raw)
|
||||
elif hasattr(t_raw, 'timestamp'):
|
||||
t_val = t_raw.timestamp()
|
||||
else:
|
||||
try:
|
||||
t_val = float(t_raw)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
flows.append({
|
||||
'slat': slat, 'slon': slon,
|
||||
'dlat': dlat, 'dlon': dlon,
|
||||
'tls': tls_label,
|
||||
'bytes': bs,
|
||||
'time': t_val,
|
||||
})
|
||||
|
||||
# Normalize timestamps to 0-1 and sort flows by time
|
||||
timed = [f for f in flows if f['time'] is not None]
|
||||
if len(timed) > 1:
|
||||
timed.sort(key=lambda f: f['time'])
|
||||
t_min, t_max = timed[0]['time'], timed[-1]['time']
|
||||
t_range = max(t_max - t_min, 1)
|
||||
for f in timed:
|
||||
f['time'] = (f['time'] - t_min) / t_range
|
||||
# Untimed flows get evenly spaced after timed ones
|
||||
untimed = [f for f in flows if f['time'] is None]
|
||||
for i, f in enumerate(untimed):
|
||||
f['time'] = min(1.0, (len(timed) + i) / max(len(flows) - 1, 1))
|
||||
flows = timed + untimed
|
||||
else:
|
||||
for i, f in enumerate(flows):
|
||||
f['time'] = i / max(len(flows) - 1, 1)
|
||||
|
||||
return flows
|
||||
|
||||
|
||||
@@ -1089,7 +1123,7 @@ def globe_view(request):
|
||||
if data_param:
|
||||
store = SessionStore()
|
||||
flows = []
|
||||
all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('id', 'csv_glob'))
|
||||
all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('display_id', 'csv_glob'))
|
||||
try:
|
||||
entry = store.get_dataset(data_param)
|
||||
if entry is not None:
|
||||
@@ -1106,7 +1140,7 @@ def globe_view(request):
|
||||
})
|
||||
|
||||
# Get all completed runs for the selector
|
||||
all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('id', 'csv_glob'))
|
||||
all_runs = list(AnalysisRun.objects.filter(status='completed').order_by('-id').values('display_id', 'csv_glob'))
|
||||
|
||||
# Get selected run IDs from query params (multiple checkboxes)
|
||||
selected_ids_list = request.GET.getlist('runs')
|
||||
@@ -1116,7 +1150,7 @@ def globe_view(request):
|
||||
sid = sid.strip()
|
||||
if sid.isdigit():
|
||||
try:
|
||||
r = AnalysisRun.objects.get(id=int(sid), status='completed')
|
||||
r = AnalysisRun.objects.get(display_id=int(sid), status='completed')
|
||||
selected_runs.append(r)
|
||||
except AnalysisRun.DoesNotExist:
|
||||
pass
|
||||
@@ -1138,7 +1172,7 @@ def globe_view(request):
|
||||
all_selected_failed = False
|
||||
break # stop after first successful load
|
||||
except Exception as exc:
|
||||
logger.warning('[GLOBE] Failed to load run %s (%s): %s', run.id, run.csv_glob, exc)
|
||||
logger.warning('[GLOBE] Failed to load run %s (%s): %s', run.display_id, run.csv_glob, exc)
|
||||
continue
|
||||
|
||||
if all_selected_failed or not flows:
|
||||
@@ -1167,7 +1201,7 @@ def globe_view(request):
|
||||
resp = render(request, 'tianxuan/globe.html', {
|
||||
'geo_flows': json.dumps(flows),
|
||||
'all_runs': all_runs,
|
||||
'selected_ids': [r.id for r in selected_runs],
|
||||
'selected_ids': [r.display_id for r in selected_runs],
|
||||
})
|
||||
resp['Cache-Control'] = 'no-cache, no-store, must-revalidate'
|
||||
resp['Pragma'] = 'no-cache'
|
||||
@@ -1312,3 +1346,36 @@ def tool_plan(request):
|
||||
|
||||
return JsonResponse({'error': 'Method not allowed'}, status=405)
|
||||
|
||||
|
||||
@csrf_exempt
|
||||
def retry_run(request, display_id):
|
||||
"""Reset a failed run and re-trigger based on run_type."""
|
||||
if request.method != 'POST':
|
||||
return JsonResponse({'error': 'POST required'}, status=405)
|
||||
|
||||
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
||||
if run.status != 'failed':
|
||||
return JsonResponse({'error': '只能重试失败的分析运行'}, status=400)
|
||||
|
||||
run.status = 'pending'
|
||||
run.error_message = ''
|
||||
run.progress_pct = 0
|
||||
run.progress_msg = ''
|
||||
run.save(update_fields=['status', 'error_message', 'progress_pct', 'progress_msg'])
|
||||
|
||||
if run.run_type == 'upload' and run.csv_glob:
|
||||
upload_dir = Path(run.csv_glob).parent
|
||||
if upload_dir.is_dir():
|
||||
t = threading.Thread(
|
||||
target=_background_process,
|
||||
args=(run.id, upload_dir),
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
else:
|
||||
run.error_message = '上传数据目录已不存在，请重新上传'
|
||||
run.status = 'failed'
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
|
||||
return redirect('analysis:run_detail', display_id=run.display_id)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user