PM-pinou
95 lines
4.0 KiB
Python
95 lines
4.0 KiB
Python
"""Batch generator: ~100 CSV files with varied column subsets for multi-upload testing.
|
||
|
||
Each file: 100 rows, 20-40% columns randomly dropped (IP/port identity columns retained).
|
||
Total: ~100 files × 100 rows = ~10000 rows.
|
||
|
||
Usage:
|
||
runtime\\python\\python.exe scripts\\gen_multi.py [--files N] [--rows N]
|
||
"""
|
||
import sys
|
||
import os
|
||
import random
|
||
import csv
|
||
from pathlib import Path
|
||
from datetime import datetime, timedelta
|
||
|
||
# Add scripts/ to path so we can import gen_test_data
|
||
sys.path.insert(0, str(Path(__file__).parent))
|
||
|
||
# Reuse gen_test_data's data pools and generation logic
|
||
from gen_test_data import (
|
||
COLUMNS, COVERAGE, generate_row, _rand_ip, _rand_hex, _maybe, _blank_or, _blank_or_plus,
|
||
TLS_VERSIONS, TLS_WEIGHTS, CIPHER_HEX, CIPHER_NAMES, NAMED_CURVES_HEX, NAMED_CURVES_TEXT,
|
||
SRC_IPS, DST_IPS, PORTS, COUNTRIES, ISPS, ORGS, CITIES, SERVICES, SRC_NODES, NODE_NAMES,
|
||
)
|
||
|
||
|
||
def main():
|
||
import argparse
|
||
parser = argparse.ArgumentParser(description='Generate multi-file test CSV dataset')
|
||
parser.add_argument('--files', type=int, default=100, help='Number of files to generate')
|
||
parser.add_argument('--rows', type=int, default=100, help='Rows per file')
|
||
parser.add_argument('--start', type=int, default=0, help='Start file index (for resume)')
|
||
args = parser.parse_args()
|
||
|
||
# Use a different seed from gen_test_data so IP pools differ
|
||
random.seed(99)
|
||
|
||
out_dir = Path('data/multi_upload')
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Identity columns that must never be dropped
|
||
_keep_cols = {':ips', ':ipd', ':prs', ':prd', 'server-ip', 'client-ip'}
|
||
_droppable = [c for c in COLUMNS if c not in _keep_cols]
|
||
|
||
base_time = datetime(2026, 6, 1, 0, 0, 0)
|
||
# Resume support: start from --start, estimate prior rows
|
||
total_rows = args.start * args.rows if args.start > 0 else 0
|
||
|
||
for fi in range(args.start, args.files):
|
||
# Coverage-based column dropping: columns with lower coverage are more likely dropped.
|
||
# Drop a column with probability (1 - coverage), capped to [0.10, 0.90] range.
|
||
dropped = set()
|
||
for col in _droppable:
|
||
cov = COVERAGE.get(col, 1.0)
|
||
drop_prob = 1.0 - cov
|
||
# Clamp to ensure some columns always have a chance to be dropped/kept
|
||
drop_prob = max(0.10, min(0.90, drop_prob))
|
||
if random.random() < drop_prob:
|
||
dropped.add(col)
|
||
# Ensure at least 1 column is dropped and not ALL droppable columns
|
||
if len(dropped) < 3:
|
||
# Add a few extra random drops to ensure variety
|
||
extra = random.sample([c for c in _droppable if c not in dropped],
|
||
k=min(3 - len(dropped), len(_droppable) - len(dropped)))
|
||
dropped.update(extra)
|
||
if len(dropped) == len(_droppable):
|
||
# Keep at least 3 droppable columns
|
||
spared = set(random.sample(list(dropped), k=min(3, len(dropped))))
|
||
dropped -= spared
|
||
file_columns = [c for c in COLUMNS if c not in dropped]
|
||
|
||
out_path = out_dir / f'{fi}.csv'
|
||
n_rows = random.randint(max(50, args.rows - 10), args.rows + 10)
|
||
|
||
with open(out_path, 'w', newline='', encoding='utf-8') as f:
|
||
writer = csv.DictWriter(f, fieldnames=file_columns, extrasaction='ignore')
|
||
writer.writeheader()
|
||
for row_i in range(n_rows):
|
||
row_time = base_time + timedelta(seconds=random.randint(0, 86400 * 30))
|
||
row = generate_row(total_rows + row_i, row_time)
|
||
filtered = {k: v for k, v in row.items() if k in file_columns}
|
||
writer.writerow(filtered)
|
||
|
||
total_rows += n_rows
|
||
|
||
if (fi + 1) % 20 == 0:
|
||
print(f' [{fi + 1}/{args.files}] {total_rows} rows generated so far...')
|
||
|
||
print(f'\nDone: {args.files} files, {total_rows} total rows → {out_dir}')
|
||
print(f'Each file: ~{args.rows} rows, {len(_droppable)} droppable columns, coverage-based dropping per file')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|