chore: bump version to v1.1.2
- New timeline UI: interleaved LLM reasoning + tool call display - Fix: run_log truncation for large tool outputs - Fix: lat/lon + values reduced to ~1% (was 99%) - Fix: @csrf_exempt on upload/delete endpoints - Fix: DATA_UPLOAD_MAX_NUMBER_FILES for 1998-file upload - Fix: dtype alignment for multi-batch CSV processing - New: tests/test_e2e_full.py with real DeepSeek LLM - New: 1998-file test data generator
This commit is contained in:
PM-pinou
@@ -156,12 +156,19 @@ entity:
|
||||
|
||||
### 运行全部测试
|
||||
```bash
|
||||
runtime\python\python.exe -m pytest tests -q
|
||||
# 单元测试（不需要Django）
|
||||
runtime\python\python.exe -m pytest tests/test_data_loader.py tests/test_type_classifier.py -q
|
||||
|
||||
# 集成测试（自动启动后端、模拟前端行为、监控stderr报错）
|
||||
runtime\python\python.exe tests/test_integration.py
|
||||
```
|
||||
```
|
||||
92 passed (unit tests); ALL INTEGRATION TESTS PASSED
|
||||
92 passed (unit tests); ALL INTEGRATION TESTS PASSED
|
||||
```
|
||||
|
||||
> ⚠️ **重要**: `tests/test_integration.py` 模拟前端HTML行为（上传CSV → 轮询状态 → 访问页面 → 删除）。
|
||||
> **当前端HTML/API发生变化时，必须同步更新此脚本**，否则集成测试将不再反映真实用户流程。
|
||||
|
||||
### 跨机一致性测试
|
||||
```bash
|
||||
# 两台机器各自执行:
|
||||
|
||||
+18
-2
@@ -288,7 +288,19 @@ def load_from_db(
|
||||
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])
|
||||
try:
|
||||
df = pl.from_dicts([dict(r) for r in rows])
|
||||
except Exception:
|
||||
# Fallback: load all columns as Utf8 to handle mixed-type SQLite columns,
|
||||
# then apply schema_overrides to restore intended types.
|
||||
import io, csv as _csv_mod
|
||||
buf = io.StringIO()
|
||||
writer = _csv_mod.writer(buf)
|
||||
writer.writerow(rows[0].keys())
|
||||
for r in rows:
|
||||
writer.writerow(str(v) if v is not None else '' for v in r)
|
||||
buf.seek(0)
|
||||
df = pl.read_csv(buf, infer_schema_length=0)
|
||||
|
||||
# Apply type overrides to restore original column types
|
||||
if schema_overrides:
|
||||
@@ -298,7 +310,11 @@ def load_from_db(
|
||||
continue
|
||||
dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype)
|
||||
if df[col].dtype != dtype_obj:
|
||||
casts[col] = pl.col(col).cast(dtype_obj)
|
||||
try:
|
||||
casts[col] = pl.col(col).cast(dtype_obj)
|
||||
except Exception:
|
||||
# If a cast fails (e.g. "A" → Int64), keep the column as-is
|
||||
pass
|
||||
if casts:
|
||||
df = df.with_columns(list(casts.values()))
|
||||
|
||||
|
||||
+174
-6
@@ -4,13 +4,13 @@ Each tool is defined by a :class:`Tool <mcp.types.Tool>` metadata object and
|
||||
an async handler function. Handlers interact with the :class:`SessionStore`
|
||||
singleton and the Django ORM.
|
||||
|
||||
All 11 tools:
|
||||
All 12 tools:
|
||||
1. load_data 7. filter_and_cluster
|
||||
2. profile_data 8. export_results
|
||||
3. filter_data 9. list_datasets
|
||||
4. preprocess_data 10. drop_dataset
|
||||
5. run_clustering 11. clone_dataset
|
||||
6. evaluate_clustering
|
||||
2. profile_data 8. build_entity_profiles
|
||||
3. filter_data 9. export_results
|
||||
4. preprocess_data 10. list_datasets
|
||||
5. run_clustering 11. drop_dataset
|
||||
6. evaluate_clustering 12. clone_dataset
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
@@ -462,6 +462,39 @@ def get_tools_meta() -> list[Tool]:
|
||||
"required": ["dataset_id", "filters"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="build_entity_profiles",
|
||||
description=(
|
||||
"按实体列（源IP、SNI等）分组原始流数据。自动检测实体列（如 src_ip, "
|
||||
"dst_ip, sni, host, mac, domain），按列分组并计算每个实体的聚合特征："
|
||||
"流数量、唯一端口数、TLS版本分布、字节总量、唯一目标数。"
|
||||
"生成的数据集每个实体一行。结果存储在会话中，可供后续 compute_scores 使用。"
|
||||
),
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"dataset_id": {
|
||||
"type": "string",
|
||||
"description": "Dataset ID from load_data or profile_data",
|
||||
},
|
||||
"entity_column": {
|
||||
"type": "string",
|
||||
"description": "Specific entity column to group by (e.g. 'src_ip')",
|
||||
},
|
||||
"entity_columns": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Multiple entity columns to group by",
|
||||
},
|
||||
"auto_detect": {
|
||||
"type": "boolean",
|
||||
"description": "Auto-detect entity columns (IP, SNI, host, MAC, domain)",
|
||||
"default": True,
|
||||
},
|
||||
},
|
||||
"required": ["dataset_id"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="compute_scores",
|
||||
description=(
|
||||
@@ -775,6 +808,7 @@ async def handle_call(name: str, arguments: dict) -> dict:
|
||||
'list_datasets': _handle_list_datasets,
|
||||
'drop_dataset': _handle_drop_dataset,
|
||||
'clone_dataset': _handle_clone_dataset,
|
||||
'build_entity_profiles': _handle_build_entity_profiles,
|
||||
'compute_scores': _handle_compute_scores,
|
||||
'filter_and_cluster': _handle_filter_and_cluster,
|
||||
'detect_anomalies': _handle_detect_anomalies,
|
||||
@@ -1808,6 +1842,140 @@ async def _handle_clone_dataset(dataset_id: str) -> dict:
|
||||
return {'original_dataset_id': dataset_id, 'cloned_dataset_id': new_id}
|
||||
|
||||
|
||||
# ── 12. build_entity_profiles ────────────────────────────────────────────
|
||||
|
||||
_ENTITY_KEYWORDS = [
|
||||
'src_ip', 'dst_ip', 'source_ip', 'destination_ip', 'src_addr', 'dst_addr',
|
||||
'sip', 'dip', 'ip_src', 'ip_dst', 'saddr', 'daddr',
|
||||
'sni', 'host', 'domain', 'uri', 'url', 'fqdn', 'server_name',
|
||||
'mac', 'mac_addr', 'src_mac', 'dst_mac',
|
||||
'user_agent', 'ua', 'asn', 'organization', 'org',
|
||||
]
|
||||
|
||||
_NUMERIC_TYPES_FOR_AGG = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||||
pl.Float32, pl.Float64)
|
||||
|
||||
|
||||
async def _handle_build_entity_profiles(
|
||||
dataset_id: str,
|
||||
entity_column: Optional[str] = None,
|
||||
entity_columns: Optional[list[str]] = None,
|
||||
auto_detect: bool = True,
|
||||
) -> dict:
|
||||
"""Group raw flow data by entity column(s) and compute per-entity aggregate features.
|
||||
|
||||
Stores the result in the session store and returns the new dataset_id.
|
||||
"""
|
||||
store = SessionStore()
|
||||
entry = store.get_dataset(dataset_id)
|
||||
if entry is None:
|
||||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
col_names = list(schema.keys())
|
||||
|
||||
# ── Resolve entity columns ──────────────────────────────────────────
|
||||
if entity_column and entity_column in col_names:
|
||||
group_cols = [entity_column]
|
||||
elif entity_columns:
|
||||
group_cols = [c for c in entity_columns if c in col_names]
|
||||
if not group_cols:
|
||||
return {'error': f'None of entity_columns {entity_columns} exist in dataset. Available: {col_names[:20]}',
|
||||
'truncated': False}
|
||||
elif auto_detect:
|
||||
# Auto-detect: find columns matching entity keywords
|
||||
matched = []
|
||||
for kw in _ENTITY_KEYWORDS:
|
||||
found = [c for c in col_names if kw.lower() in c.lower()]
|
||||
matched.extend(found)
|
||||
# Deduplicate while preserving order
|
||||
seen = set()
|
||||
group_cols = [c for c in matched if not (c in seen or seen.add(c))]
|
||||
if not group_cols:
|
||||
# Fallback: pick first non-numeric, non-internal column
|
||||
for c in col_names:
|
||||
if c.startswith('_'):
|
||||
continue
|
||||
group_cols = [c]
|
||||
break
|
||||
# Limit to at most 3 entity columns
|
||||
group_cols = group_cols[:3]
|
||||
else:
|
||||
return {'error': 'No entity column specified and auto_detect is disabled',
|
||||
'truncated': False}
|
||||
|
||||
if not group_cols:
|
||||
return {'error': 'Could not determine entity column(s) for grouping',
|
||||
'truncated': False}
|
||||
|
||||
# ── Build aggregation expressions ──────────────────────────────────
|
||||
# Always: count of flows
|
||||
agg_exprs: list[pl.Expr] = [
|
||||
pl.len().alias('flow_count'),
|
||||
pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
|
||||
]
|
||||
|
||||
# Numeric aggregations (sum, mean)
|
||||
numeric_cols = [c for c in col_names if c not in group_cols
|
||||
and schema[c].startswith(('Int', 'UInt', 'Float'))]
|
||||
for c in numeric_cols[:20]:
|
||||
agg_exprs.append(pl.sum(c).alias(f'total_{c}'))
|
||||
agg_exprs.append(pl.mean(c).alias(f'avg_{c}'))
|
||||
|
||||
# Unique counts for string-like entity columns (dst_ip, dst_port, sni, etc.)
|
||||
str_cols = [c for c in col_names if c not in group_cols
|
||||
and schema[c].startswith(('Utf8', 'String', 'Cat'))]
|
||||
for c in str_cols[:10]:
|
||||
agg_exprs.append(pl.col(c).n_unique().alias(f'unique_{c}'))
|
||||
|
||||
# Also count distinct values of specific keyword-matching columns
|
||||
for kw in ['dst_ip', 'dip', 'dst_addr', 'daddr', 'ip_dst', 'dest_ip',
|
||||
'dst_port', 'dport', 'destination_port', 'dest_port',
|
||||
'port', 'protocol', 'proto', 'sni', 'server_name',
|
||||
'cipher_suite', 'cipher', 'tls_version', 'version', '0ver',
|
||||
'country', 'scnt', 'dcnt', 'asn', 'organization']:
|
||||
found = [c for c in col_names if c not in group_cols and kw.lower() in c.lower()]
|
||||
for c in found:
|
||||
alias_name = f'unique_{c}'
|
||||
if alias_name not in [str(a.meta.output_name()) if hasattr(a, 'meta') else '' for a in agg_exprs]:
|
||||
agg_exprs.append(pl.col(c).n_unique().alias(alias_name))
|
||||
|
||||
# ── Execute grouping ───────────────────────────────────────────────
|
||||
try:
|
||||
entity_lf = lf.group_by(group_cols).agg(agg_exprs)
|
||||
# Estimate row count
|
||||
n_entities = entity_lf.select(pl.len()).collect(streaming=True).item()
|
||||
except Exception as e:
|
||||
return {'error': f'Entity aggregation failed: {e}', 'truncated': False}
|
||||
|
||||
# ── Store result ───────────────────────────────────────────────────
|
||||
new_schema = {name: str(dtype) for name, dtype in
|
||||
zip(entity_lf.collect_schema().names(), entity_lf.collect_schema().dtypes())}
|
||||
new_id = f'entity_{_generate_id("ds")}'
|
||||
store.store_dataset(
|
||||
dataset_id=new_id,
|
||||
lazyframe=entity_lf,
|
||||
schema=new_schema,
|
||||
metadata={
|
||||
'row_count': n_entities,
|
||||
'entity_columns': group_cols,
|
||||
'parent_dataset_id': dataset_id,
|
||||
'n_agg_features': len(agg_exprs),
|
||||
},
|
||||
parent_id=dataset_id,
|
||||
)
|
||||
|
||||
return {
|
||||
'dataset_id': new_id,
|
||||
'entity_columns': group_cols,
|
||||
'n_entities': n_entities,
|
||||
'n_features': len(agg_exprs),
|
||||
'aggregated_columns': list(new_schema.keys())[:20],
|
||||
}
|
||||
|
||||
|
||||
# ── Adaptive anomaly scoring helpers (inlined from entity_aggregator) ─────
|
||||
|
||||
_ANOMALY_FEATURE_COLS = [
|
||||
|
||||
+48
-10
@@ -436,8 +436,9 @@ def _background_process(run_id, upload_dir):
|
||||
table_name = f'_data_{run_id}'
|
||||
|
||||
# ── Phase 1: Scan first batch ONLY for schema consistency ──
|
||||
# Use large infer_schema_length to catch mixed-type columns early
|
||||
first_batch = csv_paths[:batch_size]
|
||||
first_lf = pl.scan_csv(first_batch, low_memory=True,
|
||||
first_lf = pl.scan_csv(first_batch, infer_schema_length=100000,
|
||||
schema_overrides=lat_lon_overrides)
|
||||
first_schema = dict(first_lf.collect_schema())
|
||||
schema = {name: str(dtype) for name, dtype in first_schema.items()}
|
||||
@@ -449,7 +450,7 @@ def _background_process(run_id, upload_dir):
|
||||
total_rows = 0
|
||||
for i in range(0, total_files, batch_size):
|
||||
batch = csv_paths[i:i+batch_size]
|
||||
batch_lf = pl.scan_csv(batch, low_memory=True,
|
||||
batch_lf = pl.scan_csv(batch, infer_schema_length=100000,
|
||||
schema_overrides=lat_lon_overrides)
|
||||
df_batch = batch_lf.collect(streaming=True)
|
||||
|
||||
@@ -462,6 +463,27 @@ def _background_process(run_id, upload_dir):
|
||||
df_batch = df_batch.drop(extra)
|
||||
df_batch = df_batch.select(reference_columns)
|
||||
|
||||
# Dtype alignment: cast each column to the reference dtype
|
||||
# to prevent mixed-type columns when loading from SQLite later
|
||||
casts = {}
|
||||
for col in reference_columns:
|
||||
ref_dtype_str = schema.get(col)
|
||||
if ref_dtype_str and col in df_batch.columns:
|
||||
# Parse the polars dtype string back to a pl.DataType
|
||||
_base = ref_dtype_str.split('(')[0].strip()
|
||||
_ref_dtype = {'Float64': pl.Float64, 'Float32': pl.Float32,
|
||||
'Int64': pl.Int64, 'Int32': pl.Int32,
|
||||
'Utf8': pl.Utf8, 'String': pl.Utf8,
|
||||
'Boolean': pl.Boolean, 'Bool': pl.Boolean}.get(_base)
|
||||
if _ref_dtype and df_batch[col].dtype != _ref_dtype:
|
||||
try:
|
||||
casts[col] = pl.col(col).cast(_ref_dtype)
|
||||
except Exception:
|
||||
# If cast fails (e.g. "A" → Int64), force to Utf8
|
||||
casts[col] = pl.col(col).cast(pl.Utf8)
|
||||
if casts:
|
||||
df_batch = df_batch.with_columns(list(casts.values()))
|
||||
|
||||
mode = 'replace' if not table_created else 'append'
|
||||
save_to_db(run.id, df_batch.lazy(), mode=mode)
|
||||
table_created = True
|
||||
@@ -499,7 +521,7 @@ def _background_process(run_id, upload_dir):
|
||||
run.total_flows = total_rows
|
||||
run.save(update_fields=['sqlite_table', 'total_flows'])
|
||||
|
||||
ds_id = f'upload_{run_id}'
|
||||
ds_id = f'upload_{run.display_id}'
|
||||
store.store_dataset(ds_id, lf, schema=schema, metadata={
|
||||
'row_count': total_rows, 'file_count': total_files, 'upload_dir': str(upload_dir),
|
||||
'csv_glob': str(upload_dir / '*.csv'),
|
||||
@@ -825,10 +847,15 @@ def run_analysis(request):
|
||||
|
||||
store = SessionStore()
|
||||
try:
|
||||
upload_ds_id = f'upload_{pk}'
|
||||
# Use display_id for dataset key (consistent with _background_process)
|
||||
display_id = ctx.get('display_id', pk)
|
||||
upload_ds_id = f'upload_{display_id}'
|
||||
entry = store.get_dataset(upload_ds_id)
|
||||
if entry is None:
|
||||
# Fallback: try entity dataset (from old pipeline)
|
||||
# Fallback: try PK-based key (old pipeline)
|
||||
entry = store.get_dataset(f'upload_{pk}')
|
||||
if entry is None:
|
||||
# Fallback: try entity dataset (from older pipeline)
|
||||
entry = store.get_dataset(f'entity_{pk}')
|
||||
if entry is None:
|
||||
run.error_message = '请先上传并等待预处理完成'
|
||||
@@ -836,7 +863,7 @@ def run_analysis(request):
|
||||
run.save(update_fields=['status', 'error_message'])
|
||||
return
|
||||
|
||||
ds_id = upload_ds_id if store.get_dataset(upload_ds_id) else f'entity_{pk}'
|
||||
ds_id = upload_ds_id if store.get_dataset(upload_ds_id) else f'upload_{pk}'
|
||||
|
||||
# Use the unified clustering pipeline (clustering → extraction → PCA)
|
||||
from analysis.views import _run_clustering_pipeline
|
||||
@@ -859,7 +886,8 @@ def run_analysis(request):
|
||||
target=_run_pipeline_worker,
|
||||
args=(run_id, _analysis_fn),
|
||||
kwargs=dict(feature_columns=feature_columns, algorithm=algorithm,
|
||||
min_cluster_size=min_cluster_size, head=head, pk=pk),
|
||||
min_cluster_size=min_cluster_size, head=head, pk=pk,
|
||||
display_id=run.display_id),
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
@@ -887,8 +915,13 @@ def run_llm_analysis_view(request):
|
||||
|
||||
from analysis.session_store import SessionStore
|
||||
store = SessionStore()
|
||||
upload_ds_id = f'upload_{pk}'
|
||||
# Use display_id for dataset key (consistent with _background_process)
|
||||
upload_ds_id = f'upload_{display_id}'
|
||||
entry = store.get_dataset(upload_ds_id)
|
||||
if entry is None:
|
||||
# Fallback: try PK-based key (old pipeline)
|
||||
upload_ds_id = f'upload_{pk}'
|
||||
entry = store.get_dataset(upload_ds_id)
|
||||
|
||||
if entry is None:
|
||||
# Try entity dataset as fallback
|
||||
@@ -922,8 +955,13 @@ def run_llm_analysis_view(request):
|
||||
else:
|
||||
# Real tool call — log it
|
||||
if result:
|
||||
summary = str(result)
|
||||
run.run_log += f'[{step}] {tool_name}: {summary}\n'
|
||||
if isinstance(result, dict):
|
||||
# Format dict as JSON, truncate large values
|
||||
import json as _json
|
||||
summary = _json.dumps(result, default=str, ensure_ascii=False)[:500]
|
||||
else:
|
||||
summary = str(result)[:500]
|
||||
run.run_log += f'[{step}] {tool_name}: {summary}...\n'
|
||||
else:
|
||||
run.run_log += f'[{step}] {tool_name}\n'
|
||||
|
||||
|
||||
+4
-4
@@ -23,10 +23,10 @@ anomaly_detection:
|
||||
suspicious: 0.95
|
||||
watch: 0.80
|
||||
llm:
|
||||
enabled: false
|
||||
base_url: ''
|
||||
api_key: ''
|
||||
model: gpt-4
|
||||
enabled: true
|
||||
base_url: 'https://api.deepseek.com'
|
||||
api_key: 'sk-ecd33e5341ff4ec9a2c5a5c97a97699b'
|
||||
model: deepseek-v4-flash
|
||||
entity:
|
||||
subnet_masks: [24, 28]
|
||||
ip_columns: ["src_ip", "dst_ip", ":ips", ":ipd", "server-ip", "client-ip"]
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
# start_ip,end_ip,lat,lon,city,country
|
||||
# === North America ===
|
||||
1.0.0.0,1.0.0.255,37.386,-122.083,Mountain View,US
|
||||
8.8.8.0,8.8.8.255,37.386,-122.083,Mountain View,US
|
||||
72.14.192.0,72.14.255.255,40.7128,-74.0060,New York,US
|
||||
74.125.0.0,74.125.255.255,40.7128,-74.0060,New York,US
|
||||
66.102.0.0,66.102.255.255,34.0522,-118.2437,Los Angeles,US
|
||||
64.0.0.0,64.255.255.255,43.6532,-79.3832,Toronto,CA
|
||||
# === East Asia ===
|
||||
1.2.4.0,1.2.4.255,39.9042,116.4074,Beijing,CN
|
||||
123.125.114.0,123.125.114.255,39.9042,116.4074,Beijing,CN
|
||||
101.95.0.0,101.95.255.255,31.2304,121.4737,Shanghai,CN
|
||||
58.0.0.0,58.255.255.255,22.3193,114.1694,Hong Kong,HK
|
||||
14.0.0.0,14.255.255.255,22.3193,114.1694,Hong Kong,HK
|
||||
1.160.0.0,1.160.255.255,25.0330,121.5654,Taipei,TW
|
||||
101.0.0.0,101.0.255.255,25.0330,121.5654,Taipei,TW
|
||||
27.0.0.0,27.0.255.255,35.6762,139.6503,Tokyo,JP
|
||||
210.130.0.0,210.130.255.255,35.6762,139.6503,Tokyo,JP
|
||||
110.0.0.0,110.0.255.255,37.5665,126.9780,Seoul,KR
|
||||
121.78.0.0,121.78.255.255,37.5665,126.9780,Seoul,KR
|
||||
# === Southeast Asia ===
|
||||
116.14.0.0,116.14.255.255,1.3521,103.8198,Singapore,SG
|
||||
58.8.0.0,58.8.255.255,13.7563,100.5018,Bangkok,TH
|
||||
# === South Asia ===
|
||||
14.96.0.0,14.96.255.255,19.0760,72.8777,Mumbai,IN
|
||||
# === Oceania ===
|
||||
1.40.0.0,1.40.255.255,-33.8688,151.2093,Sydney,AU
|
||||
149.135.0.0,149.135.255.255,-33.8688,151.2093,Sydney,AU
|
||||
# === Europe ===
|
||||
5.10.0.0,5.10.255.255,51.5074,-0.1278,London,GB
|
||||
81.2.0.0,81.2.255.255,51.5074,-0.1278,London,GB
|
||||
90.0.0.0,90.0.255.255,48.8566,2.3522,Paris,FR
|
||||
46.0.0.0,46.0.255.255,50.1109,8.6821,Frankfurt,DE
|
||||
145.0.0.0,145.0.255.255,52.3676,4.9041,Amsterdam,NL
|
||||
84.203.0.0,84.203.255.255,53.3498,-6.2603,Dublin,IE
|
||||
78.0.0.0,78.0.255.255,59.3293,18.0686,Stockholm,SE
|
||||
95.0.0.0,95.0.255.255,55.7558,37.6173,Moscow,RU
|
||||
# === Middle East ===
|
||||
91.72.0.0,91.72.255.255,25.2048,55.2708,Dubai,AE
|
||||
# === South America ===
|
||||
177.0.0.0,177.0.255.255,-23.5505,-46.6333,Sao Paulo,BR
|
||||
@@ -0,0 +1,42 @@
|
||||
import os, sys, django, urllib.request, glob, uuid, json, time
|
||||
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
||||
sys.path.insert(0, r'C:\Users\25044\Desktop\Proj\天璇')
|
||||
import django; django.setup()
|
||||
from analysis.models import AnalysisRun
|
||||
AnalysisRun.objects.all().delete()
|
||||
print('Cleaned old runs')
|
||||
|
||||
BASE = 'http://127.0.0.1:18766'
|
||||
data_dir = r'C:\Users\25044\Desktop\Proj\天璇\data\multi_upload'
|
||||
files = sorted(glob.glob(os.path.join(data_dir, '*.csv')))
|
||||
boundary = uuid.uuid4().hex
|
||||
lines = []
|
||||
for fp in files:
|
||||
with open(fp, 'rb') as f: d = f.read()
|
||||
fname = os.path.basename(fp)
|
||||
lines.append(f'--{boundary}'.encode())
|
||||
lines.append(f'Content-Disposition: form-data; name="files"; filename="{fname}"'.encode())
|
||||
lines.append(b'Content-Type: text/csv'); lines.append(b''); lines.append(d)
|
||||
lines.append(f'--{boundary}--'.encode())
|
||||
req = urllib.request.Request(f'{BASE}/upload/csv/', data=b'\r\n'.join(lines),
|
||||
headers={'Content-Type': f'multipart/form-data; boundary={boundary}'})
|
||||
r = urllib.request.urlopen(req, timeout=300)
|
||||
resp = json.loads(r.read()); r.close()
|
||||
print(f'Uploaded run_id={resp["run_id"]}')
|
||||
|
||||
run_id = resp['run_id']
|
||||
for i in range(120):
|
||||
time.sleep(2)
|
||||
try:
|
||||
req2 = urllib.request.Request(f'{BASE}/runs/{run_id}/status/')
|
||||
r2 = urllib.request.urlopen(req2, timeout=10)
|
||||
status_resp = json.loads(r2.read()); r2.close()
|
||||
status = status_resp.get('status', status_resp.get('state', ''))
|
||||
print(f' attempt {i+1}: status={status}')
|
||||
if status == 'ready':
|
||||
print(f'Run {run_id} is ready!')
|
||||
break
|
||||
except Exception as e:
|
||||
print(f' attempt {i+1}: error={e}')
|
||||
else:
|
||||
print('Timed out waiting for ready status')
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "tls-analyzer"
|
||||
version = "0.1.0"
|
||||
version = "1.1.2"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Generate 1998 CSV files (0.csv..1997.csv) with realistic TLS data matching TlsDB.csv.
|
||||
|
||||
Usage: runtime\python\python.exe scripts\gen_1998_files.py [--rows-per-file N] [--output-dir DIR]
|
||||
"""
|
||||
import random, csv, argparse
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timedelta
|
||||
from gen_test_data import generate_row, COLUMNS as columns
|
||||
|
||||
random.seed(42)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--rows-per-file', type=int, default=5, help='Rows per CSV file')
|
||||
parser.add_argument('--output-dir', type=str, default='data/multi_upload')
|
||||
args = parser.parse_args()
|
||||
|
||||
out_dir = Path(args.output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
total = 1998
|
||||
base_time = datetime(2026, 6, 1, 0, 0, 0)
|
||||
|
||||
for i in range(total):
|
||||
filepath = out_dir / f'{i}.csv'
|
||||
with open(filepath, 'w', newline='', encoding='utf-8') as f:
|
||||
writer = csv.DictWriter(f, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
for j in range(args.rows_per_file):
|
||||
row_time = base_time + timedelta(seconds=random.randint(0, 86400 * 30))
|
||||
row = generate_row(i * args.rows_per_file + j, row_time)
|
||||
writer.writerow(row)
|
||||
|
||||
print(f'Generated {total} files × {args.rows_per_file} rows = {total * args.rows_per_file} total rows')
|
||||
print(f'Output: {out_dir.resolve()}')
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+180
-140
@@ -1,8 +1,11 @@
|
||||
"""Generate synthetic TLS flow test data.
|
||||
"""Generate synthetic TLS flow test data matching TlsDB.csv column spec.
|
||||
|
||||
Coverage rates simulate real TLS traffic: IP/port always present,
|
||||
lat/lon/ISP often blank, TLS fingerprint fields sometimes blank,
|
||||
obscure fields rarely present.
|
||||
|
||||
Usage:
|
||||
runtime\\python\\python.exe scripts\\gen_test_data.py [--rows N] [--output PATH]
|
||||
runtime\\python\\python.exe scripts\\gen_test_data.py --globe [--rows N] [--output PATH]
|
||||
runtime\\python\\python.exe scripts\\gen_test_data.py [--rows N]
|
||||
"""
|
||||
import random
|
||||
import csv
|
||||
@@ -10,175 +13,212 @@ import argparse
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
# TLS version hex values
|
||||
TLS_VERSIONS = ['03 03', '03 04', '02 00', '03 01', '03 02']
|
||||
TLS_WEIGHTS = [0.60, 0.25, 0.05, 0.05, 0.05]
|
||||
random.seed(42)
|
||||
|
||||
# Cipher suites (hex)
|
||||
CIPHER_SUITES = ['c0 2b', 'c0 2f', '13 01', '13 02', 'c0 2c', 'cc a9', 'c0 23']
|
||||
CIPHER_NAMES = [
|
||||
'TLS_ECDHE_ECDSA_AES128_GCM_SHA256',
|
||||
'TLS_ECDHE_RSA_AES128_GCM_SHA256',
|
||||
'TLS_AES_128_GCM_SHA256',
|
||||
'TLS_AES_256_GCM_SHA384',
|
||||
'TLS_ECDHE_ECDSA_AES256_GCM_SHA384',
|
||||
'TLS_ECDHE_ECDSA_CHACHA20_POLY1305',
|
||||
'TLS_ECDHE_RSA_AES256_GCM_SHA384',
|
||||
# All columns from TlsDB.csv in realistic order
|
||||
COLUMNS = [
|
||||
':ips', ':ipd', ':prs', ':prd', 'scnt', 'dcnt',
|
||||
':ips.latd', ':ips.lond', ':ipd.latd', ':ipd.lond',
|
||||
':ips.ispn', ':ipd.ispn', ':ips.orgn', ':ipd.orgn',
|
||||
':ips.city', ':ipd.city', ':ips.anon', ':ipd.anon',
|
||||
'server-ip', 'client-ip', 'row', 'time', 'timestamp',
|
||||
'1ipp', '4dbn', 'tabl', '4ksz', 'cnrs', 'isrs',
|
||||
'cnam', '0ver', 'snam', '4dur', '8seq', '2tmo', 'name',
|
||||
'source-node', 'cipher-suite', 'ecdhe-named-curve',
|
||||
'0cph', '0crv', '0rnd', '0rnt',
|
||||
'8ack', '8pak', '8did', '4srs', '8ppk', '8ses', '8byt', '8dbd',
|
||||
'crcc', 'orga', 'orgu', 'eiph', '@iph',
|
||||
]
|
||||
|
||||
# ECDHE named curves (hex)
|
||||
NAMED_CURVES = ['00 1d', '00 17', '00 18', '00 19', 'x448']
|
||||
CURVE_NAMES = ['secp256r1', 'secp384r1', 'secp521r1', 'x25519', 'x448']
|
||||
# ── Data pools ───────────────────────────────────────────────────────────
|
||||
|
||||
# OCph (cipher suite ID)
|
||||
OCPH_VALS = ['c0 2b', 'c0 2f', '13 01', '13 02', 'c0 2c', 'cc a9']
|
||||
TLS_VERSIONS = ['03 03', '03 04', '02 00', '03 01', '03 02']
|
||||
TLS_WEIGHTS = [0.55, 0.30, 0.05, 0.05, 0.05]
|
||||
|
||||
# OCrv (named curve ID)
|
||||
OCRV_VALS = ['00 1d', '00 17', '00 18', '00 19', '00 1e']
|
||||
CIPHER_HEX = ['c0 2b', 'c0 2f', '13 01', '13 02', 'c0 2c', 'cc a9', 'c0 23',
|
||||
'c0 27', 'c0 13', 'c0 14', '00 9e', '00 9f', '00 35']
|
||||
CIPHER_NAMES = [
|
||||
'TLS_ECDHE_ECDSA_AES128_GCM_SHA256', 'TLS_ECDHE_RSA_AES128_GCM_SHA256',
|
||||
'TLS_AES_128_GCM_SHA256', 'TLS_AES_256_GCM_SHA384',
|
||||
'TLS_ECDHE_ECDSA_AES256_GCM_SHA384', 'TLS_ECDHE_ECDSA_CHACHA20_POLY1305',
|
||||
'TLS_ECDHE_RSA_AES256_GCM_SHA384', 'TLS_ECDHE_RSA_AES128_SHA256',
|
||||
'TLS_ECDHE_ECDSA_AES128_SHA', 'TLS_ECDHE_RSA_AES128_SHA',
|
||||
'TLS_RSA_WITH_AES_128_GCM_SHA256', 'TLS_RSA_WITH_AES_256_GCM_SHA384',
|
||||
'TLS_RSA_WITH_AES_128_CBC_SHA',
|
||||
]
|
||||
|
||||
# Random bytes for 0rnd, 0rnt
|
||||
def _rand_hex(n=32):
|
||||
NAMED_CURVES_HEX = ['00 1d', '00 17', '00 18', '00 19', '00 1e']
|
||||
NAMED_CURVES_TEXT = ['secp256r1', 'secp384r1', 'secp521r1', 'x25519', 'x448']
|
||||
|
||||
SRC_IPS = [f'{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}' for _ in range(200)]
|
||||
DST_IPS = [f'{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}' for _ in range(500)]
|
||||
|
||||
PORTS = [443, 80, 8080, 8443, 465, 993, 995, 53, 22, 3389, 25, 110, 143, 21, 12345, 50000, 443, 443, 443, 443]
|
||||
COUNTRIES = ['US','CN','KR','JP','GB','DE','FR','RU','BR','IN','SG','NL','CA','AU','HK','UA','IL','SE','NO','FI']
|
||||
ISPS = ['China Telecom','China Mobile','BT Group','Orange','Deutsche Telekom','AT&T','Verizon','Comcast','NTT','KDDI','SK Telecom','Singtel','Telstra']
|
||||
ORGS = ['Baidu Inc.','Alibaba Inc.','Amazon.com Inc.','Google LLC','Microsoft Corp.','Meta Platforms','Tencent','Samsung','Apple Inc.','Netflix Inc.']
|
||||
CITIES = ['Beijing','Shanghai','Seoul','Berlin','Paris','London','New York','San Francisco','Tokyo','Singapore','Sydney','Moscow','Seattle','Dublin','Mumbai']
|
||||
SERVICES = ['mail.example.com','api.example.com','cdn.example.com','auth.example.net','stream.example.org','login.live.com','*.cloudfront.net','*.s3.amazonaws.com','graph.facebook.com','www.google.com']
|
||||
SRC_NODES = ['packet_capture','ssl_logs','netflow','zeek','suricata']
|
||||
NODE_NAMES = ['wan-link','core-02','gw-09','edge-01','backbone-03']
|
||||
|
||||
def _rand_ip(pool):
|
||||
return random.choice(pool)
|
||||
|
||||
def _rand_hex(n):
|
||||
return ' '.join(f'{random.randint(0,255):02x}' for _ in range(n))
|
||||
|
||||
# Country pools
|
||||
SRC_COUNTRIES = ['US', 'CN', 'KR', 'JP', 'GB', 'DE', 'FR', 'RU', 'BR', 'IN', 'SG', 'NL', 'CA', 'AU', 'HK']
|
||||
DST_COUNTRIES = ['US', 'CN', 'KR', 'JP', 'GB', 'DE', 'FR', 'RU', 'BR', 'IN', 'SG', 'NL', 'CA', 'AU', 'HK']
|
||||
def _maybe(rate):
|
||||
"""Return True with given probability (0.0-1.0)."""
|
||||
return random.random() < rate
|
||||
|
||||
# IP pools
|
||||
def _rand_ip():
|
||||
return f'{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}'
|
||||
def _blank_or(rate, value, blank_marker=''):
|
||||
"""Return value with rate probability, else blank_marker."""
|
||||
return value if _maybe(rate) else blank_marker
|
||||
|
||||
# Server names
|
||||
SERVICES = [
|
||||
'service36.example.com', 'mail.example.com', 'api.example.com',
|
||||
'cdn.example.com', 'dns.example.com', 'auth.example.net',
|
||||
'stream.example.org', 'web.example.io', 'storage.example.co',
|
||||
'cdn.cloudflare.com', 'api.github.com', 'login.live.com',
|
||||
]
|
||||
|
||||
# Source nodes
|
||||
SRC_NODES = ['packet_capture', 'ssl_logs', 'netflow', 'zeek', 'suricata']
|
||||
NODE_NAMES = ['wan-link', 'core-02', 'gw-09', 'edge-01', 'backbone-03']
|
||||
NODE_TYPES = ['gw-09', 'core-02', 'wan-link', 'edge-01', 'backbone-03']
|
||||
|
||||
def _weighted_bool(true_weight=0.7):
|
||||
return random.random() < true_weight
|
||||
def _blank_or_plus(rate, value):
|
||||
"""Return value with rate probability, else '+' (null marker).
|
||||
Signature matches _blank_or(rate, value, blank_marker=...)."""
|
||||
return value if _maybe(rate) else '+'
|
||||
|
||||
def generate_row(row_id, base_time):
|
||||
"""Generate a single TLS flow row."""
|
||||
src_ip = _rand_ip()
|
||||
dst_ip = _rand_ip()
|
||||
src_port = random.choice([443, 465, 993, 995, 8443, 8080, 53, 80, 22, 12345])
|
||||
dst_port = random.choice([443, 465, 993, 995, 8443, 53, 80, 22, 8080])
|
||||
src_country = random.choice(SRC_COUNTRIES)
|
||||
dst_country = random.choice(DST_COUNTRIES)
|
||||
server_ip = dst_ip if random.random() > 0.1 else _rand_ip()
|
||||
client_ip = src_ip if random.random() > 0.1 else _rand_ip()
|
||||
|
||||
"""Generate a single TLS flow row matching all TlsDB.csv columns."""
|
||||
src_ip = _rand_ip(SRC_IPS)
|
||||
dst_ip = _rand_ip(DST_IPS)
|
||||
src_port = random.choice(PORTS)
|
||||
dst_port = random.choice(PORTS)
|
||||
src_country = random.choice(COUNTRIES)
|
||||
dst_country = random.choice(COUNTRIES)
|
||||
server_ip = _rand_ip(DST_IPS + SRC_IPS)
|
||||
client_ip = src_ip if _maybe(0.85) else _rand_ip(SRC_IPS)
|
||||
tls_ver = random.choices(TLS_VERSIONS, weights=TLS_WEIGHTS, k=1)[0]
|
||||
cipher_hex = random.choice(CIPHER_HEX)
|
||||
cipher_name = random.choice(CIPHER_NAMES)
|
||||
curve_hex = random.choice(NAMED_CURVES_HEX)
|
||||
curve_name = random.choice(NAMED_CURVES_TEXT)
|
||||
service = random.choice(SERVICES)
|
||||
duration = round(random.uniform(0.5, 120.0), 2)
|
||||
sessions = round(random.uniform(1, 2000), 4)
|
||||
tmo = round(random.uniform(0, 1000), 4)
|
||||
ks = random.choice([128, 256, 384, 512, 1024, 2048, 4096])
|
||||
cnrs = 'True' if _weighted_bool(0.65) else ''
|
||||
isrs = 'True' if _weighted_bool(0.60) else ''
|
||||
ack = random.randint(100, 200000)
|
||||
ppk = random.randint(1, 500)
|
||||
dbd = random.randint(1, 200)
|
||||
ipp = random.randint(1, 50)
|
||||
dbn = random.randint(1, 100)
|
||||
tabl = random.choice(['packet_capture', 'ssl_logs', 'netflow'])
|
||||
name = random.choice(['wan-link', 'core-02', 'gw-09'])
|
||||
src_node = random.choice(SRC_NODES)
|
||||
node_name = random.choice(NODE_NAMES)
|
||||
node_type = random.choice(NODE_TYPES)
|
||||
cipher_idx = random.randint(0, len(CIPHER_SUITES) - 1)
|
||||
curve_idx = random.randint(0, len(NAMED_CURVES) - 1)
|
||||
ocph = OCPH_VALS[random.randint(0, len(OCPH_VALS) - 1)]
|
||||
ocrv = OCRV_VALS[random.randint(0, len(OCRV_VALS) - 1)]
|
||||
rand_hex_1 = _rand_hex(32)
|
||||
rand_hex_2 = _rand_hex(32)
|
||||
cert_cn = service if _maybe(0.7) else f'*.{random.choice(SERVICES).split(".")[-2:]}'
|
||||
ts = int(base_time.timestamp()) + row_id
|
||||
time_str = base_time.strftime('%H:%M:%S')
|
||||
date_str = base_time.strftime('%Y-%m-%d')
|
||||
|
||||
# Random lat/lon for GeoIP
|
||||
# Row-local lat/lon with realistic blank rates
|
||||
src_lat = round(random.uniform(-60, 60), 6)
|
||||
src_lon = round(random.uniform(-180, 180), 6)
|
||||
dst_lat = round(random.uniform(-60, 60), 6)
|
||||
dst_lon = round(random.uniform(-180, 180), 6)
|
||||
|
||||
isps = ['China Telecom', 'China Mobile', 'BT Group', 'Orange', 'Deutsche Telekom',
|
||||
'AT&T', 'Verizon', 'Comcast', 'NTT', 'KDDI']
|
||||
orgs = ['Baidu Inc.', 'Alibaba Inc.', 'Amazon.com Inc.', 'Google LLC',
|
||||
'Microsoft Corp.', 'Meta Platforms', 'Tencent', 'Samsung']
|
||||
cities = ['Beijing', 'Shanghai', 'Seoul', 'Berlin', 'Paris', 'London',
|
||||
'New York', 'San Francisco', 'Tokyo', 'Singapore', 'Sydney', 'Moscow']
|
||||
return {
|
||||
# ── Core (always present) ──
|
||||
':ips': src_ip,
|
||||
':ipd': dst_ip,
|
||||
':prs': src_port,
|
||||
':prd': dst_port,
|
||||
'time': base_time.strftime('%H:%M:%S'),
|
||||
'timestamp': ts,
|
||||
|
||||
# ── TLS fingerprint (very high coverage) ──
|
||||
'0ver': tls_ver,
|
||||
'0cph': cipher_hex,
|
||||
'cipher-suite': _blank_or(0.85, cipher_name),
|
||||
'ecdhe-named-curve': _blank_or(0.80, curve_name),
|
||||
'0crv': _blank_or(0.80, curve_hex),
|
||||
'0rnd': _blank_or(0.75, _rand_hex(28)),
|
||||
'0rnt': _blank_or(0.70, _rand_hex(4)),
|
||||
'4ksz': _blank_or(0.85, str(random.choice([128, 256, 384, 512]))),
|
||||
|
||||
# ── Session / connection info (medium-high coverage) ──
|
||||
'snam': _blank_or(0.80, service),
|
||||
'cnam': _blank_or(0.70, cert_cn),
|
||||
'server-ip': _blank_or(0.60, server_ip),
|
||||
'client-ip': _blank_or(0.85, client_ip),
|
||||
'4dur': _blank_or(0.75, str(round(random.uniform(0.5, 120.0), 2))),
|
||||
'8ack': _blank_or(0.80, str(random.randint(100, 200000))),
|
||||
'8pak': _blank_or(0.70, str(random.randint(1, 2000))),
|
||||
'8ppk': _blank_or(0.65, str(random.randint(1, 500))),
|
||||
'8seq': _blank_or(0.40, str(round(random.uniform(1, 2000), 4))),
|
||||
'8ses': _blank_or(0.50, str(round(random.uniform(1, 2000), 4))),
|
||||
'8byt': _blank_or(0.40, str(random.randint(200, 500000))),
|
||||
'2tmo': _blank_or(0.50, str(round(random.uniform(0, 1000), 4))),
|
||||
|
||||
# ── TCP / network info (medium coverage) ──
|
||||
'1ipp': str(random.randint(1, 50)),
|
||||
'row': _blank_or(1.0, ''), # TLS协议含义
|
||||
'tabl': random.choice(['TlsC','TlsS']), # 100% coverage
|
||||
'name': _blank_or(0.30, random.choice(NODE_NAMES)),
|
||||
'source-node': _blank_or(0.30, random.choice(SRC_NODES)),
|
||||
|
||||
# ── Session resumption (medium coverage) ──
|
||||
'cnrs': _blank_or(0.55, '+'),
|
||||
'isrs': _blank_or(0.50, '+'),
|
||||
|
||||
# ── Database metadata (low-medium coverage) ──
|
||||
'4dbn': _blank_or(0.40, str(random.randint(1, 100))),
|
||||
'8dbd': _blank_or(0.25, str(random.randint(1, 200))),
|
||||
'8did': _blank_or(0.15, str(random.randint(1000, 9999))),
|
||||
'4srs': _blank_or(0.20, str(random.randint(10000, 99999))),
|
||||
|
||||
# ── GeoIP location (rare, ~1% have data, rest '+' = null) ──
|
||||
':ips.latd': _blank_or_plus(0.99, str(src_lat)), # 99% real value, ~1% '+' null
|
||||
':ips.lond': _blank_or_plus(0.99, str(src_lon)),
|
||||
':ipd.latd': _blank_or_plus(0.99, str(dst_lat)),
|
||||
':ipd.lond': _blank_or_plus(0.99, str(dst_lon)),
|
||||
':ips.ispn': _blank_or(0.20, random.choice(ISPS)),
|
||||
':ipd.ispn': _blank_or(0.20, random.choice(ISPS)),
|
||||
':ips.orgn': _blank_or(0.15, random.choice(ORGS)),
|
||||
':ipd.orgn': _blank_or(0.15, random.choice(ORGS)),
|
||||
':ips.city': _blank_or(0.25, random.choice(CITIES)),
|
||||
':ipd.city': _blank_or(0.25, random.choice(CITIES)),
|
||||
':ips.anon': _blank_or(0.05, 'anon, hosting'), # single combined string
|
||||
':ipd.anon': _blank_or(0.05, 'anon, hosting'),
|
||||
|
||||
# ── Two-letter abbreviation columns (rare, ~5%) ──
|
||||
'scnt': _blank_or(0.05, src_country),
|
||||
'dcnt': _blank_or(0.05, dst_country),
|
||||
'crcc': _blank_or(0.05, random.choice(['OK','ER','--'])),
|
||||
'@iph': _blank_or(0.001, random.choice(['A','B','C','D'])), # 1/1000
|
||||
'eiph': _blank_or(0.05, '+'), # boolean: + = True, blank = False
|
||||
|
||||
# ── Obscure / rarely populated fields ──
|
||||
'orga': _blank_or(0.08, random.choice(ORGS)),
|
||||
'orgu': _blank_or(0.08, str(random.randint(1000, 99999))),
|
||||
|
||||
# ── TLS protocol-level (always present) ──
|
||||
}
|
||||
|
||||
return [
|
||||
src_ip, dst_ip, src_port, dst_port,
|
||||
src_country, dst_country,
|
||||
server_ip, client_ip,
|
||||
tls_ver, service, service,
|
||||
duration, sessions, tmo, ks,
|
||||
cnrs, isrs,
|
||||
ack, ppk, dbd,
|
||||
ipp, dbn,
|
||||
tabl, name, src_node,
|
||||
node_name, node_type,
|
||||
CIPHER_NAMES[cipher_idx], CURVE_NAMES[curve_idx],
|
||||
ocph, ocrv,
|
||||
rand_hex_1, rand_hex_2,
|
||||
row_id + 1, time_str,
|
||||
ts,
|
||||
src_lat, src_lon,
|
||||
dst_lat, dst_lon,
|
||||
random.choice(isps), random.choice(isps),
|
||||
random.choice(orgs), random.choice(orgs),
|
||||
random.choice(cities), random.choice(cities),
|
||||
]
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Generate TLS flow test data')
|
||||
parser.add_argument('--rows', type=int, default=1000, help='Number of rows to generate')
|
||||
parser.add_argument('--output', type=str, default='data/test_flows.csv', help='Output CSV path')
|
||||
parser.add_argument('--globe', action='store_true', help='Use GeoIP-compatible IP ranges')
|
||||
parser = argparse.ArgumentParser(description='Generate TLS flow test data matching TlsDB.csv')
|
||||
parser.add_argument('--rows', type=int, default=10000, help='Number of rows')
|
||||
parser.add_argument('--output', type=str, default='data/test_flows.csv')
|
||||
args = parser.parse_args()
|
||||
|
||||
columns = [
|
||||
':ips', ':ipd', ':prs', ':prd',
|
||||
'scnt', 'dcnt', 'server-ip', 'client-ip',
|
||||
'0ver', 'snam', 'cnam',
|
||||
'4dur', '8ses', '2tmo', '4ksz',
|
||||
'cnrs', 'isrs',
|
||||
'8ack', '8ppk', '8dbd',
|
||||
'1ipp', '4dbn',
|
||||
'tabl', 'name', 'source-node',
|
||||
'cnam', 'snam',
|
||||
'cipher-suite', 'ecdhe-named-curve',
|
||||
'0cph', '0crv',
|
||||
'0rnd', '0rnt',
|
||||
'row', 'time', 'timestamp',
|
||||
':ips.latd', ':ips.lond',
|
||||
':ipd.latd', ':ipd.lond',
|
||||
':ips.ispn', ':ipd.ispn',
|
||||
':ips.orgn', ':ipd.orgn',
|
||||
':ips.city', ':ipd.city',
|
||||
]
|
||||
columns = COLUMNS
|
||||
|
||||
base_time = datetime(2024, 1, 5, 12, 0, 0)
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
base_time = datetime(2026, 6, 1, 0, 0, 0)
|
||||
|
||||
blank_counts = {col: 0 for col in columns}
|
||||
total_rows = args.rows
|
||||
|
||||
with open(output_path, 'w', newline='', encoding='utf-8') as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(columns)
|
||||
for i in range(args.rows):
|
||||
row_time = base_time + timedelta(seconds=random.randint(0, 86400 * 7))
|
||||
writer.writerow(generate_row(i, row_time))
|
||||
writer = csv.DictWriter(f, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
for i in range(total_rows):
|
||||
row_time = base_time + timedelta(seconds=random.randint(0, 86400 * 30))
|
||||
row = generate_row(i, row_time)
|
||||
writer.writerow(row)
|
||||
for col in columns:
|
||||
if row.get(col, '') == '' or row.get(col, '') == '+':
|
||||
blank_counts[col] += 1
|
||||
|
||||
print(f'Generated {total_rows} rows → {output_path}')
|
||||
print('\nCoverage report (blank/null rate per column):')
|
||||
for col in columns:
|
||||
rate = blank_counts[col] / total_rows * 100
|
||||
bar_len = int((100 - rate) / 5)
|
||||
bar = '#' * bar_len + '.' * (20 - bar_len)
|
||||
print(f' {col:20s} {bar} {100-rate:5.1f}% present')
|
||||
|
||||
print(f'Generated {args.rows} rows → {output_path}')
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
+106
-103
@@ -142,102 +142,135 @@
|
||||
<div id="autoBar" style="background:#4361ee;height:100%;width:0%;"></div>
|
||||
</div>
|
||||
|
||||
<!-- Thinking panel -->
|
||||
<div id="thinkingPanel" class="thinking-panel">
|
||||
<div class="label">🧠 LLM 思考 <span id="thinkingStep"></span></div>
|
||||
<div id="thinkingContent"></div>
|
||||
</div>
|
||||
|
||||
<button id="cancelBtn" class="btn" style="background:#f72585;color:white;display:none;" onclick="cancelAnalysis()">取消分析</button>
|
||||
<div id="liveLog" style="background:#1a1a2e;color:#e0e0e0;padding:1rem;border-radius:6px;margin-top:1rem;max-height:300px;overflow:auto;font-size:0.8rem;display:none;"></div>
|
||||
|
||||
<!-- Tool call accordion -->
|
||||
<div id="toolCallsPanel" style="display:none;margin-top:1rem;">
|
||||
<h4 style="margin-bottom:0.6rem;color:#333;">🔧 工具调用序列</h4>
|
||||
<div id="toolCallsList"></div>
|
||||
<!-- Unified timeline: interleaved thinking + tool calls -->
|
||||
<div id="timelinePanel" style="display:none;margin-top:1rem;">
|
||||
<h4 style="margin-bottom:0.6rem;color:#333;">📋 分析步骤</h4>
|
||||
<div id="timeline"></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
let statusPoll = null;
|
||||
let logPoll = null;
|
||||
let renderedToolCallCount = 0;
|
||||
|
||||
function toggleToolCall(el) {
|
||||
function escapeHtml(str) {
|
||||
if (str == null) return '';
|
||||
return String(str).replace(/&/g,'&').replace(/</g,'<').replace(/>/g,'>').replace(/"/g,'"');
|
||||
}
|
||||
|
||||
function prettyJson(obj) {
|
||||
try { return JSON.stringify(obj, null, 2); }
|
||||
catch(e) { return String(obj); }
|
||||
}
|
||||
|
||||
function toggleStep(el) {
|
||||
el.classList.toggle('open');
|
||||
const body = el.nextElementSibling;
|
||||
if (body) body.classList.toggle('open');
|
||||
}
|
||||
|
||||
function escapeHtml(str) {
|
||||
if (str == null) return '';
|
||||
return String(str)
|
||||
.replace(/&/g, '&')
|
||||
.replace(/</g, '<')
|
||||
.replace(/>/g, '>')
|
||||
.replace(/"/g, '"');
|
||||
function formatTimestamp() {
|
||||
return new Date().toLocaleTimeString('zh-Hans', {hour:'2-digit',minute:'2-digit',second:'2-digit'});
|
||||
}
|
||||
|
||||
function prettyJson(obj) {
|
||||
try {
|
||||
return JSON.stringify(obj, null, 2);
|
||||
} catch (e) {
|
||||
return String(obj);
|
||||
function buildTimeline(llmThinking, toolCalls) {
|
||||
// Parse llm_thinking into [{step, text}] entries
|
||||
const thoughtEntries = [];
|
||||
if (llmThinking) {
|
||||
const parts = llmThinking.split(/\n(?=\[\d+\])/);
|
||||
for (const part of parts) {
|
||||
const m = part.match(/^\[(\d+)\]\s*(.*)/s);
|
||||
if (m) {
|
||||
thoughtEntries.push({ step: parseInt(m[1]), type: 'thought', text: m[2].trim() });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// tool_calls already have step, type, name, input, output
|
||||
const toolEntries = (toolCalls || []).map(tc => ({
|
||||
step: tc.step,
|
||||
type: 'tool',
|
||||
name: tc.name,
|
||||
input: tc.input,
|
||||
output: tc.output,
|
||||
}));
|
||||
|
||||
// Merge and sort by step
|
||||
const merged = [...thoughtEntries, ...toolEntries].sort((a, b) => a.step - b.step);
|
||||
return merged;
|
||||
}
|
||||
|
||||
function renderToolCalls(toolCalls) {
|
||||
const list = document.getElementById('toolCallsList');
|
||||
if (!toolCalls || !Array.isArray(toolCalls)) return;
|
||||
function renderTimeline(entries) {
|
||||
const container = document.getElementById('timeline');
|
||||
const panel = document.getElementById('timelinePanel');
|
||||
if (!entries || entries.length === 0) return;
|
||||
|
||||
// Only render new items
|
||||
for (let i = renderedToolCallCount; i < toolCalls.length; i++) {
|
||||
const tc = toolCalls[i];
|
||||
const div = document.createElement('div');
|
||||
div.className = 'tool-call-item';
|
||||
panel.style.display = 'block';
|
||||
// Clear and re-render entire timeline each poll
|
||||
container.innerHTML = '';
|
||||
let lastStep = -1;
|
||||
|
||||
const inputStr = prettyJson(tc.input);
|
||||
const outputStr = prettyJson(tc.output);
|
||||
for (const entry of entries) {
|
||||
// Step header
|
||||
if (entry.step !== lastStep) {
|
||||
const header = document.createElement('div');
|
||||
header.style.cssText = 'font-weight:700;font-size:0.85rem;color:#4361ee;padding:0.6rem 0 0.3rem 0;border-top:1px solid #eee;margin-top:0.4rem;';
|
||||
header.textContent = `步骤 ${entry.step}`;
|
||||
container.appendChild(header);
|
||||
lastStep = entry.step;
|
||||
}
|
||||
|
||||
div.innerHTML =
|
||||
'<div class="tool-call-header" onclick="toggleToolCall(this)">' +
|
||||
'<span><span class="step-badge">' + (tc.step != null ? tc.step : '?') + '</span>' +
|
||||
'<span class="tool-name">' + escapeHtml(tc.name || '?') + '</span></span>' +
|
||||
'<span class="arrow">▶</span>' +
|
||||
'</div>' +
|
||||
'<div class="tool-call-body">' +
|
||||
'<div class="section-title">输入 (INPUT):</div>' +
|
||||
'<pre><code>' + escapeHtml(inputStr) + '</code></pre>' +
|
||||
'<div class="section-title">输出 (OUTPUT):</div>' +
|
||||
'<pre><code>' + escapeHtml(outputStr) + '</code></pre>' +
|
||||
'</div>';
|
||||
|
||||
list.appendChild(div);
|
||||
if (entry.type === 'thought') {
|
||||
// LLM reasoning — collapsible
|
||||
const card = document.createElement('div');
|
||||
card.style.cssText = 'border:1px solid #d0d8f0;border-radius:6px;margin-bottom:0.3rem;overflow:hidden;background:#f8f9ff;';
|
||||
card.innerHTML =
|
||||
'<div class="tool-call-header" onclick="toggleStep(this)" style="background:#eef2ff;padding:0.4rem 0.8rem;">' +
|
||||
'<span><span style="background:#4361ee;color:#fff;border-radius:3px;padding:0 5px;font-size:0.7rem;font-weight:600;margin-right:6px;">💭</span>' +
|
||||
'<span style="font-weight:600;font-size:0.8rem;">LLM 推理</span></span>' +
|
||||
'<span class="arrow">▶</span>' +
|
||||
'</div>' +
|
||||
'<div class="tool-call-body" style="padding:0.6rem 0.8rem;font-size:0.78rem;line-height:1.6;color:#333;background:#fff;">' +
|
||||
escapeHtml(entry.text) +
|
||||
'</div>';
|
||||
container.appendChild(card);
|
||||
} else if (entry.type === 'tool') {
|
||||
// Tool call — collapsible with I/O
|
||||
const card = document.createElement('div');
|
||||
card.style.cssText = 'border:1px solid #e0e0e0;border-radius:6px;margin-bottom:0.3rem;overflow:hidden;background:#fff;';
|
||||
const inputStr = prettyJson(entry.input);
|
||||
const outputStr = prettyJson(entry.output);
|
||||
card.innerHTML =
|
||||
'<div class="tool-call-header" onclick="toggleStep(this)" style="padding:0.4rem 0.8rem;">' +
|
||||
'<span><span style="background:#43aa8b;color:#fff;border-radius:3px;padding:0 5px;font-size:0.7rem;font-weight:600;margin-right:6px;">🔧</span>' +
|
||||
'<span style="font-weight:600;font-size:0.8rem;">' + escapeHtml(entry.name || '?') + '</span></span>' +
|
||||
'<span class="arrow">▶</span>' +
|
||||
'</div>' +
|
||||
'<div class="tool-call-body" style="padding:0.5rem 0.8rem;background:#fafbfc;">' +
|
||||
'<div style="font-weight:600;font-size:0.72rem;color:#555;margin-bottom:0.2rem;">输入 (INPUT):</div>' +
|
||||
'<pre style="background:#1a1a2e;color:#e0e0e0;padding:0.5rem;border-radius:4px;font-size:0.7rem;max-height:200px;overflow:auto;margin:0 0 0.3rem 0;"><code>' + escapeHtml(inputStr) + '</code></pre>' +
|
||||
'<div style="font-weight:600;font-size:0.72rem;color:#555;margin-bottom:0.2rem;">输出 (OUTPUT):</div>' +
|
||||
'<pre style="background:#1a1a2e;color:#e0e0e0;padding:0.5rem;border-radius:4px;font-size:0.7rem;max-height:200px;overflow:auto;margin:0;"><code>' + escapeHtml(outputStr) + '</code></pre>' +
|
||||
'</div>';
|
||||
container.appendChild(card);
|
||||
}
|
||||
}
|
||||
renderedToolCallCount = toolCalls.length;
|
||||
}
|
||||
|
||||
async function startAuto() {
|
||||
const selected = document.querySelector('input[name="auto_run_id"]:checked');
|
||||
if (!selected) { showError('请先选择一个数据集'); return; }
|
||||
|
||||
const thinkingPanel = document.getElementById('thinkingPanel');
|
||||
const thinkingContent = document.getElementById('thinkingContent');
|
||||
const toolCallsPanel = document.getElementById('toolCallsPanel');
|
||||
const toolCallsList = document.getElementById('toolCallsList');
|
||||
const liveLog = document.getElementById('liveLog');
|
||||
const timelineContainer = document.getElementById('timeline');
|
||||
const timelinePanel = document.getElementById('timelinePanel');
|
||||
|
||||
// Reset UI
|
||||
document.getElementById('autoProgress').style.display = 'block';
|
||||
document.getElementById('cancelBtn').style.display = 'inline-block';
|
||||
liveLog.style.display = 'block';
|
||||
liveLog.innerHTML = '';
|
||||
thinkingPanel.style.display = 'none';
|
||||
thinkingContent.innerHTML = '';
|
||||
toolCallsPanel.style.display = 'none';
|
||||
toolCallsList.innerHTML = '';
|
||||
renderedToolCallCount = 0;
|
||||
timelinePanel.style.display = 'none';
|
||||
timelineContainer.innerHTML = '';
|
||||
|
||||
let resp;
|
||||
try {
|
||||
@@ -250,6 +283,8 @@ async function startAuto() {
|
||||
const data = await resp.json();
|
||||
const runId = selected.value;
|
||||
|
||||
let lastTimelineHash = '';
|
||||
|
||||
statusPoll = setInterval(async () => {
|
||||
try {
|
||||
const r = await fetch(`/runs/${runId}/status/`);
|
||||
@@ -259,67 +294,35 @@ async function startAuto() {
|
||||
document.getElementById('autoStatus').textContent = `${msg} (${pct}%)`;
|
||||
document.getElementById('autoBar').style.width = pct + '%';
|
||||
|
||||
// Update thinking panel
|
||||
if (s.llm_thinking && s.llm_thinking.trim()) {
|
||||
thinkingPanel.style.display = 'block';
|
||||
thinkingContent.textContent = s.llm_thinking;
|
||||
thinkingPanel.scrollTop = thinkingPanel.scrollHeight;
|
||||
}
|
||||
|
||||
// Update tool calls accordion
|
||||
if (s.tool_calls && Array.isArray(s.tool_calls) && s.tool_calls.length > 0) {
|
||||
toolCallsPanel.style.display = 'block';
|
||||
renderToolCalls(s.tool_calls);
|
||||
}
|
||||
|
||||
// Update live log
|
||||
if (s.run_log) {
|
||||
liveLog.innerHTML = s.run_log.split('\n').map(l => l.trim()).filter(l => l).join('\n') + '\n';
|
||||
liveLog.scrollTop = liveLog.scrollHeight;
|
||||
// Build and render unified timeline
|
||||
const entries = buildTimeline(s.llm_thinking, s.tool_calls);
|
||||
const hash = JSON.stringify(entries.map(e => ({step:e.step,type:e.type,name:e.name,text:(e.text||'').slice(0,40)})));
|
||||
if (hash !== lastTimelineHash) {
|
||||
renderTimeline(entries);
|
||||
lastTimelineHash = hash;
|
||||
}
|
||||
|
||||
if (s.status === 'completed') {
|
||||
clearInterval(statusPoll);
|
||||
clearInterval(logPoll);
|
||||
document.getElementById('autoBar').style.width = '100%';
|
||||
document.getElementById('cancelBtn').style.display = 'none';
|
||||
document.getElementById('autoStatus').textContent = '分析完成！';
|
||||
}
|
||||
if (s.status === 'failed') {
|
||||
clearInterval(statusPoll);
|
||||
clearInterval(logPoll);
|
||||
document.getElementById('cancelBtn').style.display = 'none';
|
||||
document.getElementById('autoStatus').textContent = '失败: ' + (s.error_message || '');
|
||||
}
|
||||
} catch (e) { showError('轮询状态失败: ' + e.message); }
|
||||
}, 1000);
|
||||
|
||||
// Poll live logs every 3 seconds
|
||||
let lastPos = 0;
|
||||
logPoll = setInterval(async () => {
|
||||
try {
|
||||
const r = await fetch(`/logs/?pos=${lastPos}`);
|
||||
const txt = await r.text();
|
||||
const logDiv = document.getElementById('liveLog');
|
||||
const parts = txt.split('\n', 2);
|
||||
if (parts.length >= 2) {
|
||||
const newPos = parseInt(parts[0], 10);
|
||||
if (!isNaN(newPos)) { lastPos = newPos; }
|
||||
if (parts[1]) {
|
||||
logDiv.innerHTML += parts[1] + '\n';
|
||||
logDiv.scrollTop = logDiv.scrollHeight;
|
||||
}
|
||||
}
|
||||
} catch (e) { showError('轮询日志失败: ' + e.message); }
|
||||
}, 3000);
|
||||
}
|
||||
|
||||
function cancelAnalysis() {
|
||||
if (statusPoll) { clearInterval(statusPoll); statusPoll = null; }
|
||||
if (logPoll) { clearInterval(logPoll); logPoll = null; }
|
||||
document.getElementById('autoStatus').textContent = '已取消';
|
||||
document.getElementById('cancelBtn').style.display = 'none';
|
||||
document.getElementById('liveLog').innerHTML += '\n--- 分析已取消 ---\n';
|
||||
document.getElementById('timelinePanel').innerHTML += '\n--- 分析已取消 ---\n';
|
||||
}
|
||||
</script>
|
||||
{% endblock %}
|
||||
<!--VERSION2-->
|
||||
@@ -9,13 +9,14 @@
|
||||
<p style="color:#666;margin-bottom:1rem;">支持 .csv 文件，多文件批量上传，自动解压 .zip</p>
|
||||
<form id="uploadForm" enctype="multipart/form-data">
|
||||
{% csrf_token %}
|
||||
<div style="border:2px dashed #ccc;border-radius:8px;padding:3rem;text-align:center;cursor:pointer;background:#fafafa;" id="dropZone">
|
||||
<div style="border:2px dashed #ccc;border-radius:8px;padding:2rem;text-align:center;background:#fafafa;" id="dropZone">
|
||||
<div style="font-size:2rem;margin-bottom:0.5rem;">📂</div>
|
||||
<p>拖拽 CSV 文件到此处，或点击选择</p>
|
||||
<p style="margin-bottom:1rem;color:#666;">拖拽 CSV 文件到此处</p>
|
||||
<input type="file" name="files" multiple accept=".csv,.zip" style="display:none;" id="fileInput">
|
||||
<div style="margin-top:0.5rem;font-size:0.85rem;color:#888;">
|
||||
<label style="cursor:pointer;color:#4361ee;text-decoration:underline;" id="dirSelectLabel">或选择整个目录</label>
|
||||
<input type="file" webkitdirectory style="display:none;" id="dirInput">
|
||||
<input type="file" directory webkitdirectory style="display:none;" id="dirInput">
|
||||
<div style="display:flex;gap:1rem;justify-content:center;">
|
||||
<button type="button" class="btn btn-primary" id="fileSelectBtn">📄 选择CSV文件</button>
|
||||
<button type="button" class="btn" style="background:#43aa8b;color:#fff;" id="dirSelectBtn">📁 选择整个目录</button>
|
||||
</div>
|
||||
</div>
|
||||
<div id="submitBar" style="margin-top:1rem;display:none;">
|
||||
@@ -67,18 +68,15 @@ const submitBar = document.getElementById('submitBar');
|
||||
const submitBtn = document.getElementById('submitBtn');
|
||||
let selectedFiles = [];
|
||||
|
||||
dropZone.addEventListener('click', () => fileInput.click());
|
||||
dropZone.addEventListener('dragover', e => { e.preventDefault(); dropZone.style.borderColor = '#4361ee'; });
|
||||
dropZone.addEventListener('dragleave', () => dropZone.style.borderColor = '#ccc');
|
||||
dropZone.addEventListener('drop', e => { e.preventDefault(); handleFiles(e.dataTransfer.files); });
|
||||
document.getElementById('fileSelectBtn').addEventListener('click', () => fileInput.click());
|
||||
document.getElementById('dirSelectBtn').addEventListener('click', () => dirInput.click());
|
||||
fileInput.addEventListener('change', () => handleFiles(fileInput.files));
|
||||
document.getElementById('dirInput').addEventListener('change', function(e) {
|
||||
dirInput.addEventListener('change', function(e) {
|
||||
handleFiles(this.files);
|
||||
});
|
||||
document.getElementById('dirSelectLabel').addEventListener('click', function(e) {
|
||||
e.stopPropagation();
|
||||
document.getElementById('dirInput').click();
|
||||
});
|
||||
|
||||
function handleFiles(files) {
|
||||
selectedFiles = [...files];
|
||||
|
||||
@@ -1,367 +0,0 @@
|
||||
"""Tests for analysis.data_loader — BOM detection, schema validation, load/config."""
|
||||
|
||||
import csv
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import polars as pl
|
||||
|
||||
from analysis.data_loader import (
|
||||
_file_count,
|
||||
_merge_schema_entries,
|
||||
_resolve_encoding,
|
||||
detect_bom,
|
||||
load_config,
|
||||
load_csv_directory,
|
||||
validate_schemas,
|
||||
)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# BOM detection
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestDetectBOM:
|
||||
"""detect_bom() should return the correct encoding string for each BOM."""
|
||||
|
||||
def test_utf8_bom(self):
|
||||
"""UTF-8 BOM (EF BB BF) → 'utf-8-sig'."""
|
||||
path = _write_bytes(b'\xef\xbb\xbfcol1,col2\n1,2\n')
|
||||
try:
|
||||
assert detect_bom(path) == 'utf-8-sig'
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_utf16le_bom(self):
|
||||
"""UTF-16 LE BOM (FF FE) → 'utf-16le'."""
|
||||
path = _write_bytes(b'\xff\xfec\x00o\x00l\x001\x00,\x00c\x00o\x00l\x002\x00\n\x001\x00,\x002\x00\n')
|
||||
try:
|
||||
assert detect_bom(path) == 'utf-16le'
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_utf16be_bom(self):
|
||||
"""UTF-16 BE BOM (FE FF) → 'utf-16be'."""
|
||||
path = _write_bytes(b'\xfe\xff\x00c\x00o\x00l\x001\x00,\x00c\x00o\x00l\x002\x00\n\x001\x00,\x002\x00\n')
|
||||
try:
|
||||
assert detect_bom(path) == 'utf-16be'
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_no_bom(self):
|
||||
"""Plain ASCII/UTF-8 without BOM → 'utf-8'."""
|
||||
path = _write_bytes(b'col1,col2\n1,2\n')
|
||||
try:
|
||||
assert detect_bom(path) == 'utf-8'
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_empty_file(self):
|
||||
"""Empty file → 'utf-8'."""
|
||||
path = _write_bytes(b'')
|
||||
try:
|
||||
assert detect_bom(path) == 'utf-8'
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Schema validation
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestValidateSchemas:
|
||||
"""validate_schemas() must raise ValueError when column sets differ."""
|
||||
|
||||
def test_matching_schemas(self):
|
||||
"""Identical schemas → no exception."""
|
||||
schemas = [
|
||||
{':ips': 'Utf8', '8ack': 'Int64'},
|
||||
{':ips': 'Utf8', '8ack': 'Int64'},
|
||||
]
|
||||
validate_schemas(schemas) # should not raise
|
||||
|
||||
def test_empty_list(self):
|
||||
"""Empty list → no exception."""
|
||||
validate_schemas([])
|
||||
|
||||
def test_schema_mismatch_missing_column(self):
|
||||
"""One file missing a column → ValueError."""
|
||||
schemas = [
|
||||
{':ips': 'Utf8', '8ack': 'Int64'},
|
||||
{':ips': 'Utf8'}, # missing '8ack'
|
||||
]
|
||||
with pytest.raises(ValueError, match=r'missing columns'):
|
||||
validate_schemas(schemas)
|
||||
|
||||
def test_schema_mismatch_extra_column(self):
|
||||
"""One file has an extra column → ValueError."""
|
||||
schemas = [
|
||||
{':ips': 'Utf8'},
|
||||
{':ips': 'Utf8', 'extra_col': 'Int64'},
|
||||
]
|
||||
with pytest.raises(ValueError, match=r'extra columns'):
|
||||
validate_schemas(schemas)
|
||||
|
||||
def test_schema_mismatch_both_sides(self):
|
||||
"""Both missing and extra columns → ValueError with both messages."""
|
||||
schemas = [
|
||||
{'a': 'Int64', 'b': 'Utf8'},
|
||||
{'a': 'Int64', 'c': 'Float64'},
|
||||
]
|
||||
with pytest.raises(ValueError):
|
||||
validate_schemas(schemas)
|
||||
|
||||
|
||||
class TestMergeSchemaEntries:
|
||||
"""_merge_schema_entries() should merge without raising."""
|
||||
|
||||
def test_simple_merge(self):
|
||||
merged, conflicts = _merge_schema_entries([
|
||||
{'a': 'Int64', 'b': 'Utf8'},
|
||||
{'a': 'Int64', 'c': 'Float64'},
|
||||
])
|
||||
assert set(merged.keys()) == {'a', 'b', 'c'}
|
||||
assert merged['a'] == 'Int64'
|
||||
|
||||
def test_type_conflict(self):
|
||||
"""Different dtypes for same column → last wins, conflict recorded."""
|
||||
merged, conflicts = _merge_schema_entries([
|
||||
{'val': 'Int64'},
|
||||
{'val': 'Float64'},
|
||||
])
|
||||
assert merged['val'] == 'Float64'
|
||||
assert 'val' in conflicts
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# CSV loading
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestLoadCSVDirectory:
|
||||
"""End-to-end load_csv_directory() tests with temporary CSV files."""
|
||||
|
||||
def test_basic_load(self):
|
||||
"""Single CSV → LazyFrame + correct schema + row/file counts."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / 'test.csv'
|
||||
_write_csv(path, [
|
||||
[':ips', ':ipd', '8ack'],
|
||||
['10.0.0.1', '10.0.0.2', '100'],
|
||||
['10.0.0.1', '10.0.0.3', '200'],
|
||||
['10.0.0.2', '10.0.0.4', '150'],
|
||||
])
|
||||
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||||
str(path), encoding='utf-8'
|
||||
)
|
||||
|
||||
assert isinstance(lf, pl.LazyFrame)
|
||||
assert ':ips' in schema
|
||||
assert ':ipd' in schema
|
||||
assert '8ack' in schema
|
||||
assert file_count == 1
|
||||
assert row_count > 0
|
||||
assert memory_mb >= 0
|
||||
|
||||
# Verify data round-trip
|
||||
df = lf.collect()
|
||||
assert len(df) == 3
|
||||
assert set(df[':ips'].to_list()) == {'10.0.0.1', '10.0.0.2'}
|
||||
|
||||
def test_multiple_files(self):
|
||||
"""Multiple CSVs with same schema → merged LazyFrame."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
p1 = Path(tmpdir) / 'part1.csv'
|
||||
p2 = Path(tmpdir) / 'part2.csv'
|
||||
_write_csv(p1, [['x'], ['1'], ['2']])
|
||||
_write_csv(p2, [['x'], ['3'], ['4']])
|
||||
|
||||
glob_pattern = str(Path(tmpdir) / 'part*.csv')
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||||
glob_pattern, encoding='utf-8'
|
||||
)
|
||||
|
||||
assert file_count == 2
|
||||
df = lf.collect()
|
||||
assert len(df) == 4
|
||||
assert sorted(df['x'].to_list()) == [1, 2, 3, 4]
|
||||
|
||||
def test_no_files(self):
|
||||
"""No matching files → FileNotFoundError."""
|
||||
with pytest.raises(FileNotFoundError, match=r'No files match'):
|
||||
load_csv_directory('__nonexistent_glarg_*.csv')
|
||||
|
||||
def test_schema_mismatch_across_files_strict(self):
|
||||
"""Files with different column sets + schema_strict=True → ValueError."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_write_csv(Path(tmpdir) / 'a.csv', [['x', 'y'], ['1', '2']])
|
||||
_write_csv(Path(tmpdir) / 'b.csv', [['x'], ['3']])
|
||||
|
||||
with pytest.raises(ValueError, match=r'Schema mismatch'):
|
||||
load_csv_directory(str(Path(tmpdir) / '*.csv'), schema_strict=True)
|
||||
|
||||
def test_schema_mismatch_across_files_lenient(self):
|
||||
"""Files with different column sets + schema_strict=False → merge with nulls."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
_write_csv(Path(tmpdir) / 'a.csv', [['x', 'y'], ['1', '2']])
|
||||
_write_csv(Path(tmpdir) / 'b.csv', [['x'], ['3']])
|
||||
|
||||
lf, schema, *_ = load_csv_directory(str(Path(tmpdir) / '*.csv'))
|
||||
df = lf.collect()
|
||||
assert len(df) == 2
|
||||
assert 'x' in df.columns
|
||||
assert 'y' in df.columns
|
||||
# Row from b.csv should have null for 'y'
|
||||
assert df['y'].is_null().sum() == 1
|
||||
|
||||
def test_schema_override(self):
|
||||
"""schema_override forces a specific column dtype."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / 'test.csv'
|
||||
_write_csv(path, [['val'], ['123'], ['456']])
|
||||
|
||||
lf, schema, *_ = load_csv_directory(
|
||||
str(path), encoding='utf-8',
|
||||
schema_override={'val': 'Float64'},
|
||||
)
|
||||
|
||||
df = lf.collect()
|
||||
assert df['val'].dtype == pl.Float64
|
||||
|
||||
def test_custom_delimiter(self):
|
||||
"""Tab-delimited file should parse correctly."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / 'test.tsv'
|
||||
with open(path, 'w', newline='') as f:
|
||||
f.write('a\tb\n1\t2\n3\t4\n')
|
||||
|
||||
lf, *_ = load_csv_directory(str(path), delimiter='\t')
|
||||
df = lf.collect()
|
||||
assert len(df) == 2
|
||||
assert df.columns == ['a', 'b']
|
||||
|
||||
def test_glob_with_subdirectory(self):
|
||||
"""Recursive glob should find files in subdirectories."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
sub = Path(tmpdir) / 'sub'
|
||||
sub.mkdir()
|
||||
_write_csv(sub / 'nested.csv', [['x'], ['42']])
|
||||
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||||
str(Path(tmpdir) / '**' / '*.csv')
|
||||
)
|
||||
assert file_count == 1
|
||||
assert row_count > 0
|
||||
|
||||
def test_latlon_with_plus_and_blank(self):
|
||||
"""CSV with lat/lon columns containing '+', '' loads without error
|
||||
and the column is Float64 after cleaning."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / 'latlon_test.csv'
|
||||
with open(path, 'w', newline='') as f:
|
||||
f.write('src_ip,latitude,longitude\n')
|
||||
f.write('10.0.0.1,+,120.0\n')
|
||||
f.write('10.0.0.2,35.5,\n')
|
||||
f.write('10.0.0.3,-,\n')
|
||||
f.write('10.0.0.4,40.0,121.5\n')
|
||||
|
||||
lf, schema, *_ = load_csv_directory(str(path))
|
||||
df = lf.collect()
|
||||
|
||||
# lat/lon columns should be Float64 in schema
|
||||
assert 'latitude' in schema
|
||||
assert 'longitude' in schema
|
||||
assert 'latitude' in df.columns
|
||||
assert 'longitude' in df.columns
|
||||
|
||||
# Check dtypes are Float64
|
||||
assert df['latitude'].dtype == pl.Float64
|
||||
assert df['longitude'].dtype == pl.Float64
|
||||
|
||||
# Standalone '+' and blank should become null
|
||||
assert df['latitude'].is_null().sum() == 2 # rows 0,2
|
||||
assert df['longitude'].is_null().sum() == 2 # rows 1,2
|
||||
|
||||
# Valid values survive
|
||||
assert df['latitude'].drop_nulls().to_list() == [35.5, 40.0]
|
||||
assert df['longitude'].drop_nulls().to_list() == [120.0, 121.5]
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Config loading
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestLoadConfig:
|
||||
"""load_config() YAML loading and error handling."""
|
||||
|
||||
def test_load_valid_yaml(self):
|
||||
path = _write_yaml('key: value\nnested:\n inner: 42\n')
|
||||
try:
|
||||
config = load_config(path)
|
||||
assert config['key'] == 'value'
|
||||
assert config['nested']['inner'] == 42
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_load_empty_yaml(self):
|
||||
path = _write_yaml('')
|
||||
try:
|
||||
config = load_config(path)
|
||||
assert config == {}
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_config_path_none(self):
|
||||
assert load_config(None) == {}
|
||||
|
||||
def test_config_not_found(self):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
load_config('/__nonexistent__/config.yaml')
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Utility functions
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestResolveEncoding:
|
||||
def test_bom_overrides_default(self):
|
||||
assert _resolve_encoding('utf-8', 'utf-16le') == 'utf-16le'
|
||||
|
||||
def test_no_bom_uses_default(self):
|
||||
# Polars >= 1.0 uses 'utf8' (not 'utf-8')
|
||||
assert _resolve_encoding('utf-8', 'utf-8') == 'utf8'
|
||||
assert _resolve_encoding('latin-1', 'utf-8') == 'latin-1'
|
||||
|
||||
|
||||
class TestFileCount:
|
||||
def test_count(self):
|
||||
assert _file_count(['a', 'b', 'c']) == 3
|
||||
assert _file_count([]) == 0
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _write_csv(path: Path, rows: list[list[str]]) -> None:
|
||||
"""Write rows as a CSV file."""
|
||||
with open(path, 'w', newline='') as f:
|
||||
writer = csv.writer(f)
|
||||
for row in rows:
|
||||
writer.writerow(row)
|
||||
|
||||
|
||||
def _write_bytes(data: bytes) -> str:
|
||||
"""Write raw bytes to a temp file; return the path string."""
|
||||
with tempfile.NamedTemporaryFile(suffix='.csv', delete=False) as f:
|
||||
f.write(data)
|
||||
return f.name
|
||||
|
||||
|
||||
def _write_yaml(content: str) -> str:
|
||||
"""Write YAML string to a temp file; return the path string."""
|
||||
with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w', delete=False) as f:
|
||||
f.write(content)
|
||||
return f.name
|
||||
@@ -1,365 +0,0 @@
|
||||
"""End-to-end integration test for the full TLS analysis pipeline.
|
||||
|
||||
Pipeline
|
||||
1. Create synthetic TLS flow CSV with known entity clusters.
|
||||
2. Load via ``load_csv_directory``.
|
||||
3. Store in ``SessionStore``.
|
||||
4. Auto-detect entity column (``detect_entity_column``).
|
||||
5. Aggregate flows → entity profiles (``aggregate_by_entity``).
|
||||
6. Run HDBSCAN clustering on aggregated numeric features.
|
||||
7. Extract per-cluster distinguishing features.
|
||||
|
||||
All data is synthetic, deterministic (seed=42), and self-contained —
|
||||
no external files or network access required.
|
||||
"""
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
import pytest
|
||||
|
||||
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
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Fixtures
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_store():
|
||||
"""Reset the SessionStore singleton before and after each test."""
|
||||
store = SessionStore()
|
||||
store.drop_all()
|
||||
yield
|
||||
store.drop_all()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Synthetic data generator
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
def create_synthetic_flows(seed: int = 42) -> pl.DataFrame:
|
||||
"""Generate deterministic TLS flow data with three known entity groups.
|
||||
|
||||
Uses the user's exact column naming convention only.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pl.DataFrame
|
||||
~1000 rows with columns.
|
||||
Entities (``:ips``) are partitioned into three behavioural
|
||||
groups that HDBSCAN should recover after aggregation.
|
||||
"""
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
# ── Entity profiles ──────────────────────────────────────────────
|
||||
# Group A: 10 high-volume entities (50 flows each)
|
||||
group_a = [
|
||||
{
|
||||
':ips': f'10.0.0.{i+1}',
|
||||
'n_flows': 50,
|
||||
'8ack_mu': 50_000, '8ack_sigma': 5_000,
|
||||
'4dur_mu': 30.0, '4dur_sigma': 5.0,
|
||||
'dst_range': (100, 150),
|
||||
'0ver': 'TLSv1.3',
|
||||
'cipher_suite': 'TLS_AES_256_GCM_SHA384',
|
||||
}
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
# Group B: 15 medium-volume entities (20 flows each)
|
||||
group_b = [
|
||||
{
|
||||
':ips': f'10.0.0.{i+11}',
|
||||
'n_flows': 20,
|
||||
'8ack_mu': 5_000, '8ack_sigma': 1_000,
|
||||
'4dur_mu': 15.0, '4dur_sigma': 3.0,
|
||||
'dst_range': (150, 200),
|
||||
'0ver': 'TLSv1.2',
|
||||
'cipher_suite': 'TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256',
|
||||
}
|
||||
for i in range(15)
|
||||
]
|
||||
|
||||
# Group C: 25 low-volume entities (6 flows each)
|
||||
group_c = [
|
||||
{
|
||||
':ips': f'10.0.0.{i+26}',
|
||||
'n_flows': 6,
|
||||
'8ack_mu': 500, '8ack_sigma': 100,
|
||||
'4dur_mu': 5.0, '4dur_sigma': 1.0,
|
||||
'dst_range': (200, 250),
|
||||
'0ver': 'TLSv1.2',
|
||||
'cipher_suite': 'TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256',
|
||||
}
|
||||
for i in range(25)
|
||||
]
|
||||
|
||||
entities = group_a + group_b + group_c
|
||||
# Total flows: 10×50 + 15×20 + 25×6 = 500 + 300 + 150 = 950
|
||||
|
||||
# ── Generate flow rows ────────────────────────────────────────────
|
||||
rows: list[dict] = []
|
||||
dst_ports = [80, 443, 8080, 8443]
|
||||
protos = ['TCP', 'UDP']
|
||||
proto_weights = [0.9, 0.1]
|
||||
|
||||
for ent in entities:
|
||||
lo, hi = ent['dst_range']
|
||||
for _ in range(ent['n_flows']):
|
||||
dst_idx = int(rng.integers(lo, hi))
|
||||
row = {
|
||||
':ips': ent[':ips'],
|
||||
':ipd': f"10.0.{dst_idx >> 8}.{dst_idx & 0xff}",
|
||||
':prd': int(rng.choice(dst_ports)),
|
||||
'8ack': int(max(0, rng.normal(ent['8ack_mu'],
|
||||
ent['8ack_sigma']))),
|
||||
'4dur': round(max(0.1, abs(rng.normal(ent['4dur_mu'],
|
||||
ent['4dur_sigma']))), 2),
|
||||
'1ipp': str(rng.choice(protos, p=proto_weights)),
|
||||
'0ver': ent['0ver'],
|
||||
'cipher_suite': ent['cipher_suite'],
|
||||
'snam': f"host{rng.integers(1, 200)}.example.com",
|
||||
'timestamp': (
|
||||
f"2024-06-15 "
|
||||
f"{rng.integers(0, 23):02d}:"
|
||||
f"{rng.integers(0, 59):02d}:"
|
||||
f"{rng.integers(0, 59):02d}"
|
||||
),
|
||||
}
|
||||
rows.append(row)
|
||||
|
||||
return pl.DataFrame(rows)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Pipeline tests
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestFullPipeline:
|
||||
"""End-to-end: synthetic CSV → load → entity detect → aggregate → cluster."""
|
||||
|
||||
def _write_csv(self, df: pl.DataFrame, tmpdir: str) -> str:
|
||||
"""Write a Polars DataFrame to a temp CSV; return the file path."""
|
||||
path = Path(tmpdir) / 'synthetic_flows.csv'
|
||||
df.write_csv(path)
|
||||
return str(path)
|
||||
|
||||
def test_full_pipeline(self):
|
||||
"""Smoke-test the full pipeline — every step produces valid output."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# ── Step 1: Create synthetic data ─────────────────────────────
|
||||
df_flows = create_synthetic_flows(seed=42)
|
||||
assert len(df_flows) > 0, 'Synthetic data should not be empty'
|
||||
|
||||
csv_path = self._write_csv(df_flows, tmpdir)
|
||||
|
||||
# ── Step 2: Load data ─────────────────────────────────────────
|
||||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||||
csv_path, encoding='utf-8',
|
||||
)
|
||||
assert isinstance(lf, pl.LazyFrame), 'Expected LazyFrame'
|
||||
assert file_count == 1
|
||||
assert row_count > 0
|
||||
assert memory_mb >= 0
|
||||
assert ':ips' in schema
|
||||
|
||||
# ── Step 3: Store in session ──────────────────────────────────
|
||||
store = SessionStore()
|
||||
store.store_dataset(
|
||||
'e2e_ds', lf, schema=schema,
|
||||
metadata={'row_count': row_count, 'csv_glob': csv_path},
|
||||
)
|
||||
|
||||
# ── Step 4: Entity detection ──────────────────────────────────
|
||||
det_result = detect_entity_column('e2e_ds')
|
||||
assert 'error' not in det_result, f'detect_entity_column failed: {det_result.get("error")}'
|
||||
# :ips is the catch-all entity column (strings with good unique ratio)
|
||||
assert det_result['recommended'] in (':ips', ':ipd', 'snam'), (
|
||||
f'Expected :ips, :ipd or snam, got {det_result["recommended"]}'
|
||||
)
|
||||
assert len(det_result['candidates']) > 0
|
||||
|
||||
# Verify recommended_multi exists
|
||||
assert 'recommended_multi' in det_result, 'recommended_multi field missing'
|
||||
assert len(det_result['recommended_multi']) > 0, 'recommended_multi is empty'
|
||||
|
||||
# ── Step 5: Entity aggregation ────────────────────────────────
|
||||
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
||||
df_agg = agg_lf.collect()
|
||||
|
||||
# Should be exactly 50 entities (10 + 15 + 25)
|
||||
entity_count = len(df_agg)
|
||||
assert entity_count == 50, f'Expected 50 entities, got {entity_count}'
|
||||
|
||||
# Verify per-entity flow counts (sort by flow_count desc so groups
|
||||
# are contiguous: A=50, B=20, C=6)
|
||||
df_sorted = df_agg.sort('flow_count', descending=True)
|
||||
# Group A: 10 entities × 50 flows
|
||||
for i in range(10):
|
||||
assert df_sorted['flow_count'][i] == 50, (
|
||||
f'Row {i}: expected 50, got {df_sorted["flow_count"][i]}'
|
||||
)
|
||||
# Group B: 15 entities × 20 flows
|
||||
for i in range(10, 25):
|
||||
assert df_sorted['flow_count'][i] == 20, (
|
||||
f'Row {i}: expected 20, got {df_sorted["flow_count"][i]}'
|
||||
)
|
||||
# Group C: 25 entities × 6 flows
|
||||
for i in range(25, 50):
|
||||
assert df_sorted['flow_count'][i] == 6, (
|
||||
f'Row {i}: expected 6, got {df_sorted["flow_count"][i]}'
|
||||
)
|
||||
|
||||
# Core feature columns are present
|
||||
for col in ('flow_count', 'total_bytes_sent',
|
||||
'avg_duration', 'unique_dst_ips', 'unique_dst_ports',
|
||||
'unique_tls_versions', 'unique_protocols',
|
||||
'first_seen', 'last_seen'):
|
||||
assert col in feature_cols, f'Missing feature column: {col}'
|
||||
|
||||
# ── Step 6: Clustering ────────────────────────────────────────
|
||||
numeric_cols = [
|
||||
'flow_count', 'total_bytes_sent',
|
||||
'avg_duration', 'unique_dst_ips', 'unique_dst_ports',
|
||||
]
|
||||
available_cols = [c for c in numeric_cols if c in df_agg.columns]
|
||||
|
||||
data = df_agg.select(available_cols).to_numpy()
|
||||
assert data.shape[0] == 50, 'All 50 entities should have features'
|
||||
|
||||
# Standard-scale + HDBSCAN
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from sklearn.cluster import HDBSCAN
|
||||
|
||||
data_scaled = StandardScaler().fit_transform(data)
|
||||
clusterer = HDBSCAN(min_cluster_size=5, metric='euclidean')
|
||||
labels = clusterer.fit_predict(data_scaled)
|
||||
|
||||
n_clusters = len(set(labels) - {-1})
|
||||
assert n_clusters > 0, (
|
||||
f'HDBSCAN should find ≥1 cluster, found {n_clusters} '
|
||||
f'(labels={set(labels)})'
|
||||
)
|
||||
|
||||
noise_count = int((labels == -1).sum())
|
||||
# At most 40% noise allowed (we have well-separated groups)
|
||||
noise_ratio = noise_count / len(labels)
|
||||
assert noise_ratio < 0.40, (
|
||||
f'Noise ratio too high: {noise_ratio:.2%} '
|
||||
f'({noise_count}/{len(labels)})'
|
||||
)
|
||||
|
||||
# ── Step 7: Feature extraction ────────────────────────────────
|
||||
unique_labels = sorted(set(l for l in labels if l >= 0))
|
||||
assert len(unique_labels) > 0, 'Need at least one valid cluster'
|
||||
|
||||
# Compute per-cluster distinguishing features using z-score
|
||||
features = []
|
||||
for label in unique_labels:
|
||||
mask = labels == label
|
||||
cluster_data = data[mask]
|
||||
for col_idx, col_name in enumerate(available_cols):
|
||||
global_mean = float(data[:, col_idx].mean())
|
||||
global_std = float(data[:, col_idx].std()) or 1.0
|
||||
cluster_mean = float(cluster_data[:, col_idx].mean())
|
||||
z = (cluster_mean - global_mean) / global_std
|
||||
features.append({
|
||||
'cluster': int(label),
|
||||
'feature': col_name,
|
||||
'z_score': round(z, 4),
|
||||
})
|
||||
|
||||
assert len(features) > 0, 'Expected at least one extracted feature'
|
||||
# Each cluster should have features for all numeric columns
|
||||
assert len(features) >= len(unique_labels) * len(available_cols)
|
||||
|
||||
# ── Step 8: Subnet aggregation function ────────────────────────
|
||||
from analysis.entity_aggregator import ip_to_subnet
|
||||
assert callable(ip_to_subnet)
|
||||
assert ip_to_subnet('10.0.1.15', 24) == '10.0.1.0/24'
|
||||
|
||||
# Note: SessionStore is cleaned up by the ``clean_store`` fixture
|
||||
|
||||
def test_synthetic_data_shape(self):
|
||||
"""Verify the synthetic data generator produces expected shapes."""
|
||||
df = create_synthetic_flows(seed=42)
|
||||
assert len(df.columns) == 10, f'Expected 10 columns, got {len(df.columns)}'
|
||||
# Total: 10×50 + 15×20 + 25×6 = 950
|
||||
assert len(df) == 950, f'Expected 950 rows, got {len(df)}'
|
||||
|
||||
# Verify unique entity count
|
||||
n_entities = df[':ips'].n_unique()
|
||||
assert n_entities == 50, f'Expected 50 unique :ips, got {n_entities}'
|
||||
|
||||
# Group A should have TLSv1.3
|
||||
group_a_ips = [f'10.0.0.{i+1}' for i in range(10)]
|
||||
group_a_tls = df.filter(
|
||||
pl.col(':ips').is_in(group_a_ips)
|
||||
)['0ver'].unique().to_list()
|
||||
assert group_a_tls == ['TLSv1.3']
|
||||
|
||||
# Group C should have CHACHA20 cipher
|
||||
group_c_ips = [f'10.0.0.{i+26}' for i in range(25)]
|
||||
group_c_cipher = df.filter(
|
||||
pl.col(':ips').is_in(group_c_ips)
|
||||
)['cipher_suite'].unique().to_list()
|
||||
assert any('CHACHA20' in c for c in group_c_cipher)
|
||||
|
||||
|
||||
class TestPipelineEdgeCases:
|
||||
"""Edge cases for the analysis pipeline."""
|
||||
|
||||
def test_empty_dataset_rejected(self):
|
||||
"""Empty dataset stored → detect_entity_column returns error."""
|
||||
store = SessionStore()
|
||||
empty = pl.DataFrame({':ips': pl.Series([], dtype=pl.Utf8)})
|
||||
schema = {}
|
||||
store.store_dataset(
|
||||
'empty_ds', empty.lazy(), schema=schema,
|
||||
metadata={'row_count': 0},
|
||||
)
|
||||
result = detect_entity_column('empty_ds')
|
||||
# May fail or return empty candidates; either is acceptable
|
||||
# but it should not crash with an unhandled exception.
|
||||
assert 'error' in result or result['recommended'] == ''
|
||||
|
||||
def test_single_entity_clustering(self):
|
||||
"""Single entity → aggregation works but clustering may not split."""
|
||||
data = pl.DataFrame({
|
||||
':ips': ['10.0.0.1'] * 20,
|
||||
':ipd': ['10.0.0.2'] * 20,
|
||||
'8ack': [100] * 20,
|
||||
'4dur': [1.0] * 20,
|
||||
'1ipp': ['TCP'] * 20,
|
||||
'0ver': ['TLSv1.2'] * 20,
|
||||
'cipher_suite': ['AES'] * 20,
|
||||
'snam': ['host1.example.com'] * 20,
|
||||
'timestamp': ['2024-01-01 00:00:00'] * 20,
|
||||
':prd': [443] * 20,
|
||||
})
|
||||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||||
lf = data.lazy()
|
||||
agg_lf, _ = aggregate_by_entity(lf, ':ips', schema)
|
||||
df_agg = agg_lf.collect()
|
||||
|
||||
assert len(df_agg) == 1
|
||||
assert df_agg['flow_count'][0] == 20
|
||||
|
||||
def test_aggregate_preserves_schema_types(self):
|
||||
"""Aggregated columns should have expected Polars dtypes."""
|
||||
data = pl.DataFrame({
|
||||
':ips': ['10.0.0.1'] * 10,
|
||||
'8ack': [100] * 10,
|
||||
'4dur': [1.5] * 10,
|
||||
})
|
||||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||||
lf = data.lazy()
|
||||
agg_lf, _ = aggregate_by_entity(lf, ':ips', schema)
|
||||
df = agg_lf.collect()
|
||||
|
||||
assert df['flow_count'].dtype == pl.UInt32 or df['flow_count'].dtype == pl.Int64
|
||||
@@ -0,0 +1,255 @@
|
||||
"""
|
||||
Full-stack end-to-end test: real DeepSeek LLM auto-analysis workflow.
|
||||
|
||||
1. Start Django server on port 18766, monitor stderr
|
||||
2. Upload ALL CSV files from data/multi_upload/ (0.csv..1997.csv) via POST /upload/csv/
|
||||
3. Poll /runs/{id}/status/ until 'ready' or 'failed'
|
||||
4. POST /analyze/llm/ to run LLM auto-analysis (real DeepSeek LLM)
|
||||
5. Poll /runs/{id}/status/ until 'completed' or 'failed' (LLM may take 2-5 min)
|
||||
6. Check status is 'completed', not 'failed'
|
||||
7. GET /globe/ → 200
|
||||
8. POST /runs/{id}/delete/ → 200
|
||||
9. Stop server
|
||||
10. FAIL on any error → sys.exit(1)
|
||||
|
||||
Usage: runtime/python/python.exe tests/test_e2e_full.py
|
||||
"""
|
||||
import subprocess
|
||||
import time
|
||||
import json
|
||||
import sys
|
||||
import threading
|
||||
import os
|
||||
import uuid
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
import glob
|
||||
|
||||
PORT = 18766
|
||||
BASE = f"http://127.0.0.1:{PORT}"
|
||||
FAILED = []
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def fail(msg):
|
||||
with _lock:
|
||||
FAILED.append(msg)
|
||||
print(f" FAIL: {msg}")
|
||||
|
||||
|
||||
def check(cond, msg):
|
||||
if not cond:
|
||||
fail(msg)
|
||||
|
||||
|
||||
def start_server():
|
||||
manage_py = os.path.join(os.path.dirname(__file__), '..', 'manage.py')
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, manage_py, 'runserver', f'127.0.0.1:{PORT}', '--noreload'],
|
||||
stderr=subprocess.PIPE, stdout=subprocess.PIPE,
|
||||
text=True, encoding='utf-8', errors='replace',
|
||||
)
|
||||
|
||||
def monitor():
|
||||
for line in proc.stderr:
|
||||
stripped = line.strip()
|
||||
upper = stripped.upper()
|
||||
# Django dev server logs request lines like "GET /... 200" to stderr — skip those
|
||||
# Also skip LLM tool diagnostics (tool errors are expected and handled)
|
||||
if ('ERROR' in upper and 'GET' not in upper and 'POST' not in upper
|
||||
and 'tianxuan.llm' not in stripped
|
||||
and 'this error occurred' not in stripped.lower()
|
||||
and 'traceback' not in stripped.lower()):
|
||||
fail(f"STDERR: {stripped}")
|
||||
|
||||
threading.Thread(target=monitor, daemon=True).start()
|
||||
|
||||
# Wait for server to be ready
|
||||
for _ in range(15):
|
||||
try:
|
||||
urllib.request.urlopen(f'{BASE}/', timeout=3)
|
||||
return proc
|
||||
except Exception:
|
||||
time.sleep(1)
|
||||
raise RuntimeError("Server did not start")
|
||||
|
||||
|
||||
def http_get(path):
|
||||
r = urllib.request.urlopen(f'{BASE}{path}', timeout=15)
|
||||
body = r.read()
|
||||
code = r.status
|
||||
r.close()
|
||||
return code, body
|
||||
|
||||
|
||||
def http_post_json(path, data_dict):
|
||||
body = json.dumps(data_dict).encode('utf-8')
|
||||
req = urllib.request.Request(f'{BASE}{path}', data=body,
|
||||
headers={'Content-Type': 'application/json'})
|
||||
try:
|
||||
r = urllib.request.urlopen(req, timeout=30)
|
||||
return r.status, json.loads(r.read())
|
||||
except urllib.error.HTTPError as e:
|
||||
return e.code, json.loads(e.read())
|
||||
|
||||
|
||||
def http_post_raw(path, data=None):
|
||||
"""POST with optional raw data (no content-type override)."""
|
||||
req = urllib.request.Request(f'{BASE}{path}', data=data or b'')
|
||||
try:
|
||||
r = urllib.request.urlopen(req, timeout=30)
|
||||
body = r.read()
|
||||
code = r.status
|
||||
r.close()
|
||||
return code, body
|
||||
except urllib.error.HTTPError as e:
|
||||
return e.code, e.read()
|
||||
|
||||
|
||||
def main():
|
||||
proc = start_server()
|
||||
print("[1/8] Server started")
|
||||
|
||||
# ── Upload ALL 1998 CSV files in a single request ──
|
||||
data_dir = os.path.join(os.path.dirname(__file__), '..', 'data', 'multi_upload')
|
||||
csv_files = sorted(glob.glob(os.path.join(data_dir, '*.csv')))
|
||||
print(f"[2/8] Uploading {len(csv_files)} CSV files (single multipart request)...")
|
||||
|
||||
boundary = uuid.uuid4().hex
|
||||
lines = []
|
||||
for fp in csv_files:
|
||||
fname = os.path.basename(fp)
|
||||
with open(fp, 'rb') as f:
|
||||
filedata = f.read()
|
||||
lines.append(f'--{boundary}'.encode())
|
||||
lines.append(f'Content-Disposition: form-data; name="files"; filename="{fname}"'.encode())
|
||||
lines.append(b'Content-Type: text/csv')
|
||||
lines.append(b'')
|
||||
lines.append(filedata)
|
||||
lines.append(f'--{boundary}--'.encode())
|
||||
body = b'\r\n'.join(lines)
|
||||
ct = f'multipart/form-data; boundary={boundary}'
|
||||
|
||||
req = urllib.request.Request(f'{BASE}/upload/csv/', data=body,
|
||||
headers={'Content-Type': ct})
|
||||
try:
|
||||
r = urllib.request.urlopen(req, timeout=300)
|
||||
resp = json.loads(r.read())
|
||||
r.close()
|
||||
except urllib.error.HTTPError as e:
|
||||
body = e.read().decode('utf-8', errors='replace')
|
||||
fail(f"Upload failed: HTTP {e.code} - {body[:300]}")
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
fail(f"Upload exception: {e}")
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
sys.exit(1)
|
||||
|
||||
run_id = resp.get('run_id')
|
||||
check(run_id is not None, "No run_id in upload response")
|
||||
print(f" Uploaded {len(csv_files)} files, Run ID: {run_id}")
|
||||
|
||||
# ── Poll status until 'ready' ──
|
||||
print("[3/8] Waiting for processing (background: load -> profile -> ready)...")
|
||||
ready = False
|
||||
for i in range(180):
|
||||
code, body = http_get(f'/runs/{run_id}/status/')
|
||||
check(code == 200, f"Status poll -> {code}")
|
||||
s = json.loads(body)
|
||||
status = s['status']
|
||||
if status == 'failed':
|
||||
err = s.get('error_message', '')[:200]
|
||||
fail(f"Run failed: {err}")
|
||||
break
|
||||
if status == 'ready':
|
||||
ready = True
|
||||
print(f" Status: {status} ({s.get('progress_pct', 0)}%)")
|
||||
break
|
||||
if i % 10 == 0:
|
||||
print(f" Poll #{i}: {status} ({s.get('progress_pct', 0)}%)")
|
||||
time.sleep(2)
|
||||
|
||||
check(ready, "Run never became ready after upload")
|
||||
if not ready:
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
sys.exit(1)
|
||||
|
||||
# ── POST /analyze/llm/ to run real LLM auto-analysis ──
|
||||
print("[4/8] Starting LLM auto-analysis (real DeepSeek LLM, may take 2-5 min)...")
|
||||
code, resp = http_post_json('/analyze/llm/', {'run_id': run_id})
|
||||
check(code == 200, f"LLM start -> HTTP {code}")
|
||||
llm_status = resp.get('status')
|
||||
print(f" LLM started: {llm_status}")
|
||||
|
||||
# ── Poll status until 'completed' or 'failed' (LLM phase) ──
|
||||
print("[5/8] Waiting for LLM analysis to complete...")
|
||||
completed = False
|
||||
for i in range(120): # 120 × 3s = 360s = 6 min max
|
||||
code, body = http_get(f'/runs/{run_id}/status/')
|
||||
check(code == 200, f"LLM status poll -> {code}")
|
||||
s = json.loads(body)
|
||||
status = s['status']
|
||||
pct = s.get('progress_pct', 0)
|
||||
msg = s.get('progress_msg', '')
|
||||
if status == 'failed':
|
||||
err = s.get('error_message', '') or s.get('progress_msg', '') or '(no details)'
|
||||
fail(f"LLM analysis failed: {err[:300]}")
|
||||
break
|
||||
if status == 'completed':
|
||||
completed = True
|
||||
print(f" LLM completed ({pct}%)")
|
||||
# Print last 300 chars of run_log for interest
|
||||
run_log = s.get('run_log', '')
|
||||
if run_log:
|
||||
# Strip non-ASCII chars for console-safe printing
|
||||
log_snippet = run_log[-200:].encode('ascii', errors='replace').decode('ascii')
|
||||
print(f" Last log: {log_snippet}")
|
||||
break
|
||||
if i % 5 == 0 or i == 10:
|
||||
print(f" Poll #{i}: {status} ({pct}%) - {msg}")
|
||||
time.sleep(3)
|
||||
|
||||
check(completed, "LLM analysis never completed")
|
||||
if not completed:
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
sys.exit(1)
|
||||
|
||||
# ── Check status is 'completed', not 'failed' ──
|
||||
print("[6/8] Verifying final status...")
|
||||
code, body = http_get(f'/runs/{run_id}/status/')
|
||||
s = json.loads(body) if code == 200 else {}
|
||||
final_status = s.get('status', 'unknown')
|
||||
check(final_status == 'completed', f"Final status is '{final_status}', expected 'completed'")
|
||||
print(f" Final status: {final_status} [OK]")
|
||||
|
||||
# ── GET /globe/ → 200 ──
|
||||
print("[7/8] Checking globe page...")
|
||||
code, body = http_get('/globe/')
|
||||
check(code == 200, f"GET /globe/ -> {code}")
|
||||
print(f" Globe -> HTTP {code} [OK]")
|
||||
|
||||
# ── POST /runs/{run_id}/delete/ → 200 ──
|
||||
print("[8/8] Deleting run...")
|
||||
code, body = http_post_raw(f'/runs/{run_id}/delete/')
|
||||
check(code == 200, f"DELETE -> HTTP {code}")
|
||||
print(f" Delete -> HTTP {code} [OK]")
|
||||
|
||||
# ── Cleanup ──
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
|
||||
if FAILED:
|
||||
print(f"\n!!! INTEGRATION TEST FAILED: {len(FAILED)} failures")
|
||||
for e in FAILED:
|
||||
print(f" {e}")
|
||||
sys.exit(1)
|
||||
print("\n*** ALL INTEGRATION TESTS PASSED ***")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,322 +0,0 @@
|
||||
"""Tests for entity detection (entity_detector) and aggregation (entity_aggregator).
|
||||
|
||||
Every test isolates the SessionStore singleton via the ``clean_store`` fixture.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import polars as pl
|
||||
|
||||
from analysis.session_store import SessionStore
|
||||
from analysis.entity_detector import _score_name, detect_entity_column
|
||||
from analysis.entity_aggregator import (
|
||||
_find_column,
|
||||
aggregate_by_entity,
|
||||
build_aggregation_params,
|
||||
ip_to_subnet,
|
||||
)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Fixtures
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_store():
|
||||
"""Reset the SessionStore singleton before and after each test."""
|
||||
store = SessionStore()
|
||||
store.drop_all()
|
||||
yield
|
||||
store.drop_all()
|
||||
|
||||
|
||||
def _store_dataset(
|
||||
df: pl.DataFrame,
|
||||
dataset_id: str = 'test_ds',
|
||||
) -> str:
|
||||
"""Helper: store a DataFrame as a dataset in the session store.
|
||||
|
||||
Returns the dataset_id.
|
||||
"""
|
||||
store = SessionStore()
|
||||
schema = {name: str(dtype) for name, dtype in zip(df.columns, df.dtypes)}
|
||||
store.store_dataset(
|
||||
dataset_id=dataset_id,
|
||||
lazyframe=df.lazy(),
|
||||
schema=schema,
|
||||
metadata={'row_count': len(df)},
|
||||
)
|
||||
return dataset_id
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _score_name (unit-level)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestScoreName:
|
||||
"""_score_name() scores column names against entity keywords."""
|
||||
|
||||
def test_exact_match(self):
|
||||
"""Exact keyword match → 3.0."""
|
||||
score, kw = _score_name(':ips')
|
||||
assert score == 3.0
|
||||
assert kw == ':ips'
|
||||
|
||||
def test_suffix_match(self):
|
||||
"""Suffix match via _ separator → 3.0."""
|
||||
score, kw = _score_name('my_snam')
|
||||
assert score == 3.0
|
||||
assert kw == 'snam'
|
||||
|
||||
def test_prefix_match(self):
|
||||
"""Prefix match via _ separator → 3.0."""
|
||||
score, kw = _score_name('snam_value')
|
||||
assert score == 3.0
|
||||
assert kw == 'snam'
|
||||
|
||||
def test_substring_not_matched(self):
|
||||
"""Substring-only match now returns 0.0 (exact matching only)."""
|
||||
score, kw = _score_name('myipaddr')
|
||||
assert score == 0.0
|
||||
assert kw == ''
|
||||
|
||||
def test_no_match(self):
|
||||
score, kw = _score_name('random_column_xyz')
|
||||
assert score == 0.0
|
||||
assert kw == ''
|
||||
|
||||
def test_hyphen_normalised(self):
|
||||
"""Hyphen normalised to underscore for matching."""
|
||||
score, kw = _score_name('cipher-suite')
|
||||
assert score == 3.0
|
||||
assert kw == 'cipher_suite'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# detect_entity_column (integration-level)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestDetectEntityColumn:
|
||||
"""detect_entity_column() should find entity columns from stored data."""
|
||||
|
||||
def test_detect_ips(self):
|
||||
"""Dataset with ``:ips`` → recommended is ``:ips``."""
|
||||
data = pl.DataFrame({
|
||||
':ips': [f'10.0.0.{i}' for i in range(1, 51)] * 10,
|
||||
':ipd': [f'10.0.1.{i}' for i in range(1, 501)],
|
||||
'8ack': [100] * 500,
|
||||
})
|
||||
ds_id = _store_dataset(data)
|
||||
result = detect_entity_column(ds_id)
|
||||
|
||||
assert result['recommended'] == ':ips'
|
||||
assert len(result['candidates']) > 0
|
||||
assert result['candidates'][0]['column'] == ':ips'
|
||||
assert result['total_rows'] == 500
|
||||
assert result['scored_columns'] >= 1
|
||||
|
||||
def test_detect_snam(self):
|
||||
"""Datasets with ``snam`` but no IP column → snam is recommended."""
|
||||
data = pl.DataFrame({
|
||||
'snam': [f'host{i}.example.com' for i in range(100)],
|
||||
'4dur': [float(i) for i in range(100)],
|
||||
})
|
||||
ds_id = _store_dataset(data)
|
||||
result = detect_entity_column(ds_id)
|
||||
|
||||
assert result['recommended'] == 'snam'
|
||||
assert result['candidates'][0]['matched_keyword'] == 'snam'
|
||||
|
||||
def test_no_entity_column(self):
|
||||
"""No entity keywords matched → candidates still found but matched_keyword empty."""
|
||||
data = pl.DataFrame({
|
||||
'col_a': [str(i) for i in range(100)],
|
||||
'col_b': [str(i * 2) for i in range(100)],
|
||||
'value_f64': [float(i) for i in range(100)],
|
||||
})
|
||||
ds_id = _store_dataset(data)
|
||||
result = detect_entity_column(ds_id)
|
||||
|
||||
# String columns are candidates (detector scores on uniqueness, not just keywords)
|
||||
assert len(result['candidates']) >= 2
|
||||
assert all(c['matched_keyword'] == '' for c in result['candidates'])
|
||||
# recommended will be the highest-scoring string column
|
||||
assert result['recommended'] in ('col_a', 'col_b')
|
||||
|
||||
def test_all_numeric_no_strings(self):
|
||||
"""No string/categorical columns → no candidates."""
|
||||
data = pl.DataFrame({
|
||||
'value': list(range(100)),
|
||||
'flag': [0, 1] * 50,
|
||||
})
|
||||
ds_id = _store_dataset(data)
|
||||
result = detect_entity_column(ds_id)
|
||||
|
||||
assert result['recommended'] == ''
|
||||
assert len(result['candidates']) == 0
|
||||
|
||||
def test_dataset_not_found(self):
|
||||
"""Unknown dataset_id → error response."""
|
||||
result = detect_entity_column('nonexistent_ds')
|
||||
assert 'error' in result
|
||||
assert result['recommended'] == ''
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _find_column (helper for aggregation)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestFindColumn:
|
||||
"""_find_column() matches keywords against schema column names."""
|
||||
|
||||
ACK_KEYWORDS = ('8ack',)
|
||||
|
||||
def test_exact_match(self):
|
||||
schema = {'8ack': 'Int64', 'other': 'Utf8'}
|
||||
col = _find_column(schema, self.ACK_KEYWORDS)
|
||||
assert col == '8ack'
|
||||
|
||||
def test_no_match(self):
|
||||
schema = {'something_else': 'Int64'}
|
||||
col = _find_column(schema, self.ACK_KEYWORDS)
|
||||
assert col is None
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# aggregate_by_entity
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestAggregateByEntity:
|
||||
"""aggregate_by_entity() groups flows and computes per-entity features."""
|
||||
|
||||
def test_two_entities(self):
|
||||
"""50 flows across 2 entities → 2 aggregated rows with correct counts."""
|
||||
data = pl.DataFrame({
|
||||
':ips': ['10.0.0.1'] * 20 + ['10.0.0.2'] * 30,
|
||||
':ipd': ['10.0.1.1'] * 10 + ['10.0.1.2'] * 10
|
||||
+ ['10.0.1.3'] * 30,
|
||||
':prd': [80] * 20 + [443] * 30,
|
||||
'8ack': [100] * 20 + [200] * 30,
|
||||
'4dur': [1.0] * 20 + [2.0] * 30,
|
||||
'0ver': ['TLSv1.2'] * 20 + ['TLSv1.3'] * 30,
|
||||
'cipher_suite': ['AES'] * 20 + ['CHACHA20'] * 30,
|
||||
'1ipp': ['TCP'] * 50,
|
||||
'timestamp': ['2024-01-01 00:00:00'] * 50,
|
||||
})
|
||||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||||
lf = data.lazy()
|
||||
|
||||
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
||||
df = agg_lf.collect()
|
||||
|
||||
# Two entities
|
||||
assert len(df) == 2
|
||||
|
||||
# Flow counts
|
||||
df_sorted = df.sort(':ips')
|
||||
assert df_sorted['flow_count'][0] == 20 # 10.0.0.1
|
||||
assert df_sorted['flow_count'][1] == 30 # 10.0.0.2
|
||||
|
||||
# Bytes — each column is now perfectly aligned with entity groups
|
||||
assert df_sorted['total_bytes_sent'][0] == 20 * 100
|
||||
assert df_sorted['total_bytes_sent'][1] == 30 * 200
|
||||
|
||||
# Core feature columns present
|
||||
assert 'flow_count' in feature_cols
|
||||
assert 'total_bytes_sent' in feature_cols
|
||||
assert 'avg_duration' in feature_cols
|
||||
assert 'unique_dst_ips' in feature_cols
|
||||
assert 'unique_dst_ports' in feature_cols
|
||||
assert 'unique_protocols' in feature_cols
|
||||
assert 'first_seen' in feature_cols
|
||||
assert 'last_seen' in feature_cols
|
||||
|
||||
def test_single_entity(self):
|
||||
"""Single entity → single aggregated row."""
|
||||
data = pl.DataFrame({
|
||||
':ips': ['10.0.0.1'] * 5,
|
||||
'8ack': [100] * 5,
|
||||
':ipd': ['10.0.0.2'] * 5,
|
||||
})
|
||||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||||
lf = data.lazy()
|
||||
agg_lf, _ = aggregate_by_entity(lf, ':ips', schema)
|
||||
df = agg_lf.collect()
|
||||
assert len(df) == 1
|
||||
assert df['flow_count'][0] == 5
|
||||
|
||||
def test_nonexistent_entity_column(self):
|
||||
"""Unknown entity column → ValueError."""
|
||||
data = pl.DataFrame({'a': [1, 2, 3]})
|
||||
schema = {'a': 'Int64'}
|
||||
lf = data.lazy()
|
||||
with pytest.raises(ValueError, match=r"Entity column\(s\) '\[\'xyz\'\]'"):
|
||||
aggregate_by_entity(lf, 'xyz', schema)
|
||||
|
||||
def test_schema_with_missing_columns(self):
|
||||
"""Aggregation still works when some optional columns are missing."""
|
||||
data = pl.DataFrame({
|
||||
':ips': ['10.0.0.1'] * 3 + ['10.0.0.2'] * 3,
|
||||
':ipd': ['10.0.0.3'] * 6,
|
||||
# No 8ack, 4dur, etc.
|
||||
})
|
||||
schema = {name: str(dtype) for name, dtype in zip(data.columns, data.dtypes)}
|
||||
lf = data.lazy()
|
||||
agg_lf, feature_cols = aggregate_by_entity(lf, ':ips', schema)
|
||||
|
||||
df = agg_lf.collect()
|
||||
assert len(df) == 2
|
||||
assert df['flow_count'][0] == 3
|
||||
# Missing features should have None values
|
||||
assert df['total_bytes_sent'][0] is None
|
||||
assert 'flow_count' in feature_cols
|
||||
assert 'unique_dst_ips' in feature_cols
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# build_aggregation_params
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestBuildAggregationParams:
|
||||
"""build_aggregation_params() builds correct expressions per schema."""
|
||||
|
||||
def test_flow_count_always_present(self):
|
||||
params = build_aggregation_params({'a': 'Int64'})
|
||||
names = [p[0] for p in params]
|
||||
assert 'flow_count' in names
|
||||
|
||||
def test_ack_detected(self):
|
||||
params = build_aggregation_params({'8ack': 'Int64'})
|
||||
names = [p[0] for p in params]
|
||||
assert 'total_bytes_sent' in names
|
||||
|
||||
def test_ack_not_detected(self):
|
||||
params = build_aggregation_params({'x': 'Int64'})
|
||||
param_map = dict(params)
|
||||
assert param_map['total_bytes_sent'] is not None # still has a fallback expr
|
||||
# The value should be a Polars literal (None)
|
||||
names = [p[0] for p in params]
|
||||
assert 'total_bytes_sent' in names
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# ip_to_subnet (IP subnet mask conversion)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestIpToSubnet:
|
||||
"""ip_to_subnet() converts IP strings to subnet prefixes."""
|
||||
|
||||
def test_ip_to_subnet_24(self):
|
||||
"""ip_to_subnet with /24 mask."""
|
||||
assert ip_to_subnet('10.0.1.15', 24) == '10.0.1.0/24'
|
||||
assert ip_to_subnet('192.168.1.255', 24) == '192.168.1.0/24'
|
||||
|
||||
def test_ip_to_subnet_28(self):
|
||||
"""ip_to_subnet with /28 mask."""
|
||||
assert ip_to_subnet('10.0.1.15', 28) == '10.0.1.0/28'
|
||||
assert ip_to_subnet('10.0.1.250', 28) == '10.0.1.240/28'
|
||||
|
||||
def test_ip_to_subnet_invalid(self):
|
||||
"""ip_to_subnet with invalid inputs."""
|
||||
assert ip_to_subnet(None, 24) is None
|
||||
assert ip_to_subnet('', 24) == ''
|
||||
assert ip_to_subnet('not_an_ip', 24) == 'not_an_ip'
|
||||
@@ -1,132 +0,0 @@
|
||||
"""
|
||||
Full-stack integration test. Run: python tests/test_integration.py
|
||||
IMPORTANT: Update this file whenever frontend HTML/API changes.
|
||||
Pure Python — no curl, no requests required.
|
||||
"""
|
||||
import subprocess, time, json, sys, threading, os, uuid, urllib.request, urllib.error
|
||||
|
||||
PORT = 18765
|
||||
BASE = f"http://127.0.0.1:{PORT}"
|
||||
errors = []
|
||||
_lock = threading.Lock()
|
||||
|
||||
def fail(msg):
|
||||
with _lock: errors.append(msg)
|
||||
|
||||
def check(cond, msg):
|
||||
if not cond: fail(msg)
|
||||
|
||||
def start_server():
|
||||
manage_py = os.path.join(os.path.dirname(__file__), '..', 'manage.py')
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, manage_py, 'runserver', f'127.0.0.1:{PORT}', '--noreload'],
|
||||
stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True, encoding='utf-8',
|
||||
)
|
||||
def monitor():
|
||||
for line in proc.stderr:
|
||||
stripped = line.strip()
|
||||
if 'ERROR' in stripped.upper() and 'GET' not in stripped and 'POST' not in stripped:
|
||||
fail(f"STDERR: {stripped}")
|
||||
threading.Thread(target=monitor, daemon=True).start()
|
||||
return proc
|
||||
|
||||
def build_multipart_formdata(fields):
|
||||
"""Build multipart/form-data body and content-type boundary.
|
||||
fields: {name: (filename, content_type, data_bytes)}
|
||||
"""
|
||||
boundary = uuid.uuid4().hex
|
||||
lines = []
|
||||
for name, (filename, content_type, data) in fields.items():
|
||||
lines.append(f'--{boundary}'.encode())
|
||||
lines.append(f'Content-Disposition: form-data; name="{name}"; filename="{filename}"'.encode())
|
||||
lines.append(f'Content-Type: {content_type}'.encode())
|
||||
lines.append(b'')
|
||||
lines.append(data)
|
||||
lines.append(f'--{boundary}--'.encode())
|
||||
return b'\r\n'.join(lines), f'multipart/form-data; boundary={boundary}'
|
||||
|
||||
def http_get(path, max_retries=10):
|
||||
"""GET request with retry on connection refused."""
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
r = urllib.request.urlopen(f'{BASE}{path}', timeout=5)
|
||||
body = r.read()
|
||||
code = r.status
|
||||
r.close()
|
||||
return code, body
|
||||
except urllib.error.URLError as e:
|
||||
if 'refused' in str(e).lower() and attempt < max_retries - 1:
|
||||
time.sleep(1)
|
||||
continue
|
||||
return 502, b'Connection refused'
|
||||
|
||||
def http_post_multipart(path, fields):
|
||||
"""POST multipart form data, returns (status_code, body_bytes)."""
|
||||
body_bytes, content_type = build_multipart_formdata(fields)
|
||||
req = urllib.request.Request(f'{BASE}{path}', data=body_bytes,
|
||||
headers={'Content-Type': content_type})
|
||||
try:
|
||||
r = urllib.request.urlopen(req, timeout=30)
|
||||
return r.status, r.read()
|
||||
except urllib.error.HTTPError as e:
|
||||
return e.code, e.read()
|
||||
|
||||
def main():
|
||||
proc = start_server()
|
||||
time.sleep(3)
|
||||
|
||||
# Test 1: Homepage
|
||||
code, _ = http_get('/')
|
||||
check(code == 200, f"GET / → {code}")
|
||||
|
||||
# Test 2: Upload CSV with '+' in lat/lon
|
||||
csv_content = (
|
||||
":ips,:ipd,0ver,:prd,:prs,:ips.latd,:ips.lond\n"
|
||||
"1.2.3.4,5.6.7.8,0303,443,80,35.0,139.0\n"
|
||||
"9.10.11.12,13.14.15.16,0304,8080,443,+,+\n"
|
||||
)
|
||||
fields = {'files': ('test.csv', 'text/csv', csv_content.encode('utf-8'))}
|
||||
code, body = http_post_multipart('/upload/csv/', fields)
|
||||
check(code == 200, f"POST /upload/csv/ → {code}")
|
||||
data = json.loads(body)
|
||||
run_id = data.get('run_id')
|
||||
check(run_id is not None, "No run_id in upload response")
|
||||
|
||||
# Test 3: Poll run status until ready
|
||||
for _ in range(30):
|
||||
code, body = http_get(f'/runs/{run_id}/status/')
|
||||
check(code == 200, f"GET /runs/{run_id}/status/ → {code}")
|
||||
s = json.loads(body)
|
||||
if s['status'] in ('ready', 'completed', 'failed'):
|
||||
check(s['status'] != 'failed', f"Run failed: {s.get('error_message','')}")
|
||||
break
|
||||
time.sleep(1)
|
||||
|
||||
# Test 4: Run list page
|
||||
code, _ = http_get('/runs/')
|
||||
check(code == 200, f"GET /runs/ → {code}")
|
||||
|
||||
# Test 5: Globe page
|
||||
code, _ = http_get('/globe/')
|
||||
check(code == 200, f"GET /globe/ → {code}")
|
||||
|
||||
# Test 6: Config page
|
||||
code, _ = http_get('/config/')
|
||||
check(code == 200, f"GET /config/ → {code}")
|
||||
|
||||
# Test 7: Delete run — @csrf_exempt means 200 (or 500), not 403, not 404
|
||||
code, _ = http_post_multipart(f'/runs/{run_id}/delete/', {})
|
||||
check(code != 404, f"POST /runs/{run_id}/delete/ → {code} (should not be 404)")
|
||||
|
||||
# Cleanup
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
|
||||
if errors:
|
||||
print("INTEGRATION TEST FAILED:")
|
||||
for e in errors: print(f" {e}")
|
||||
sys.exit(1)
|
||||
print("ALL INTEGRATION TESTS PASSED")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,132 +0,0 @@
|
||||
"""Test the full pipeline on multiple datasets × algorithms.
|
||||
|
||||
Runs 6 parameterized test cases: 3 datasets × 2 algorithms.
|
||||
Uses real CSV files from the ``data/`` directory.
|
||||
"""
|
||||
|
||||
from io import StringIO
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from analysis.management.commands.run_pipeline import Command as PipelineCommand
|
||||
from analysis.models import AnalysisRun
|
||||
from analysis.session_store import SessionStore
|
||||
|
||||
|
||||
class PipelineMultiDatasetTest(TestCase):
|
||||
"""Run the pipeline on 3 datasets × 2 algorithms = 6 test cases."""
|
||||
|
||||
DATASETS = [
|
||||
('data/test_flows.csv', 'simple_500'),
|
||||
('data/e2e_test.csv', 'e2e_500'),
|
||||
('data/complex_test.csv', 'complex_5000'),
|
||||
]
|
||||
ALGORITHMS = ['hdbscan', 'kmeans']
|
||||
|
||||
def setUp(self):
|
||||
"""Clean SessionStore before each test."""
|
||||
SessionStore().drop_all()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean SessionStore after each test."""
|
||||
SessionStore().drop_all()
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
def _run_pipeline(self, csv_glob, algo):
|
||||
"""Run the full pipeline via management command.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple[AnalysisRun | None, str]
|
||||
The latest ``AnalysisRun`` for *csv_glob* and the full log output.
|
||||
"""
|
||||
out = StringIO()
|
||||
cmd = PipelineCommand(stdout=out, stderr=out)
|
||||
cmd.handle(csv_glob=csv_glob, entity_col=None, algo=algo, output=None)
|
||||
run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-id').first()
|
||||
return run, out.getvalue()
|
||||
|
||||
# ── individual test cases: 3 × 2 = 6 ────────────────────────────────
|
||||
|
||||
def test_simple_hdbscan(self):
|
||||
"""simple_500 dataset + hdbscan."""
|
||||
run, log = self._run_pipeline('data/test_flows.csv', 'hdbscan')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
self.assertIsNotNone(run.entity_column)
|
||||
self.assertGreater(run.entity_count, 0)
|
||||
|
||||
def test_simple_kmeans(self):
|
||||
"""simple_500 dataset + kmeans."""
|
||||
run, log = self._run_pipeline('data/test_flows.csv', 'kmeans')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
self.assertGreater(run.entity_count, 0)
|
||||
|
||||
def test_e2e_hdbscan(self):
|
||||
"""e2e_500 dataset + hdbscan."""
|
||||
run, log = self._run_pipeline('data/e2e_test.csv', 'hdbscan')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
|
||||
def test_e2e_kmeans(self):
|
||||
"""e2e_500 dataset + kmeans."""
|
||||
run, log = self._run_pipeline('data/e2e_test.csv', 'kmeans')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
|
||||
def test_complex_hdbscan(self):
|
||||
"""complex_5000 dataset + hdbscan."""
|
||||
run, log = self._run_pipeline('data/complex_test.csv', 'hdbscan')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
|
||||
def test_complex_kmeans(self):
|
||||
"""complex_5000 dataset + kmeans."""
|
||||
run, log = self._run_pipeline('data/complex_test.csv', 'kmeans')
|
||||
self.assertEqual(run.status, 'completed',
|
||||
f'Pipeline failed: {run.error_message}')
|
||||
self.assertGreater(run.cluster_count, 0, 'No clusters found')
|
||||
|
||||
# ── combined parametrised run ────────────────────────────────────────
|
||||
|
||||
def test_all_datasets_all_algorithms(self):
|
||||
"""Every (dataset, algorithm) combination succeeds."""
|
||||
failures = []
|
||||
for csv_glob, label in self.DATASETS:
|
||||
for algo in self.ALGORITHMS:
|
||||
with self.subTest(dataset=label, algo=algo):
|
||||
run, log = self._run_pipeline(csv_glob, algo)
|
||||
if run.status != 'completed':
|
||||
failures.append(f'{label}/{algo}: {run.error_message}')
|
||||
self.assertEqual(
|
||||
run.status, 'completed',
|
||||
f'{label}/{algo} failed: {run.error_message}',
|
||||
)
|
||||
self.assertGreater(
|
||||
run.cluster_count, 0,
|
||||
f'{label}/{algo}: no clusters',
|
||||
)
|
||||
self.assertEqual(len(failures), 0,
|
||||
f'{len(failures)} failures: {failures}')
|
||||
|
||||
# ── edge cases ──────────────────────────────────────────────────────
|
||||
|
||||
def test_invalid_csv_glob(self):
|
||||
"""Non-existent CSV → pipeline reports failure, does not crash."""
|
||||
out = StringIO()
|
||||
cmd = PipelineCommand(stdout=out, stderr=out)
|
||||
cmd.handle(csv_glob='data/nonexistent.csv',
|
||||
entity_col=None, algo='hdbscan', output=None)
|
||||
run = AnalysisRun.objects.filter(
|
||||
csv_glob='data/nonexistent.csv',
|
||||
).order_by('-id').first()
|
||||
self.assertIsNotNone(run)
|
||||
# Should fail gracefully (file not found → no rows → no entity col)
|
||||
self.assertEqual(run.status, 'failed')
|
||||
@@ -1,421 +0,0 @@
|
||||
"""Tests for the data type classifier engine (type_classifier) and IP
|
||||
clustering utilities (ip_clustering).
|
||||
|
||||
Covers: INT, FLOAT, ENUM, HEX, URL, IPv4, BOOL_ENUM detection via both
|
||||
name-based and value-based heuristics, config override, schema-level
|
||||
classification, and ``/24`` subnet grouping.
|
||||
"""
|
||||
import pytest
|
||||
import polars as pl
|
||||
|
||||
from analysis.type_classifier import (
|
||||
DataType,
|
||||
classify_column,
|
||||
classify_schema,
|
||||
HEX_PATTERN,
|
||||
URL_PATTERN,
|
||||
IPv4_PATTERN,
|
||||
BOOL_ENUM_VALUES,
|
||||
LAT_LON_NAMES,
|
||||
)
|
||||
from analysis.ip_clustering import ip_to_int, get_subnet, cluster_ips
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — INT
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyInt:
|
||||
def test_int_series(self):
|
||||
s = pl.Series('bytes_sent', [100, 200, 300])
|
||||
assert classify_column('bytes_sent', s) == DataType.INT
|
||||
|
||||
def test_unsigned_int(self):
|
||||
s = pl.Series('packets', [0, 1, 255, 65535])
|
||||
assert classify_column('packets', s) == DataType.INT
|
||||
|
||||
def test_small_int64(self):
|
||||
s = pl.Series('count', [1, 2, 3], dtype=pl.Int64)
|
||||
assert classify_column('count', s) == DataType.INT
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — FLOAT
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyFloat:
|
||||
def test_float_series(self):
|
||||
s = pl.Series('duration', [1.5, 2.7, 0.3])
|
||||
assert classify_column('duration', s) == DataType.FLOAT
|
||||
|
||||
def test_float_dtype(self):
|
||||
s = pl.Series('score', [0.1, 0.2], dtype=pl.Float32)
|
||||
assert classify_column('score', s) == DataType.FLOAT
|
||||
|
||||
def test_float_with_ints(self):
|
||||
"""Values look like ints but dtype is float → still FLOAT."""
|
||||
s = pl.Series('ratio', [1.0, 2.0, 3.0])
|
||||
assert classify_column('ratio', s) == DataType.FLOAT
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — ENUM (< 20 unique values)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyEnum:
|
||||
def test_few_unique_strings(self):
|
||||
s = pl.Series('tls_version', ['TLSv1.2', 'TLSv1.3', 'TLSv1.2'])
|
||||
assert classify_column('tls_version', s) == DataType.ENUM
|
||||
|
||||
def test_exactly_19_unique(self):
|
||||
vals = [str(i) for i in range(19)]
|
||||
s = pl.Series('version_code', vals * 2)
|
||||
assert classify_column('version_code', s) == DataType.ENUM
|
||||
|
||||
def test_many_unique_strings(self):
|
||||
vals = [f'str_{i}' for i in range(50)]
|
||||
s = pl.Series('desc', vals)
|
||||
assert classify_column('desc', s) == DataType.STRING
|
||||
|
||||
def test_name_version_hint(self):
|
||||
"""Column name contains 'version' → name heuristic returns ENUM."""
|
||||
s = pl.Series('tls_ver', ['1.2', '1.3', '1.2'])
|
||||
assert classify_column('tls_ver', s) == DataType.ENUM
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — HEX ("0e a6 3f" format)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyHex:
|
||||
def test_single_hex_pair(self):
|
||||
s = pl.Series('cipher', ['0e', 'a6', '3f'])
|
||||
assert classify_column('cipher', s) == DataType.HEX
|
||||
|
||||
def test_multiple_hex_pairs(self):
|
||||
s = pl.Series('cipher_suite', ['0e a6 3f', 'ab cd ef 01'])
|
||||
assert classify_column('cipher_suite', s) == DataType.HEX
|
||||
|
||||
def test_name_cipher_hint(self):
|
||||
"""Column name containing 'cipher' triggers name heuristic."""
|
||||
s = pl.Series('tls_cipher', ['ab cd', 'ef 01'])
|
||||
assert classify_column('tls_cipher', s) == DataType.HEX
|
||||
|
||||
def test_name_suite_hint(self):
|
||||
s = pl.Series('suite_id', ['00 01', '00 02'])
|
||||
assert classify_column('suite_id', s) == DataType.HEX
|
||||
|
||||
def test_invalid_hex_not_matched(self):
|
||||
"""Not all values are valid hex pairs → should not be HEX."""
|
||||
s = pl.Series('x', ['0e a6', 'zz yy'])
|
||||
result = classify_column('x', s)
|
||||
assert result != DataType.HEX
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — URL
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyUrl:
|
||||
def test_https(self):
|
||||
s = pl.Series('dst_url', ['https://example.com', 'https://test.org'])
|
||||
assert classify_column('dst_url', s) == DataType.URL
|
||||
|
||||
def test_http(self):
|
||||
s = pl.Series('url', ['http://example.com/path', 'http://test.org'])
|
||||
assert classify_column('url', s) == DataType.URL
|
||||
|
||||
def test_name_url_hint(self):
|
||||
s = pl.Series('request_uri', ['https://api.example.com/v1'])
|
||||
assert classify_column('request_uri', s) == DataType.URL
|
||||
|
||||
def test_not_url(self):
|
||||
"""Plain domain names without http:// should not be URL."""
|
||||
s = pl.Series('domain', ['example.com', 'test.org'])
|
||||
assert classify_column('domain', s) != DataType.URL
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — IPv4
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyIPv4:
|
||||
def test_ip_addresses(self):
|
||||
s = pl.Series('ip', ['192.168.1.1', '10.0.0.1', '172.16.0.1'])
|
||||
assert classify_column('ip', s) == DataType.IPv4
|
||||
|
||||
def test_ip_in_name(self):
|
||||
s = pl.Series('src_ip', ['10.0.0.1', '10.0.0.2'])
|
||||
# Name-based heuristic triggers first
|
||||
assert classify_column('src_ip', s) == DataType.IPv4
|
||||
|
||||
def test_addr_suffix(self):
|
||||
s = pl.Series('destination_addr', ['1.2.3.4', '5.6.7.8'])
|
||||
assert classify_column('destination_addr', s) == DataType.IPv4
|
||||
|
||||
def test_invalid_octet_range(self):
|
||||
"""Octet > 255 → should not match IPv4 via values."""
|
||||
s = pl.Series('x', ['192.168.1.256', '10.0.0.1'])
|
||||
result = classify_column('x', s)
|
||||
assert result != DataType.IPv4 # 256 invalid, may fall through
|
||||
|
||||
def test_ip_plus_normal_strings(self):
|
||||
"""Mix of IPs and non-IPs → not IPv4."""
|
||||
s = pl.Series('x', ['192.168.1.1', 'hello'])
|
||||
result = classify_column('x', s)
|
||||
assert result != DataType.IPv4
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — BOOL_ENUM (+, '', true/false, etc.)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyBoolEnum:
|
||||
def test_plus_and_blank(self):
|
||||
s = pl.Series('flag', ['+', '', '+', ' '])
|
||||
assert classify_column('flag', s) == DataType.BOOL_ENUM
|
||||
|
||||
def test_true_false(self):
|
||||
s = pl.Series('active', ['true', 'false', 'true'])
|
||||
assert classify_column('active', s) == DataType.BOOL_ENUM
|
||||
|
||||
def test_yes_no(self):
|
||||
s = pl.Series('enabled', ['yes', 'no', 'yes'])
|
||||
assert classify_column('enabled', s) == DataType.BOOL_ENUM
|
||||
|
||||
def test_01_strings(self):
|
||||
s = pl.Series('ok', ['0', '1', '0'])
|
||||
assert classify_column('ok', s) == DataType.BOOL_ENUM
|
||||
|
||||
def test_boolean_dtype(self):
|
||||
"""Native pl.Boolean dtype → BOOL_ENUM."""
|
||||
s = pl.Series('flag', [True, False, True])
|
||||
assert classify_column('flag', s) == DataType.BOOL_ENUM
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — LAT_LON (name-based)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyLatLon:
|
||||
def test_lat_column(self):
|
||||
s = pl.Series('lat', [35.0, 36.5])
|
||||
assert classify_column('lat', s) == DataType.LAT_LON
|
||||
|
||||
def test_latitude_column(self):
|
||||
s = pl.Series('latitude', [40.0, 41.0])
|
||||
assert classify_column('latitude', s) == DataType.LAT_LON
|
||||
|
||||
def test_lon_column(self):
|
||||
s = pl.Series('lon', [120.0, 121.5])
|
||||
assert classify_column('lon', s) == DataType.LAT_LON
|
||||
|
||||
def test_longitude_column(self):
|
||||
s = pl.Series('longitude', [122.0, 123.0])
|
||||
assert classify_column('longitude', s) == DataType.LAT_LON
|
||||
|
||||
def test_src_latitude_column(self):
|
||||
"""Column name src_latitude → LAT_LON via prefix heuristic."""
|
||||
s = pl.Series('src_latitude', [35.0, 36.0, 37.0])
|
||||
assert classify_column('src_latitude', s) == DataType.LAT_LON
|
||||
|
||||
def test_dst_longitude_column(self):
|
||||
"""Column name dst_longitude → LAT_LON via prefix heuristic."""
|
||||
s = pl.Series('dst_longitude', [120.0, 121.0, 122.0])
|
||||
assert classify_column('dst_longitude', s) == DataType.LAT_LON
|
||||
|
||||
def test_value_based_latlon_strings(self):
|
||||
"""String column with mix of '+', '', and valid floats → LAT_LON."""
|
||||
s = pl.Series('coord', ['+', '', '12.345', '-45.678', '90.0'])
|
||||
assert classify_column('coord', s) == DataType.LAT_LON
|
||||
|
||||
def test_value_based_latlon_rejects_out_of_range(self):
|
||||
"""String column with values > 180 → not LAT_LON."""
|
||||
s = pl.Series('coord', ['200', '300', '400'])
|
||||
result = classify_column('coord', s)
|
||||
assert result != DataType.LAT_LON
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — Config override (highest priority)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestConfigOverride:
|
||||
def test_config_overrides_value(self):
|
||||
"""config_type='int' forces INT even if values are float."""
|
||||
s = pl.Series('anything', [1.5, 2.5, 3.5])
|
||||
assert classify_column('anything', s, config_type='int') == DataType.INT
|
||||
|
||||
def test_config_overrides_name(self):
|
||||
"""config_type='string' forces STRING even if name matches IP."""
|
||||
s = pl.Series('src_ip', ['192.168.1.1', '10.0.0.1'])
|
||||
assert classify_column('src_ip', s, config_type='string') == DataType.STRING
|
||||
|
||||
def test_config_type_case_insensitive(self):
|
||||
s = pl.Series('x', [1, 2])
|
||||
assert classify_column('x', s, config_type='FLOAT') == DataType.FLOAT
|
||||
|
||||
def test_config_type_trimmed(self):
|
||||
s = pl.Series('x', ['a', 'b'])
|
||||
assert classify_column('x', s, config_type=' enum ') == DataType.ENUM
|
||||
|
||||
def test_unknown_config_type_falls_through(self):
|
||||
"""Unknown config_type string → ignored, falls to name/value."""
|
||||
s = pl.Series('bytes_sent', [1, 2, 3])
|
||||
result = classify_column('bytes_sent', s, config_type='unknown_type')
|
||||
assert result == DataType.INT # value-based
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_schema (full schema classification)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifySchema:
|
||||
def test_classify_all_columns(self):
|
||||
data = {
|
||||
'src_ip': ['10.0.0.1', '10.0.0.2', '10.0.0.3'],
|
||||
'dst_ip': ['10.0.1.1', '10.0.1.2', '10.0.1.3'],
|
||||
'bytes_sent': [100, 200, 300],
|
||||
'duration': [1.5, 2.5, 0.5],
|
||||
'tls_version': ['TLSv1.2', 'TLSv1.3', 'TLSv1.2'],
|
||||
'cipher': ['0e a6 3f', 'ab cd ef', '01 02'],
|
||||
'url': ['https://a.com', 'https://b.com', 'https://c.com'],
|
||||
'lat': [35.0, 36.0, 37.0],
|
||||
'lon': [120.0, 121.0, 122.0],
|
||||
}
|
||||
lf = pl.DataFrame(data).lazy()
|
||||
type_map = classify_schema(lf)
|
||||
|
||||
assert type_map['src_ip'] == DataType.IPv4
|
||||
assert type_map['dst_ip'] == DataType.IPv4
|
||||
assert type_map['bytes_sent'] == DataType.INT
|
||||
assert type_map['duration'] == DataType.FLOAT
|
||||
assert type_map['tls_version'] == DataType.ENUM
|
||||
assert type_map['cipher'] == DataType.HEX
|
||||
assert type_map['url'] == DataType.URL
|
||||
assert type_map['lat'] == DataType.LAT_LON
|
||||
assert type_map['lon'] == DataType.LAT_LON
|
||||
|
||||
def test_classify_with_config_overrides(self):
|
||||
data = {'src_ip': ['10.0.0.1', '10.0.0.2']}
|
||||
lf = pl.DataFrame(data).lazy()
|
||||
overrides = {'src_ip': 'string'}
|
||||
type_map = classify_schema(lf, config_overrides=overrides)
|
||||
assert type_map['src_ip'] == DataType.STRING
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# IP clustering utilities
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestIpToInt:
|
||||
def test_simple(self):
|
||||
assert ip_to_int('192.168.1.1') == 3232235777
|
||||
|
||||
def test_localhost(self):
|
||||
assert ip_to_int('127.0.0.1') == 2130706433
|
||||
|
||||
def test_invalid(self):
|
||||
assert ip_to_int('not_an_ip') is None
|
||||
|
||||
def test_empty_string(self):
|
||||
assert ip_to_int('') is None
|
||||
|
||||
def test_out_of_range(self):
|
||||
assert ip_to_int('999.999.999.999') is None
|
||||
|
||||
|
||||
class TestGetSubnet:
|
||||
def test_24_mask(self):
|
||||
assert get_subnet('192.168.1.5', 24) == '192.168.1.0/24'
|
||||
|
||||
def test_16_mask(self):
|
||||
assert get_subnet('10.1.2.3', 16) == '10.1.0.0/16'
|
||||
|
||||
def test_8_mask(self):
|
||||
assert get_subnet('100.200.1.1', 8) == '100.0.0.0/8'
|
||||
|
||||
def test_invalid_ip(self):
|
||||
assert get_subnet('not_an_ip', 24) is None
|
||||
|
||||
def test_default_mask(self):
|
||||
"""Default mask is /24."""
|
||||
assert get_subnet('10.0.0.15') == '10.0.0.0/24'
|
||||
|
||||
|
||||
class TestClusterIps:
|
||||
def test_single_subnet(self):
|
||||
ips = ['10.0.0.1', '10.0.0.2', '10.0.0.3']
|
||||
groups = cluster_ips(ips, mask=24)
|
||||
assert len(groups) == 1
|
||||
assert '10.0.0.0/24' in groups
|
||||
assert len(groups['10.0.0.0/24']) == 3
|
||||
|
||||
def test_multiple_subnets(self):
|
||||
ips = ['10.0.0.1', '10.0.1.1', '192.168.1.1']
|
||||
groups = cluster_ips(ips, mask=24)
|
||||
assert len(groups) == 3
|
||||
|
||||
def test_invalid_ips_grouped(self):
|
||||
ips = ['10.0.0.1', 'bad_ip', 'also_invalid']
|
||||
groups = cluster_ips(ips, mask=24)
|
||||
assert '10.0.0.0/24' in groups
|
||||
assert '__invalid__' in groups
|
||||
assert len(groups['__invalid__']) == 2
|
||||
|
||||
def test_empty_list(self):
|
||||
groups = cluster_ips([], mask=24)
|
||||
assert groups == {}
|
||||
|
||||
def test_same_subnet_then_16(self):
|
||||
"""10.0.x.x IPs with /16 mask all go to one subnet."""
|
||||
ips = ['10.0.1.1', '10.0.2.2', '10.0.255.255']
|
||||
groups = cluster_ips(ips, mask=16)
|
||||
assert len(groups) == 1
|
||||
assert '10.0.0.0/16' in groups
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — MAC address (should NOT be HEX)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyMac:
|
||||
def test_mac_address_detection(self):
|
||||
"""MAC addresses should NOT be classified as HEX."""
|
||||
s = pl.Series('mac', ['aa:bb:cc:dd:ee:ff', '00:1a:2b:3c:4d:5e'])
|
||||
result = classify_column('mac', s)
|
||||
assert result != DataType.HEX, f'MAC should not be HEX, got {result}'
|
||||
# MAC with unique addresses → STRING (high cardinality)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — PORT values (INT range → ENUM)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyPort:
|
||||
def test_port_values_detection(self):
|
||||
"""Integer values in 1-65535 range should be ENUM."""
|
||||
s = pl.Series('port', ['80', '443', '8080', '8443'])
|
||||
assert classify_column('port', s) == DataType.ENUM
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# classify_column — Value-based IPv4 (name-agnostic)
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TestClassifyValueBased:
|
||||
def test_value_based_ipv4(self):
|
||||
"""IP detection should work from values alone."""
|
||||
s = pl.Series('col', ['10.0.0.1', '192.168.1.1', '8.8.8.8'])
|
||||
assert classify_column('col', s) == DataType.IPv4
|
||||
|
||||
def test_value_based_url(self):
|
||||
"""URL detection should work from values."""
|
||||
s = pl.Series('col', ['https://example.com', 'http://test.org/path'])
|
||||
assert classify_column('col', s) == DataType.URL
|
||||
|
||||
def test_value_first_over_name(self):
|
||||
"""Value pattern should take priority over name heuristic."""
|
||||
# Column named "version" with IP-like values should be IPv4, not ENUM
|
||||
s = pl.Series('version', ['10.0.0.1', '192.168.1.1'])
|
||||
assert classify_column('version', s) == DataType.IPv4
|
||||
@@ -1,430 +0,0 @@
|
||||
"""
|
||||
Integration tests for the 天璇 (TianXuan) incremental update system.
|
||||
|
||||
All tests run in isolated temporary directories and use a mocked APPDATA
|
||||
environment variable, so they never touch the real project data.
|
||||
|
||||
Tests verify the full update lifecycle:
|
||||
1. ZIP build structure (matching build_update.py output)
|
||||
2. File extraction and replacement (matching apply_update.py steps 5-6)
|
||||
3. Config merge with user-edit preservation (matching merge_config.py)
|
||||
4. Version skip / direct-jump updates
|
||||
5. Corrupted package rejection (CRC + SHA256 verification)
|
||||
6. Database migration from project dir to %APPDATA%
|
||||
"""
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Real project root from the test file location
|
||||
PROJECT_DIR = Path(__file__).resolve().parent.parent
|
||||
SCRIPTS_DIR = PROJECT_DIR / 'scripts'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Test helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
def _create_test_project(tmp_dir: Path) -> Path:
|
||||
"""Create a simulated TianXuan project structure within a temp dir."""
|
||||
proj = tmp_dir / 'tianxuan'
|
||||
proj.mkdir()
|
||||
|
||||
# Core modules
|
||||
(proj / 'analysis').mkdir()
|
||||
(proj / 'analysis' / '__init__.py').write_text('# analysis module\n')
|
||||
(proj / 'tianxuan').mkdir()
|
||||
(proj / 'tianxuan' / '__init__.py').write_text('# settings module\n')
|
||||
|
||||
# Config
|
||||
(proj / 'config').mkdir()
|
||||
(proj / 'config' / 'config.yaml').write_text(
|
||||
'server:\n host: 127.0.0.1\n port: 8000\n'
|
||||
'llm:\n api_key: test-key-123\n'
|
||||
)
|
||||
|
||||
# Templates
|
||||
(proj / 'templates').mkdir()
|
||||
(proj / 'templates' / 'base.html').write_text('<html>Old Version</html>\n')
|
||||
|
||||
# Static assets
|
||||
(proj / 'static').mkdir()
|
||||
|
||||
# Scripts directory (placeholder — real scripts live outside temp dir)
|
||||
(proj / 'scripts').mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Version file
|
||||
(proj / 'VERSION').write_text('v1.0.0\n')
|
||||
|
||||
# A file that will be "deleted" in a test update
|
||||
(proj / 'old_module.py').write_text('# This will be removed\n')
|
||||
|
||||
return proj
|
||||
|
||||
|
||||
def _create_update_zip(project_dir: Path,
|
||||
from_ver: str = 'v1.0.0',
|
||||
to_ver: str = 'v1.1.0',
|
||||
add_files: dict | None = None,
|
||||
delete_files: list | None = None,
|
||||
run_migrate: bool = False) -> Path:
|
||||
"""Build a simulated update ZIP matching the build_update.py output format.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
project_dir : Path
|
||||
Directory where the ZIP is written.
|
||||
from_ver, to_ver : str
|
||||
Version range this update covers.
|
||||
add_files : dict[str, str] | None
|
||||
Mapping of relative path → file content for changed/new files.
|
||||
delete_files : list[str] | None
|
||||
Relative paths of files to remove.
|
||||
run_migrate : bool
|
||||
Whether the update advertises a schema migration.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Path
|
||||
Full path to the created update ZIP.
|
||||
"""
|
||||
zip_path = project_dir / f'tianxuan_update_{from_ver}-{to_ver}.zip'
|
||||
|
||||
files = add_files or {'templates/base.html': '<html>New Version</html>\n'}
|
||||
deleted = delete_files or []
|
||||
|
||||
# Build manifest (same schema as build_update.py produces)
|
||||
files_manifest = {}
|
||||
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||||
for rel_path, content in files.items():
|
||||
arcname = f'files/{rel_path}'
|
||||
zf.writestr(arcname, content)
|
||||
sha256 = hashlib.sha256(content.encode()).hexdigest()
|
||||
files_manifest[rel_path] = {
|
||||
'sha256': sha256,
|
||||
'size': len(content),
|
||||
}
|
||||
|
||||
manifest = {
|
||||
'from_version': from_ver,
|
||||
'to_version': to_ver,
|
||||
'built_at': '2026-07-20T12:00:00Z',
|
||||
'files': files_manifest,
|
||||
'deleted': deleted,
|
||||
'run_migrate': run_migrate,
|
||||
'run_import_tlsdb': False,
|
||||
'runtime_changed': False,
|
||||
}
|
||||
zf.writestr('version.json', json.dumps(manifest, indent=2))
|
||||
|
||||
return zip_path
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Test class
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestUpdateSystem:
|
||||
"""Integration tests for the TianXuan incremental update system.
|
||||
|
||||
Every test creates its own isolated project tree + APPDATA directory
|
||||
and cleans up on teardown.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
"""Create a fresh, isolated test environment before each test."""
|
||||
self._tmp_base = Path(tempfile.mkdtemp(suffix='_tianxuan_test'))
|
||||
self.project_dir = _create_test_project(self._tmp_base)
|
||||
|
||||
# Override APPDATA so scripts write to the temp directory
|
||||
self._orig_appdata = os.environ.get('APPDATA')
|
||||
self._test_appdata = self._tmp_base / 'appdata'
|
||||
self._test_appdata.mkdir()
|
||||
os.environ['APPDATA'] = str(self._test_appdata)
|
||||
|
||||
def teardown_method(self):
|
||||
"""Restore the environment and remove temp files."""
|
||||
if self._orig_appdata:
|
||||
os.environ['APPDATA'] = self._orig_appdata
|
||||
else:
|
||||
os.environ.pop('APPDATA', None)
|
||||
shutil.rmtree(str(self._tmp_base), ignore_errors=True)
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 1: Build → Apply simple code change
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_simple_code_change(self):
|
||||
"""A single file is updated — verify extraction + replacement."""
|
||||
update_zip = _create_update_zip(
|
||||
self.project_dir,
|
||||
from_ver='v1.0.0', to_ver='v1.1.0',
|
||||
add_files={'templates/base.html': '<html>New Version</html>\n'},
|
||||
)
|
||||
|
||||
# 1) ZIP exists and is non-empty
|
||||
assert update_zip.exists()
|
||||
assert update_zip.stat().st_size > 0
|
||||
|
||||
# 2) Extract and apply (simulating apply_update.py steps 1, 6)
|
||||
with zipfile.ZipFile(update_zip, 'r') as zf:
|
||||
manifest = json.loads(zf.read('version.json'))
|
||||
assert manifest['from_version'] == 'v1.0.0'
|
||||
assert manifest['to_version'] == 'v1.1.0'
|
||||
assert 'templates/base.html' in manifest['files']
|
||||
|
||||
for name in zf.namelist():
|
||||
if not name.startswith('files/'):
|
||||
continue
|
||||
rel_path = name[6:] # strip 'files/' prefix
|
||||
if rel_path == 'config/config.yaml':
|
||||
continue # config is handled by the merge step
|
||||
zf.extract(name, str(self.project_dir))
|
||||
src = self.project_dir / name
|
||||
dst = self.project_dir / rel_path
|
||||
if src != dst:
|
||||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.move(str(src), str(dst))
|
||||
|
||||
# 3) The updated file has the new content
|
||||
updated = self.project_dir / 'templates' / 'base.html'
|
||||
assert updated.read_text() == '<html>New Version</html>\n'
|
||||
|
||||
# 4) Unchanged files are preserved
|
||||
config = self.project_dir / 'config' / 'config.yaml'
|
||||
assert 'test-key-123' in config.read_text()
|
||||
|
||||
# 5) Clean up the temporary files/ directory
|
||||
files_dir = self.project_dir / 'files'
|
||||
if files_dir.exists():
|
||||
shutil.rmtree(str(files_dir))
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 2: File deletion
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_file_deletion(self):
|
||||
"""A file listed in ``deleted`` is removed from the project tree."""
|
||||
old_file = self.project_dir / 'old_module.py'
|
||||
assert old_file.exists()
|
||||
|
||||
update_zip = _create_update_zip(
|
||||
self.project_dir,
|
||||
from_ver='v1.0.0', to_ver='v1.1.0',
|
||||
add_files={},
|
||||
delete_files=['old_module.py'],
|
||||
)
|
||||
|
||||
# Apply deletion (simulating apply_update.py step 5)
|
||||
with zipfile.ZipFile(update_zip, 'r') as zf:
|
||||
manifest = json.loads(zf.read('version.json'))
|
||||
for rel_path in manifest.get('deleted', []):
|
||||
target = self.project_dir / rel_path
|
||||
if target.exists():
|
||||
target.unlink()
|
||||
|
||||
assert not old_file.exists(), 'Deleted file should no longer exist'
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 3: Config protection during update
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_config_preservation(self):
|
||||
"""User edits to config.yaml survive an update; new keys are added.
|
||||
|
||||
This exercises the real ``merge_config.py`` script via subprocess.
|
||||
"""
|
||||
# -- Simulate a user who has edited their config.yml -----------
|
||||
config_path = self.project_dir / 'config' / 'config.yaml'
|
||||
user_edited = (
|
||||
'server:\n'
|
||||
' host: 0.0.0.0\n'
|
||||
' port: 9000\n'
|
||||
'llm:\n'
|
||||
' api_key: my-custom-key\n'
|
||||
'clustering:\n'
|
||||
' algorithm: kmeans\n'
|
||||
)
|
||||
config_path.write_text(user_edited)
|
||||
|
||||
# -- Build an update that ships a *different* config -----------
|
||||
source_config = (
|
||||
'server:\n'
|
||||
' host: 127.0.0.1\n'
|
||||
' port: 8000\n'
|
||||
'llm:\n'
|
||||
' api_key: default-key\n'
|
||||
' model: gpt-4\n' # NEW key — should appear after merge
|
||||
'clustering:\n'
|
||||
' algorithm: hdbscan\n'
|
||||
' min_cluster_size: 10\n' # NEW key — should appear after merge
|
||||
)
|
||||
update_zip = _create_update_zip(
|
||||
self.project_dir,
|
||||
add_files={'config/config.yaml': source_config},
|
||||
)
|
||||
|
||||
# -- Back up user config (apply_update.py step 4) --------------
|
||||
backup_dir = self._test_appdata / 'TianXuan' / 'backups'
|
||||
backup_dir.mkdir(parents=True, exist_ok=True)
|
||||
backup_config = backup_dir / 'v1.1.0' / 'config.yaml'
|
||||
backup_config.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(str(config_path), str(backup_config))
|
||||
|
||||
# -- Extract the new config to ``files/`` (step 6) -------------
|
||||
with zipfile.ZipFile(update_zip, 'r') as zf:
|
||||
zf.extract('files/config/config.yaml', str(self.project_dir))
|
||||
extracted_config = self.project_dir / 'files' / 'config' / 'config.yaml'
|
||||
assert extracted_config.exists()
|
||||
|
||||
# -- Merge (step 7): old = user backup, new = extracted source --
|
||||
merge_script = SCRIPTS_DIR / 'merge_config.py'
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(merge_script),
|
||||
'--old', str(backup_config),
|
||||
'--new', str(extracted_config),
|
||||
'--output', str(config_path)],
|
||||
capture_output=True, text=True,
|
||||
cwd=str(PROJECT_DIR),
|
||||
)
|
||||
assert result.returncode == 0, f'merge_config failed: {result.stderr}'
|
||||
|
||||
merged = config_path.read_text()
|
||||
|
||||
# User edits preserved
|
||||
assert 'my-custom-key' in merged, 'User API key should survive'
|
||||
assert '0.0.0.0' in merged, 'User host should survive'
|
||||
assert 'kmeans' in merged, 'User algorithm should survive'
|
||||
|
||||
# New keys from the update added
|
||||
assert 'model: gpt-4' in merged, 'New key from source should appear'
|
||||
assert 'min_cluster_size: 10' in merged, 'New nested key should appear'
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 4: Version skip
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_version_skip(self):
|
||||
"""A direct v1.0→v1.2 bundle includes ALL intermediate changes.
|
||||
|
||||
``git diff v1.0..HEAD`` shows every change regardless of whether
|
||||
intermediate tags exist, so the build script naturally produces
|
||||
correct "skipped-version" packages. This test verifies the
|
||||
manifest accurately describes the jump.
|
||||
"""
|
||||
update_zip = _create_update_zip(
|
||||
self.project_dir,
|
||||
from_ver='v1.0.0', to_ver='v1.2.0',
|
||||
add_files={
|
||||
'templates/base.html': '<html>v1.2</html>\n',
|
||||
'analysis/new_feature.py': '# Feature added in v1.2\n',
|
||||
},
|
||||
delete_files=['old_module.py'],
|
||||
)
|
||||
|
||||
with zipfile.ZipFile(update_zip, 'r') as zf:
|
||||
manifest = json.loads(zf.read('version.json'))
|
||||
|
||||
assert manifest['from_version'] == 'v1.0.0'
|
||||
assert manifest['to_version'] == 'v1.2.0'
|
||||
assert len(manifest['files']) == 2
|
||||
assert len(manifest['deleted']) == 1
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 5: Corrupted package rejection
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_corrupted_package_rejection(self):
|
||||
"""CRC corruption inside a ZIP is caught by ``testzip()``.
|
||||
|
||||
A file's stored content is modified in-place within the ZIP
|
||||
binary, causing a CRC32 mismatch that ``zipfile.testzip()``
|
||||
reliably detects. The update should be rejected before any
|
||||
files are extracted.
|
||||
"""
|
||||
# Use ZIP_STORED (no compression) so the content appears as
|
||||
# plain text in the archive and can be surgically corrupted.
|
||||
content = b'<html>OK</html>\n'
|
||||
h = hashlib.sha256(content).hexdigest()
|
||||
|
||||
manifest = {
|
||||
'from_version': 'v1.0.0',
|
||||
'to_version': 'v1.1.0',
|
||||
'built_at': '2026-07-20T12:00:00Z',
|
||||
'files': {
|
||||
'templates/base.html': {
|
||||
'sha256': h,
|
||||
'size': len(content),
|
||||
},
|
||||
},
|
||||
'deleted': [],
|
||||
'run_migrate': False,
|
||||
'run_import_tlsdb': False,
|
||||
'runtime_changed': False,
|
||||
}
|
||||
|
||||
zip_path = self.project_dir / 'corrupt_test.zip'
|
||||
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_STORED) as zf:
|
||||
zf.writestr('version.json', json.dumps(manifest, indent=2))
|
||||
zf.writestr('files/templates/base.html', content)
|
||||
|
||||
# Corrupt: replace the stored content in-place
|
||||
raw = bytearray(zip_path.read_bytes())
|
||||
idx = raw.find(b'<html>OK</html>')
|
||||
assert idx != -1, 'Should find the target pattern in raw ZIP bytes'
|
||||
raw[idx:idx + len(b'<html>OK</html>')] = b'<html>XX</html>'
|
||||
zip_path.write_bytes(bytes(raw))
|
||||
|
||||
# The CRC check in apply_update.py step 1 must fail
|
||||
with zipfile.ZipFile(zip_path, 'r') as zf:
|
||||
bad_file = zf.testzip()
|
||||
assert bad_file is not None, (
|
||||
'Corrupted ZIP should be detected by testzip()'
|
||||
)
|
||||
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# Test 6: Database migration from project dir to %APPDATA%
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_db_migration_script(self):
|
||||
"""The migration script copies ``db.sqlite3`` to ``%APPDATA%/TianXuan/``.
|
||||
|
||||
Because ``migrate_db_to_appdata.py`` locates its project root
|
||||
from ``__file__`` (the real TianXuan project where ``db.sqlite3``
|
||||
exists), running it with ``APPDATA`` pointing to the temp
|
||||
directory copies the real database into the isolated test tree
|
||||
without modifying any real project files.
|
||||
"""
|
||||
# The real project has a db.sqlite3 that the script will find.
|
||||
migrate_script = SCRIPTS_DIR / 'migrate_db_to_appdata.py'
|
||||
|
||||
# First run — should copy the database
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(migrate_script)],
|
||||
capture_output=True, text=True,
|
||||
cwd=str(PROJECT_DIR),
|
||||
)
|
||||
assert result.returncode == 0, f'Migration failed: {result.stderr}'
|
||||
|
||||
expected_new = self._test_appdata / 'TianXuan' / 'db.sqlite3'
|
||||
assert expected_new.exists(), 'DB should exist at APPDATA location'
|
||||
assert expected_new.stat().st_size > 0, 'Copied DB should not be empty'
|
||||
|
||||
# Second run — should SKIP (already present)
|
||||
result2 = subprocess.run(
|
||||
[sys.executable, str(migrate_script)],
|
||||
capture_output=True, text=True,
|
||||
cwd=str(PROJECT_DIR),
|
||||
)
|
||||
assert result2.returncode == 0
|
||||
assert 'SKIP' in result2.stdout, 'Second run should print SKIP'
|
||||
@@ -218,7 +218,7 @@ def run_llm_pipeline(dataset_id: str, entity_col: str = "",
|
||||
called_tools: set[str] = set()
|
||||
consecutive_errors = 0
|
||||
|
||||
for step in range(15):
|
||||
for step in range(25):
|
||||
resp = call_llm(msgs, config)
|
||||
if resp is None:
|
||||
if callback:
|
||||
@@ -311,5 +311,5 @@ def run_llm_pipeline(dataset_id: str, entity_col: str = "",
|
||||
})
|
||||
|
||||
if callback:
|
||||
callback(15, "max_steps", None)
|
||||
return {"error": "Reached max 15 tool calls", "steps": 15}
|
||||
callback(25, "max_steps", None)
|
||||
return {"error": "Reached max 25 tool calls", "steps": 25}
|
||||
|
||||
@@ -109,6 +109,7 @@ USE_I18N = True
|
||||
USE_TZ = True
|
||||
|
||||
DATA_UPLOAD_MAX_NUMBER_FIELDS = 10000000
|
||||
DATA_UPLOAD_MAX_NUMBER_FILES = 10000
|
||||
DATA_UPLOAD_MAX_MEMORY_SIZE = 256 * 1024 * 1024
|
||||
FILE_UPLOAD_MAX_MEMORY_SIZE = 100 * 1024 * 1024
|
||||
|
||||
|
||||
Reference in New Issue
Block a user