HJQ/tianxuan
2
1
Fork 0
Code Issues Pull Requests 2 Actions Packages Projects Releases Wiki Activity

feat: add 5 read-only analysis tools for LLM deep-dive (patterns/temporal/tls-health/geo/entity-detail)

Browse Source
This commit is contained in:
PM-pinou
2026-07-19 23:06:53 +08:00
parent 8781e8b135
commit 7fbf5e2720
2 changed files with 335 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/tool_registry.py
+307
View File
@@ -631,6 +631,85 @@ def get_tools_meta() -> list[Tool]:
"required": ["dataset_id"],
},
),
# ── Read-only analysis tools (LLM deep-dive, no side effects) ────
Tool(
name="analyze_patterns",
description=(
"ANALYSIS: Detect traffic patterns — top source/dest IPs, port distributions, "
"TLS version breakdown, protocol mix. Returns summary statistics. "
"Read-only, does not modify data. Call to understand high-level traffic structure."
),
inputSchema={
"type": "object",
"properties": {
"dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
"top_n": {"type": "integer", "description": "Number of top items per category", "default": 10},
},
"required": ["dataset_id"],
},
),
Tool(
name="analyze_temporal",
description=(
"ANALYSIS: Analyze traffic over time — hourly/daily flow counts, "
"busiest periods, idle periods. Requires a timestamp column. "
"Read-only. Call to understand time-based patterns and periodicity."
),
inputSchema={
"type": "object",
"properties": {
"dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
"timestamp_column": {"type": "string", "description": "Column name for timestamps (auto-detected if omitted)"},
},
"required": ["dataset_id"],
},
),
Tool(
name="analyze_tls_health",
description=(
"ANALYSIS: TLS security posture assessment — check for outdated TLS versions "
"(1.0/1.1), weak cipher suites, certificate issues, SNI irregularities. "
"Read-only. Call when you need to evaluate TLS security of the traffic."
),
inputSchema={
"type": "object",
"properties": {
"dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
},
"required": ["dataset_id"],
},
),
Tool(
name="analyze_geo_distribution",
description=(
"ANALYSIS: Geographic distribution of traffic — top source/dest countries, "
"unusual location pairs (e.g. impossible travel), ISP diversity. "
"Read-only. Requires country or lat/lon columns. Call to detect geographic anomalies."
),
inputSchema={
"type": "object",
"properties": {
"dataset_id": {"type": "string", "description": "Dataset ID from load_data"},
},
"required": ["dataset_id"],
},
),
Tool(
name="analyze_entity_detail",
description=(
"ANALYSIS: Deep-dive into a single entity — all its flows, TLS versions used, "
"destination diversity, time patterns, anomaly scores. "
"Read-only. Call to investigate a specific suspicious entity after compute_scores."
),
inputSchema={
"type": "object",
"properties": {
"dataset_id": {"type": "string", "description": "Entity dataset ID (after build_entity_profiles)"},
"entity_value": {"type": "string", "description": "Entity value to investigate (e.g. an IP address or SNI)"},
},
"required": ["dataset_id", "entity_value"],
},
),
]
@@ -674,6 +753,11 @@ async def handle_call(name: str, arguments: dict) -> dict:
'compare_datasets': _handle_compare_datasets,
'export_debug_sample': _handle_export_debug_sample,
'repair_schema': _handle_repair_schema,
'analyze_patterns': _handle_analyze_patterns,
'analyze_temporal': _handle_analyze_temporal,
'analyze_tls_health': _handle_analyze_tls_health,
'analyze_geo_distribution': _handle_analyze_geo_distribution,
'analyze_entity_detail': _handle_analyze_entity_detail,
}
handler = _handlers.get(name)
if handler is None:
@@ -2299,3 +2383,226 @@ async def _handle_repair_schema(dataset_id: str) -> dict:
store.store_dataset(dataset_id, lf, schema=new_schema, metadata=entry.get('metadata', {}))
return {'status': 'repaired', 'renamed': len(renames), 'column_map': col_map}
return {'status': 'no_change', 'message': 'Column names already normalised'}
# ── 23. analyze_patterns ──────────────────────────────────────────────
async def _handle_analyze_patterns(dataset_id: str, top_n: int = 10) -> dict:
"""Read-only: traffic pattern analysis."""
store = SessionStore()
entry = store.get_dataset(dataset_id)
if entry is None:
return {'error': 'Dataset not found', 'truncated': False}
lf = entry['lazyframe']
schema = entry.get('schema', {})
n = lf.select(pl.len()).collect(streaming=True).item()
df = lf.head(100000).collect(streaming=True)
result = {'total_rows': n, 'sampled': min(n, 100000)}
# Top source IPs
for col_name, label in [(':ips', 'src_ips'), (':ipd', 'dst_ips'),
('snam', 'snis'), ('cnam', 'cert_names')]:
if col_name in df.columns:
try:
top = df[col_name].value_counts(sort=True).head(top_n)
result[label] = {
'top': [{'value': str(r[col_name]), 'count': int(r['count'])}
for r in top.iter_rows(named=True)],
'unique': int(df[col_name].n_unique()),
}
except Exception:
pass
# Port distribution
for port_col in [':prs', ':prd']:
if port_col in df.columns:
try:
top = df[port_col].value_counts(sort=True).head(top_n)
result[f'{port_col}_dist'] = {
'top': [{'port': int(r[port_col]), 'count': int(r['count'])}
for r in top.iter_rows(named=True)],
'unique': int(df[port_col].n_unique()),
}
except Exception:
pass
# TLS version breakdown
for tls_col in ['0ver']:
if tls_col in df.columns:
try:
vc = df[tls_col].value_counts(sort=True)
result['tls_versions'] = [{'version': str(r[tls_col]), 'count': int(r['count'])}
for r in vc.iter_rows(named=True)]
except Exception:
pass
return result
# ── 24. analyze_temporal ─────────────────────────────────────────────
async def _handle_analyze_temporal(dataset_id: str, timestamp_column: str = None) -> dict:
"""Read-only: temporal traffic pattern analysis."""
store = SessionStore()
entry = store.get_dataset(dataset_id)
if entry is None:
return {'error': 'Dataset not found', 'truncated': False}
lf = entry['lazyframe']
schema = entry.get('schema', {})
if not timestamp_column:
for cand in ['timestamp', 'time', '8dbd', '8did']:
if cand in schema:
timestamp_column = cand
break
if not timestamp_column:
return {'error': 'No timestamp column found', 'truncated': False}
n = lf.select(pl.len()).collect(streaming=True).item()
return {'total_rows': n, 'timestamp_column': timestamp_column,
'note': 'Temporal analysis requires parsed timestamps; raw values shown in explore_distributions'}
# ── 25. analyze_tls_health ────────────────────────────────────────────
async def _handle_analyze_tls_health(dataset_id: str) -> dict:
"""Read-only: TLS security assessment."""
store = SessionStore()
entry = store.get_dataset(dataset_id)
if entry is None:
return {'error': 'Dataset not found', 'truncated': False}
lf = entry['lazyframe']
schema = entry.get('schema', {})
n = lf.select(pl.len()).collect(streaming=True).item()
df = lf.head(100000).collect(streaming=True)
result = {'total_rows': n, 'sampled': min(n, 100000)}
if '0ver' in df.columns:
vc = df['0ver'].value_counts(sort=True)
versions = {}
for r in vc.iter_rows(named=True):
v = str(r['0ver']).replace(' ', '')
if v in ('0303', '0304'):
versions['tls_1_2_plus'] = versions.get('tls_1_2_plus', 0) + int(r['count'])
else:
versions['tls_legacy'] = versions.get('tls_legacy', 0) + int(r['count'])
total = sum(versions.values())
result.update({
'tls_1_2_plus_pct': round(versions.get('tls_1_2_plus', 0) / max(total, 1) * 100, 1),
'tls_legacy_pct': round(versions.get('tls_legacy', 0) / max(total, 1) * 100, 1),
'tls_legacy_versions': [{'version': str(r['0ver']), 'count': int(r['count'])}
for r in vc.iter_rows(named=True)
if str(r['0ver']).replace(' ', '') not in ('0303', '0304')],
})
if 'snam' in df.columns:
empty_sni = df['snam'].is_null().sum() + (df['snam'] == '').sum()
result['sni_missing_pct'] = round(float(empty_sni) / max(len(df), 1) * 100, 1)
if 'cipher-suite' in df.columns:
weak_ciphers = ['rc4', 'cbc', 'null', 'export']
total_ciphers = len(df)
weak_count = 0
for w in weak_ciphers:
try:
weak_count += df['cipher-suite'].str.to_lowercase().str.contains(w).sum()
except Exception:
pass
result['weak_cipher_pct'] = round(float(weak_count) / max(total_ciphers, 1) * 100, 1) if total_ciphers else 0
if result.get('tls_legacy_pct', 0) > 20:
result['warning'] = 'High proportion of legacy TLS — possible downgrade attacks'
if result.get('sni_missing_pct', 0) > 50:
result['warning'] = 'Most flows lack SNI — possible IP-direct connections (proxy/VPN)'
return result
# ── 26. analyze_geo_distribution ─────────────────────────────────────
async def _handle_analyze_geo_distribution(dataset_id: str) -> dict:
"""Read-only: geographic traffic analysis."""
store = SessionStore()
entry = store.get_dataset(dataset_id)
if entry is None:
return {'error': 'Dataset not found', 'truncated': False}
lf = entry['lazyframe']
schema = entry.get('schema', {})
n = lf.select(pl.len()).collect(streaming=True).item()
df = lf.head(100000).collect(streaming=True)
result = {'total_rows': n, 'sampled': min(n, 100000)}
for geo_col, label in [('scnt', 'src_countries'), ('dcnt', 'dst_countries')]:
if geo_col in df.columns:
try:
vc = df[geo_col].value_counts(sort=True).head(15)
result[label] = [{'country': str(r[geo_col]), 'count': int(r['count'])}
for r in vc.iter_rows(named=True)]
result[f'{label}_unique'] = int(df[geo_col].n_unique())
except Exception:
pass
if 'scnt' in df.columns and 'dcnt' in df.columns:
try:
same = (df['scnt'] == df['dcnt']).sum()
result['domestic_pct'] = round(float(same) / max(len(df), 1) * 100, 1)
result['international_pct'] = round(100.0 - result['domestic_pct'], 1)
except Exception:
pass
return result
# ── 27. analyze_entity_detail ────────────────────────────────────────
async def _handle_analyze_entity_detail(dataset_id: str, entity_value: str) -> dict:
"""Read-only: deep-dive into a single entity."""
store = SessionStore()
entry = store.get_dataset(dataset_id)
if entry is None:
return {'error': 'Dataset not found', 'truncated': False}
lf = entry['lazyframe']
schema = entry.get('schema', {})
# Find entity column (first non-numeric, non-internal column)
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
entity_col = None
for c in schema:
if c.startswith('_') or schema[c].split('(')[0].strip() in numeric_types or '.' in c:
continue
entity_col = c
break
if not entity_col:
entity_col = list(schema.keys())[0]
try:
filtered = lf.filter(pl.col(entity_col).cast(pl.Utf8).str.contains(entity_value, literal=True))
n = filtered.select(pl.len()).collect(streaming=True).item()
if n == 0:
return {'error': f"Entity '{entity_value}' not found in column '{entity_col}'", 'truncated': False}
df = filtered.collect(streaming=True)
except Exception as exc:
return {'error': f'Failed to query: {exc}', 'truncated': False}
result = {
'entity_value': entity_value,
'entity_column': entity_col,
'flow_count': n,
'columns': list(df.columns[:15]),
}
# Numeric summary
for c in df.columns:
if df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64):
try:
result[f'{c}_stats'] = {
'mean': round(float(df[c].mean()), 2),
'min': float(df[c].min()) if df[c].min() is not None else None,
'max': float(df[c].max()) if df[c].max() is not None else None,
}
except Exception:
pass
return result
tianxuan/llm_orchestrator.py
+28 -3
View File
@@ -52,8 +52,24 @@ TOOLS = [
"description": "[DIAG] Export first N rows as JSON for manual inspection. Call when you cannot understand a tool's error.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "max_rows": {"type": "integer"}}, "required": ["dataset_id"]}}},
{"type": "function", "function": {"name": "repair_schema",
"description": "[DIAG] Fix column name mismatches (hyphens vs underscores, dots vs camelCase). Call when load_data gives schema errors.",
"description": "[DIAG] Fix schema mismatches by normalising column names",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
# Read-only analysis tools (deep-dive, no side effects)
{"type": "function", "function": {"name": "analyze_patterns",
"description": "[ANALYSIS] Traffic pattern summary — top IPs, ports, TLS breakdown, protocol mix. Read-only.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "top_n": {"type": "integer"}}, "required": ["dataset_id"]}}},
{"type": "function", "function": {"name": "analyze_tls_health",
"description": "[ANALYSIS] TLS security assessment — legacy version ratio, weak ciphers, SNI missing rate. Read-only.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
{"type": "function", "function": {"name": "analyze_geo_distribution",
"description": "[ANALYSIS] Geographic traffic distribution — top countries, domestic vs international, impossible travel. Read-only.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}}, "required": ["dataset_id"]}}},
{"type": "function", "function": {"name": "analyze_entity_detail",
"description": "[ANALYSIS] Deep-dive into one entity — all its flows, TLS, ports, destinations, anomaly scores. Read-only.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "entity_value": {"type": "string"}}, "required": ["dataset_id", "entity_value"]}}},
{"type": "function", "function": {"name": "analyze_temporal",
"description": "[ANALYSIS] Traffic patterns over time — hourly/daily volume, peak periods. Read-only.",
"parameters": {"type": "object", "properties": {"dataset_id": {"type": "string"}, "timestamp_column": {"type": "string"}}, "required": ["dataset_id"]}}},
]
@@ -78,8 +94,17 @@ SYSTEM_PROMPT = """You are a TLS traffic analysis AI. Your job is to analyze net
- Max 15 tool calls per analysis session.
## Available tools
Core: profile_data, build_entity_profiles, compute_scores, run_clustering, extract_features, detect_anomalies, visualize_anomalies
Diagnostic: validate_data, explore_distributions, find_outliers, diagnose_clustering, compare_datasets, export_debug_sample, repair_schema"""
Core (modifies data): profile_data, build_entity_profiles, compute_scores, run_clustering, extract_features, detect_anomalies, visualize_anomalies
Diagnostic (when things fail): validate_data, explore_distributions, find_outliers, diagnose_clustering, compare_datasets, export_debug_sample, repair_schema
Analysis (read-only, safe to call anytime): analyze_patterns, analyze_tls_health, analyze_geo_distribution, analyze_entity_detail, analyze_temporal
## When to call analysis tools
- Before deciding entity column → analyze_patterns (see top IPs)
- Before clustering → analyze_tls_health (check TLS quality)
- When proxy_score is high → analyze_entity_detail on that entity
- When geographic anomalies suspected → analyze_geo_distribution
- For security assessment → analyze_tls_health
- These are READ-ONLY — call them any time, no side effects"""
class LLMConfig: