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

analysis/tool_registry.py

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

@@ -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: