PM-pinou
505 lines
23 KiB
Python
505 lines
23 KiB
Python
"""Cluster views: overview, detail, globe flows, and the clustering pipeline."""
|
|
import asyncio
|
|
import json
|
|
import traceback
|
|
import logging
|
|
|
|
from django.shortcuts import render, get_object_or_404
|
|
from django.http import Http404
|
|
|
|
from analysis.models import AnalysisRun, ClusterResult, EntityProfile
|
|
from analysis import nl_describe
|
|
from analysis.constants import RANDOM_SEED, UMAP_BATCH_SIZE, UMAP_TRAIN_SAMPLE
|
|
from analysis.profile_util import profile
|
|
|
|
from .helpers import _extract_lat, _extract_lon
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def cluster_overview(request, display_id):
|
|
"""Overview of all clusters for a run, with UMAP scatter data and geo scatter data."""
|
|
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()
|
|
|
|
# Build UMAP scatter data for Canvas / Three.js 3D
|
|
entities = run.entities.exclude(embedding_x=None).exclude(embedding_y=None).only(
|
|
'embedding_x', 'embedding_y', 'embedding_z', 'cluster_label', 'entity_value'
|
|
)
|
|
has_z = any(e.embedding_z is not None for e in entities)
|
|
scatter_data = [
|
|
{'x': e.embedding_x, 'y': e.embedding_y, 'z': e.embedding_z if has_z else 0,
|
|
'label': e.cluster_label if e.cluster_label is not None else -1, 'entity': e.entity_value}
|
|
for e in entities
|
|
]
|
|
|
|
# Noise cluster: features + entity detail
|
|
noise_features = []
|
|
noise_summary = ''
|
|
noise_entity_count = 0
|
|
if noise:
|
|
noise_features_qs = noise.features.all().order_by('-distinguishing_score')[:15]
|
|
noise_features = [
|
|
{'feature_name': f.feature_name, 'score': f.distinguishing_score, 'mean': f.mean,
|
|
'std': f.std, 'p25': f.p25, 'p75': f.p75}
|
|
for f in noise_features_qs
|
|
]
|
|
if noise_features:
|
|
noise_summary = nl_describe.describe_cluster(noise_features)
|
|
noise_entity_count = noise.size
|
|
|
|
# LLM workflow timeline for auto runs
|
|
llm_timeline_json = ''
|
|
if run.run_type == 'auto':
|
|
llm_thinking = run.llm_thinking or ''
|
|
tool_calls = run.tool_calls_json or []
|
|
if tool_calls or llm_thinking:
|
|
llm_timeline_json = json.dumps({
|
|
'thinking': llm_thinking,
|
|
'tool_calls': tool_calls,
|
|
})
|
|
|
|
# Build geo scatter data from feature_json (lat/lon)
|
|
geo_data = []
|
|
geo_skipped = 0
|
|
for e in run.entities.all().only('feature_json', 'cluster_label', 'entity_value'):
|
|
fj = e.feature_json
|
|
if not fj:
|
|
continue
|
|
lat = _extract_lat(fj.get('avg_latitude'))
|
|
lon = _extract_lon(fj.get('avg_longitude'))
|
|
if lat is not None and lon is not None:
|
|
geo_data.append({
|
|
'x': lon,
|
|
'y': lat,
|
|
'label': e.cluster_label if e.cluster_label is not None else -1,
|
|
'entity': e.entity_value,
|
|
})
|
|
else:
|
|
geo_skipped += 1
|
|
|
|
# Cluster sizes for bar chart
|
|
cluster_sizes = [{'label': c.cluster_label, 'size': c.size, 'silhouette': c.silhouette_score}
|
|
for c in clusters]
|
|
|
|
# Top features per cluster from ClusterFeature + NL summaries
|
|
top_features = {}
|
|
for c in clusters:
|
|
features_qs = c.features.all().order_by('-distinguishing_score')[:10]
|
|
feats = [
|
|
{'feature_name': f.feature_name, 'score': f.distinguishing_score, 'mean': f.mean,
|
|
'std': f.std}
|
|
for f in features_qs
|
|
]
|
|
top_features[str(c.cluster_label)] = feats
|
|
# Attach features + NL summary directly to cluster object for template
|
|
c._top_features = feats
|
|
c.nl_summary = nl_describe.describe_cluster(feats) if feats else ''
|
|
|
|
return render(request, 'analysis/cluster_overview.html', {
|
|
'run': run,
|
|
'clusters': clusters,
|
|
'noise': noise,
|
|
'scatter_data_json': json.dumps(scatter_data),
|
|
'has_z': has_z,
|
|
'geo_data_json': json.dumps(geo_data),
|
|
'geo_skipped': geo_skipped,
|
|
'geo_count': len(geo_data),
|
|
'cluster_sizes_json': json.dumps(cluster_sizes),
|
|
'top_features_json': json.dumps(top_features),
|
|
'noise_features': noise_features,
|
|
'noise_summary': noise_summary,
|
|
'noise_entity_count': noise_entity_count,
|
|
'llm_timeline_json': llm_timeline_json,
|
|
# Globe flow data for the sidebar
|
|
'globe_flows_json': json.dumps(_get_globe_flows(run)),
|
|
'selected_ds_id': f'upload_{run.display_id}',
|
|
})
|
|
|
|
|
|
def _get_globe_flows(run, max_rows=500):
|
|
"""Extract flow data for 3D globe visualization from a run."""
|
|
from analysis.data_loader import load_csv_directory, load_from_db
|
|
try:
|
|
lf = None
|
|
if run.sqlite_table:
|
|
lf = load_from_db(run.sqlite_table)
|
|
if lf is None and run.csv_glob:
|
|
lf, _, _, _, _ = load_csv_directory(run.csv_glob)
|
|
if lf is None:
|
|
return []
|
|
df = lf.head(max_rows).collect()
|
|
from analysis.type_classifier import TLS_HEX_MAP
|
|
from analysis.geoip import lookup as geo_lookup
|
|
flows = []
|
|
src_ip_col = next((c for c in df.columns if 'src_ip' in c.lower() or ':ips' in c.lower()), None)
|
|
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)
|
|
for row in df.iter_rows(named=True):
|
|
sip = str(row.get(src_ip_col, '')) if src_ip_col else ''
|
|
dip = str(row.get(dst_ip_col, '')) if dst_ip_col else ''
|
|
if not sip or not dip:
|
|
continue
|
|
sg = geo_lookup(sip)
|
|
dg = geo_lookup(dip)
|
|
if not sg or not dg:
|
|
continue
|
|
tls_val = str(row.get(tls_col, '') or '') if tls_col else ''
|
|
tls_ver = TLS_HEX_MAP.get(tls_val.replace(' ', ''), tls_val)
|
|
flows.append({
|
|
'slat': sg['lat'], 'slon': sg['lon'],
|
|
'dlat': dg['lat'], 'dlon': dg['lon'],
|
|
'tls': tls_ver,
|
|
'bytes': float(row.get(bytes_col, 0) or 0) if bytes_col else 0,
|
|
})
|
|
return flows
|
|
except Exception:
|
|
logger.error("unknown failed: {}".format(traceback.format_exc()))
|
|
return []
|
|
|
|
|
|
def cluster_detail(request, display_id, cluster_label):
|
|
"""Detailed view of a single cluster with entity list, SVD features, and NL summary."""
|
|
run = get_object_or_404(AnalysisRun, display_id=display_id)
|
|
try:
|
|
cluster_label = int(cluster_label)
|
|
except (ValueError, TypeError):
|
|
raise Http404("Invalid cluster label")
|
|
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()[:100]
|
|
nl_summary = ''
|
|
feats_list = [{'feature_name': f.feature_name, 'score': f.distinguishing_score,
|
|
'mean': f.mean, 'std': f.std} for f in features[:10]]
|
|
if feats_list:
|
|
nl_summary = nl_describe.describe_cluster(feats_list)
|
|
|
|
# Build entity data with feature_json values
|
|
entity_data = []
|
|
for e in entities:
|
|
entry = {'id': e.id, 'value': e.entity_value,
|
|
'embedding_x': e.embedding_x, 'embedding_y': e.embedding_y}
|
|
if e.feature_json:
|
|
entry['features'] = {k: round(float(v), 4) if isinstance(v, (int, float)) else str(v)
|
|
for k, v in e.feature_json.items()}
|
|
entity_data.append(entry)
|
|
|
|
return render(request, 'analysis/cluster_detail.html', {
|
|
'run': run,
|
|
'cluster': cluster,
|
|
'features': features,
|
|
'entities': entities,
|
|
'entity_data_json': json.dumps(entity_data),
|
|
'nl_summary': nl_summary,
|
|
})
|
|
|
|
|
|
@profile
|
|
def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
|
|
min_cluster_size, run_umap=True, head=None):
|
|
"""Clustering pipeline on raw rows: clustering → feature extraction → UMAP → ORM save.
|
|
|
|
Operates directly on raw data rows (no entity aggregation).
|
|
Per-row results stored in EntityProfile with row index as identifier.
|
|
|
|
Args:
|
|
run: AnalysisRun ORM object.
|
|
store: SessionStore instance.
|
|
ds_id: dataset ID in SessionStore pointing to raw data.
|
|
feature_columns: list of feature column names (None = auto-detect all numeric).
|
|
algorithm: clustering algorithm.
|
|
min_cluster_size: int for HDBSCAN.
|
|
run_umap: bool, whether to compute and save UMAP-2D embeddings.
|
|
head: optional int, downsample to at most this many rows (min 100).
|
|
"""
|
|
import warnings
|
|
warnings.filterwarnings('ignore', category=RuntimeWarning, module='sklearn')
|
|
|
|
try:
|
|
from analysis.tool_registry import _handle_run_clustering, _handle_extract_features
|
|
import polars as pl
|
|
|
|
entry = store.get_dataset(ds_id)
|
|
if entry is None:
|
|
run.status = 'failed'
|
|
run.error_message = '数据集在会话中不存在'
|
|
run.save(update_fields=['status', 'error_message'])
|
|
return
|
|
|
|
schema = entry.get('schema', {})
|
|
|
|
# ── Downsampling ──
|
|
if head is not None and head > 0:
|
|
lf = entry['lazyframe']
|
|
try:
|
|
row_count = lf.select(pl.len()).collect(streaming=True).item()
|
|
except Exception:
|
|
logger.error("unknown failed: {}".format(traceback.format_exc()))
|
|
row_count = 0
|
|
|
|
if row_count > head:
|
|
head = max(head, 100) # ensure minimum rows for clustering
|
|
run.progress_msg = f'正在采样 {head}/{row_count} 行...'
|
|
run.save(update_fields=['progress_msg'])
|
|
lf = lf.head(head)
|
|
store.store_dataset(
|
|
ds_id, lf,
|
|
schema=schema,
|
|
metadata=entry.get('metadata', {}),
|
|
)
|
|
entry = store.get_dataset(ds_id)
|
|
row_count = head
|
|
|
|
# Resolve feature columns — use actual LazyFrame dtypes (not stored schema dict
|
|
# which may be stale after Utf8 coercion in _background_process).
|
|
numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
|
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
|
pl.Float32, pl.Float64)
|
|
|
|
lf = entry['lazyframe']
|
|
live_schema = lf.collect_schema()
|
|
live_names = live_schema.names()
|
|
live_dtypes = list(live_schema.dtypes())
|
|
|
|
if feature_columns:
|
|
feature_cols = [
|
|
c for c in feature_columns
|
|
if c in live_names and live_dtypes[live_names.index(c)] in numeric_types
|
|
]
|
|
else:
|
|
feature_cols = [
|
|
live_names[i] for i, dt in enumerate(live_dtypes)
|
|
if dt in numeric_types and not live_names[i].startswith('_')
|
|
]
|
|
|
|
# ── Step 1: Clustering ────────────────────────────────────────────────
|
|
# If no numeric columns detected, pass empty list to let
|
|
# _handle_run_clustering apply its Utf8→Float64 coercion fallback.
|
|
run.status = 'clustering'
|
|
run.progress_pct = 70
|
|
run.progress_msg = '正在进行聚类分析...'
|
|
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
|
|
|
result = asyncio.run(_handle_run_clustering(
|
|
dataset_id=ds_id,
|
|
cluster_columns=feature_cols,
|
|
algorithm=algorithm,
|
|
params={'min_cluster_size': min_cluster_size},
|
|
random_state=RANDOM_SEED,
|
|
))
|
|
|
|
if 'error' in result:
|
|
err = result['error']
|
|
if 'numeric' in err.lower():
|
|
available = [
|
|
f"{k}({v})" for k, v in schema.items()
|
|
if v.split('(')[0].strip() in numeric_types and not k.startswith('_')
|
|
]
|
|
err += f"。可用数值列: {available}"
|
|
raise Exception(err)
|
|
|
|
cluster_id = result.get('cluster_result_id', '')
|
|
run.cluster_count = result.get('n_clusters', 0)
|
|
run.save(update_fields=['cluster_count'])
|
|
|
|
# ── Step 2: Feature extraction ──────────────────────────────────────
|
|
run.status = 'extracting'
|
|
run.progress_pct = 90
|
|
run.progress_msg = '正在提取聚类特征...'
|
|
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
|
|
|
feat_result = asyncio.run(_handle_extract_features(
|
|
dataset_id=ds_id,
|
|
cluster_result_id=cluster_id,
|
|
top_k=10, method='zscore', save_to_db=True,
|
|
run_id=run.id,
|
|
))
|
|
|
|
if 'error' in feat_result:
|
|
raise Exception(feat_result['error'])
|
|
|
|
# ── Step 2b: Cluster SVD feature extraction ──────────────────────
|
|
run.progress_msg = '正在进行聚类SVD特征提取...'
|
|
run.save(update_fields=['progress_msg'])
|
|
try:
|
|
from analysis.distance import cluster_svd_extract
|
|
lf = entry['lazyframe']
|
|
cluster_entry = store.get_cluster_result(cluster_id)
|
|
labels_list = cluster_entry.get('labels', []) if cluster_entry else []
|
|
|
|
if labels_list and feature_cols:
|
|
svd_results = cluster_svd_extract(lf, labels_list, feature_cols)
|
|
|
|
# Persist SVD features to ClusterFeature (method='cluster_svd')
|
|
from analysis.models import ClusterFeature as CF
|
|
from analysis.models import ClusterResult as CR
|
|
for label, svd_info in svd_results.items():
|
|
cr = CR.objects.filter(run=run, cluster_label=label).first()
|
|
if cr is None:
|
|
continue
|
|
# Delete old zscore entries for this cluster, replace with SVD
|
|
CF.objects.filter(cluster=cr, distinguishing_method='zscore').delete()
|
|
for feat_name in svd_info.get('features', [])[:10]:
|
|
idx = svd_info['features'].index(feat_name) if feat_name in svd_info['features'] else -1
|
|
strength = (svd_info.get('feature_strength', [])[idx]
|
|
if idx >= 0 and idx < len(svd_info.get('feature_strength', []))
|
|
else None)
|
|
CF.objects.get_or_create(
|
|
cluster=cr,
|
|
feature_name=feat_name,
|
|
defaults={
|
|
'mean': None, 'std': None, 'median': None,
|
|
'p25': None, 'p75': None, 'missing_rate': None,
|
|
'distinguishing_score': float(strength) if strength is not None else None,
|
|
'distinguishing_method': 'cluster_svd',
|
|
},
|
|
)
|
|
except Exception as svd_err:
|
|
logger.warning(f'Cluster SVD extraction skipped: {svd_err}')
|
|
|
|
# ── Step 2c: Save per-row cluster labels to EntityProfile ──
|
|
# Always persist row→cluster mapping so cluster_detail can show data rows
|
|
# even when UMAP is skipped or fails.
|
|
try:
|
|
cluster_entry_save = store.get_cluster_result(cluster_id)
|
|
labels_list_save = cluster_entry_save.get('labels', []) if cluster_entry_save else []
|
|
if labels_list_save:
|
|
from analysis.models import EntityProfile as EP_save
|
|
from analysis.models import ClusterResult as CR_save
|
|
# Delete stale EP records for this run first
|
|
EP_save.objects.filter(run=run).delete()
|
|
cr_cache_save = {}
|
|
row_profiles = []
|
|
for r_idx in range(len(labels_list_save)):
|
|
lbl = labels_list_save[r_idx]
|
|
ck = f'{run.id}_{lbl}'
|
|
if ck not in cr_cache_save:
|
|
cr_cache_save[ck] = CR_save.objects.filter(run=run, cluster_label=lbl).first()
|
|
row_profiles.append(EP_save(
|
|
entity_value=f'row_{r_idx}', run=run, cluster_label=lbl,
|
|
cluster=cr_cache_save[ck],
|
|
embedding_x=None, embedding_y=None, embedding_z=None,
|
|
))
|
|
if row_profiles:
|
|
EP_save.objects.bulk_create(row_profiles, ignore_conflicts=True, batch_size=1000)
|
|
run.entity_count = len(row_profiles)
|
|
run.save(update_fields=['entity_count'])
|
|
except Exception as save_err:
|
|
logger.warning(f'Row labels save skipped: {save_err}')
|
|
|
|
# ── Step 3: UMAP-3D embedding (2D fallback) ───────────────────────
|
|
if run_umap:
|
|
try:
|
|
lf = entry['lazyframe']
|
|
# UMAP: sample max 10K for training, batch-transform rest in 1K batches
|
|
MAX_UMAP_TRAIN = UMAP_TRAIN_SAMPLE
|
|
BATCH_SIZE = UMAP_BATCH_SIZE
|
|
try:
|
|
n_total = lf.select(pl.len()).collect(streaming=True).item()
|
|
except Exception:
|
|
logger.error("unknown failed: {}".format(traceback.format_exc()))
|
|
n_total = 0
|
|
from sklearn.preprocessing import StandardScaler
|
|
import umap
|
|
import numpy as np
|
|
|
|
live_schema = lf.collect_schema()
|
|
num_cols = [name for name, dt in zip(live_schema.names(), live_schema.dtypes())
|
|
if dt in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
|
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
|
pl.Float32, pl.Float64) and not name.startswith('_')]
|
|
|
|
if len(num_cols) >= 2:
|
|
if n_total > MAX_UMAP_TRAIN:
|
|
run.progress_msg = f'UMAP 采样 {MAX_UMAP_TRAIN}/{n_total} 行训练...'
|
|
run.save(update_fields=['progress_msg'])
|
|
# LazyFrame.sample() is unavailable in Polars 1.42 — collect then sample
|
|
df_sample = df_umap.sample(n=MAX_UMAP_TRAIN, seed=RANDOM_SEED)
|
|
df_sample = df_umap
|
|
|
|
umap_components = 3 # try 3D first
|
|
coords = None
|
|
reducer = None
|
|
|
|
if len(df_umap) > 2:
|
|
for attempt_n in (3, 2):
|
|
try:
|
|
if n_total > MAX_UMAP_TRAIN:
|
|
# Train UMAP on 10K sample, batch-transform full dataset
|
|
mat_sample = df_sample.select(num_cols).to_numpy()
|
|
scaler = StandardScaler().fit(mat_sample)
|
|
mat_sample_scaled = np.nan_to_num(
|
|
scaler.transform(mat_sample), nan=0.0)
|
|
reducer = umap.UMAP(n_components=attempt_n,
|
|
random_state=RANDOM_SEED,
|
|
n_neighbors=15, min_dist=0.1,
|
|
metric='euclidean')
|
|
reducer.fit(mat_sample_scaled)
|
|
# Batch transform in 1K-row batches
|
|
coords_list = []
|
|
for start in range(0, len(df_umap), BATCH_SIZE):
|
|
end = min(start + BATCH_SIZE, len(df_umap))
|
|
mat_batch = scaler.transform(
|
|
df_umap[start:end].select(num_cols).to_numpy())
|
|
mat_batch = np.nan_to_num(mat_batch, nan=0.0)
|
|
coords_list.append(reducer.transform(mat_batch))
|
|
coords = np.vstack(coords_list)
|
|
else:
|
|
mat = np.nan_to_num(StandardScaler().fit_transform(
|
|
df_umap.select(num_cols).to_numpy()), nan=0.0)
|
|
reducer = umap.UMAP(n_components=attempt_n,
|
|
random_state=RANDOM_SEED,
|
|
n_neighbors=15, min_dist=0.1,
|
|
metric='euclidean')
|
|
coords = reducer.fit_transform(mat)
|
|
umap_components = attempt_n
|
|
break
|
|
except Exception as e_umap:
|
|
if attempt_n == 2:
|
|
raise
|
|
logger.warning(f'UMAP 3D failed ({e_umap}), falling back to 2D')
|
|
|
|
if coords is None:
|
|
raise Exception('UMAP produced no output')
|
|
|
|
cluster_entry = store.get_cluster_result(cluster_id)
|
|
labels_list = cluster_entry.get('labels', []) if cluster_entry else []
|
|
|
|
from analysis.models import EntityProfile as EP
|
|
|
|
# Update existing EP records with UMAP coordinates
|
|
for i in range(len(df_umap)):
|
|
ev = f'row_{i}'
|
|
has_z = umap_components >= 3 and coords.shape[1] >= 3
|
|
update_fields = {
|
|
'embedding_x': float(coords[i, 0]) if i < len(coords) else None,
|
|
'embedding_y': float(coords[i, 1]) if i < len(coords) else None,
|
|
}
|
|
if has_z:
|
|
update_fields['embedding_z'] = float(coords[i, 2]) if i < len(coords) else None
|
|
EP.objects.filter(run=run, entity_value=ev).update(**update_fields)
|
|
run.entity_count = len(df_umap)
|
|
run.save(update_fields=['entity_count'])
|
|
except Exception as umap_err:
|
|
logger.warning(f'UMAP embedding skipped: {umap_err}')
|
|
|
|
# ── Done ──
|
|
run.status = 'completed'
|
|
run.progress_pct = 100
|
|
run.progress_msg = '分析完成'
|
|
run.save(update_fields=['status', 'progress_pct', 'progress_msg', 'entity_count'])
|
|
|
|
except Exception:
|
|
tb = traceback.format_exc()
|
|
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=['status', 'error_message', 'progress_msg', 'run_log'])
|
|
except Exception as e:
|
|
logger.warning('save failed after pipeline error: %s', e)
|