PM-pinou
187 lines
9.0 KiB
Python
187 lines
9.0 KiB
Python
"""Generate test CSVs with alternating high/low column coverage.
|
|
Output files: 0.csv .. N-1.csv in specified directory.
|
|
Columns that are all-blank in a CSV are dropped from that file.
|
|
Usage:
|
|
runtime\python\python.exe scripts\gen_coverage_test.py --files 2500 --rows 10000 --output-dir ../test_data
|
|
"""
|
|
import argparse, csv, os, random, sys
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
# ── Column spec with coverage profiles ──
|
|
# HIGH = present in 95-100% of files, MEDIUM = 40-60%, LOW = 5-15%
|
|
COLUMNS = [
|
|
# HIGH coverage — always present
|
|
(':ips', 1.0), (':ipd', 1.0), (':prs', 1.0), (':prd', 1.0),
|
|
('scnt', 1.0), ('dcnt', 1.0), ('tabl', 1.0), ('name', 1.0),
|
|
('source-node', 1.0), ('server-ip', 1.0), ('client-ip', 1.0),
|
|
('timestamp', 1.0), ('time', 1.0),
|
|
('1ipp', 1.0), ('4dbn', 1.0),
|
|
# MEDIUM-HIGH coverage
|
|
(':ips.latd', 0.85), (':ips.lond', 0.85),
|
|
(':ipd.latd', 0.85), (':ipd.lond', 0.85),
|
|
(':ips.ispn', 0.80), (':ipd.ispn', 0.80),
|
|
(':ips.orgn', 0.70), (':ipd.orgn', 0.70),
|
|
(':ips.city', 0.85), (':ipd.city', 0.85),
|
|
('4dur', 0.95), ('8seq', 0.95), ('2tmo', 0.90),
|
|
('8ack', 0.95), ('8pak', 0.90), ('8byt', 0.85), ('8ppk', 0.80),
|
|
('8ses', 0.80), ('8did', 0.85), ('4srs', 0.80),
|
|
# MEDIUM coverage — present in about half the files
|
|
('0rnd', 0.55), ('0rnt', 0.55), ('0ver', 0.60),
|
|
('snam', 0.50), ('cnam', 0.35),
|
|
('0cph', 0.45), ('0crv', 0.35),
|
|
('cnrs', 0.50), ('isrs', 0.50),
|
|
# LOW coverage — rarely present
|
|
(':ips.anon', 0.12), (':ipd.anon', 0.12),
|
|
(':ips.doma', 0.15), (':ipd.doma', 0.15),
|
|
('4ksz', 0.20), ('cipher-suite', 0.15),
|
|
('ecdhe-named-curve', 0.12),
|
|
('crcc', 0.06), ('orga', 0.06), ('orgu', 0.05),
|
|
('eiph', 0.02), ('@iph', 0.02),
|
|
]
|
|
|
|
def generate_row(row_idx: int, base_time: datetime, file_rng: random.Random,
|
|
present_cols: set) -> dict:
|
|
"""Generate a single row of TLS flow data."""
|
|
row = {}
|
|
src_octets = file_rng.randint(1, 223)
|
|
src_ip = f'{src_octets}.{file_rng.randint(0,255)}.{file_rng.randint(0,255)}.{file_rng.randint(1,254)}'
|
|
dst_octets = file_rng.randint(1, 223)
|
|
dst_ip = f'{dst_octets}.{file_rng.randint(0,255)}.{file_rng.randint(0,255)}.{file_rng.randint(1,254)}'
|
|
ts_val = base_time.timestamp() + file_rng.randint(0, 86400 * 7) + row_idx * 0.1
|
|
|
|
for col, _ in COLUMNS:
|
|
if col not in present_cols:
|
|
continue
|
|
v = ''
|
|
if col in (':ips', 'server-ip'):
|
|
v = src_ip
|
|
elif col in (':ipd', 'client-ip'):
|
|
v = dst_ip
|
|
elif col == ':prs':
|
|
v = str(file_rng.randint(1024, 65535))
|
|
elif col == ':prd':
|
|
v = str(file_rng.choice([80, 443, 8080, 8443, 53, 22, 25, 993, 3306, 6379]))
|
|
elif col in ('scnt', 'dcnt'):
|
|
v = file_rng.choice(['CN', 'US', 'JP', 'KR', 'GB', 'DE', 'SG', 'HK', 'AU', 'RU'])
|
|
elif col == 'tabl':
|
|
v = file_rng.choice(['TlsC', 'TlsS'])
|
|
elif col == 'name':
|
|
v = file_rng.choice(['流量_01', '流量_02', '流量_03', '流量_04', '流量_05'])
|
|
elif col == 'source-node':
|
|
v = file_rng.choice(['node-a', 'node-b', 'node-c', 'node-d'])
|
|
elif col in ('4dur', '8ses', '2tmo'):
|
|
v = f'{file_rng.uniform(1, 300):.2f}' if file_rng.random() < 0.95 else ''
|
|
elif col in ('8ack', '8pak', '8byt', '8ppk', '8did', '4srs'):
|
|
v = str(file_rng.randint(50, 100000))
|
|
elif col == '8seq':
|
|
v = f'{file_rng.uniform(1, 1000):.2f}' if file_rng.random() < 0.95 else ''
|
|
elif col in ('timestamp',):
|
|
v = f'{ts_val:.2f}'
|
|
elif col == 'time':
|
|
dt = datetime.fromtimestamp(ts_val)
|
|
v = dt.strftime('%Y-%m-%d-%H-%M-%S')
|
|
elif col == '1ipp':
|
|
v = str(file_rng.choice([6, 17, 1, 58]))
|
|
elif col == '4dbn':
|
|
v = str(file_rng.randint(1, 200))
|
|
elif col in (':ips.latd', ':ipd.latd'):
|
|
v = f'{file_rng.uniform(-90, 90):.4f}' if file_rng.random() < 0.98 else ''
|
|
elif col in (':ips.lond', ':ipd.lond'):
|
|
v = f'{file_rng.uniform(-180, 180):.4f}' if file_rng.random() < 0.98 else ''
|
|
elif col in (':ips.ispn', ':ipd.ispn'):
|
|
v = file_rng.choices(['电信', '联通', '移动', 'AWS', 'Azure', 'GCP', 'Cloudflare', ''], weights=[0.25,0.2,0.15,0.1,0.1,0.05,0.05,0.1], k=1)[0]
|
|
elif col in (':ips.orgn', ':ipd.orgn'):
|
|
v = file_rng.choices(['阿里巴巴', '腾讯', '华为', 'Amazon', 'Microsoft', 'Google', ''], weights=[0.2,0.15,0.1,0.15,0.1,0.05,0.25], k=1)[0]
|
|
elif col in (':ips.city', ':ipd.city'):
|
|
v = file_rng.choice(['北京', '上海', '深圳', '广州', '杭州', '东京', '新加坡', '伦敦', '纽约', '硅谷', '首尔', '悉尼', '法兰克福', ''])
|
|
elif col in (':ips.anon', ':ipd.anon'):
|
|
v = file_rng.choices(['', '+'], weights=[0.87, 0.13], k=1)[0]
|
|
elif col in (':ips.doma', ':ipd.doma'):
|
|
v = file_rng.choice(['example.com', 'api.example.com', 'cdn.example.net', 'db.internal', ''])
|
|
elif col in ('cnrs', 'isrs'):
|
|
v = file_rng.choices(['', '+', ''], weights=[0.5, 0.2, 0.3], k=1)[0]
|
|
elif col in ('0ver',):
|
|
v = file_rng.choices(['0303', '0304', '0302', ''], weights=[0.6, 0.3, 0.05, 0.05], k=1)[0]
|
|
elif col in ('snam',):
|
|
v = file_rng.choice(['www.example.com', 'api.server.com', 'cdn.cdn.net', 'mail.host.com', 'db.internal.net', ''])
|
|
elif col in ('cnam',):
|
|
v = file_rng.choice(['*.example.com', '*.server.com', 'www', ''])
|
|
elif col in ('0rnd',):
|
|
v = ''.join(file_rng.choices('0123456789abcdef', k=28)) if file_rng.random() < 0.99 else ''
|
|
elif col in ('0rnt',):
|
|
v = f'{file_rng.uniform(1, 2**32):.0f}' if file_rng.random() < 0.99 else ''
|
|
elif col in ('4ksz',):
|
|
v = str(file_rng.choice([128, 256, 1024, 2048, 4096]))
|
|
elif col in ('0cph', '0crv'):
|
|
v = file_rng.choice(['0e a6 3f 2b', 'ab cd ef 01', '12 34 56 78', '00 00 00 00', ''])
|
|
elif col in ('cipher-suite',):
|
|
v = file_rng.choice(['TLS_AES_128_GCM_SHA256', 'TLS_AES_256_GCM_SHA384', 'TLS_CHACHA20_POLY1305_SHA256', ''])
|
|
elif col in ('ecdhe-named-curve',):
|
|
v = file_rng.choice(['X25519', 'secp256r1', 'secp384r1', ''])
|
|
elif col in ('crcc',):
|
|
v = file_rng.choices(['AB', 'CD', 'EF', ''], weights=[0.4, 0.3, 0.1, 0.2], k=1)[0]
|
|
elif col in ('orga', 'orgu'):
|
|
v = file_rng.choice(['ORG_A', 'ORG_B', 'ORG_C', 'ORG_D', ''])
|
|
elif col in ('eiph', '@iph'):
|
|
v = file_rng.choices(['', '+', '', ''], weights=[0.02, 0.02, 0.48, 0.48], k=1)[0]
|
|
row[col] = v
|
|
return row
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Generate test CSVs with alternating column coverage')
|
|
parser.add_argument('--files', type=int, default=100, help='Number of CSV files')
|
|
parser.add_argument('--rows', type=int, default=10000, help='Rows per file')
|
|
parser.add_argument('--output-dir', type=str, default='../test_data', help='Output directory')
|
|
parser.add_argument('--seed', type=int, default=42, help='Random seed')
|
|
args = parser.parse_args()
|
|
|
|
out_dir = Path(args.output_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
base_time = datetime(2026, 1, 1, 0, 0, 0)
|
|
total_rows = 0
|
|
col_names = [c[0] for c in COLUMNS]
|
|
|
|
for file_idx in range(args.files):
|
|
file_rng = random.Random(args.seed + file_idx * 7)
|
|
# Decide per-file which LOW/MEDIUM columns are present
|
|
present = set()
|
|
for col, cov in COLUMNS:
|
|
if file_rng.random() < cov:
|
|
present.add(col)
|
|
|
|
# Generate rows, track which columns actually have data
|
|
non_blank = {c: False for c in present}
|
|
rows = []
|
|
for row_i in range(args.rows):
|
|
row = generate_row(total_rows + row_i, base_time, file_rng, present)
|
|
rows.append(row)
|
|
for col in present:
|
|
if row.get(col, '') and not non_blank[col]:
|
|
non_blank[col] = True
|
|
|
|
# Determine final column list: only columns that have at least some data
|
|
final_cols = [c for c in col_names if non_blank.get(c, False)]
|
|
if not final_cols:
|
|
final_cols = [c for c in present] # fallback
|
|
|
|
# Write file
|
|
out_path = out_dir / f'{file_idx}.csv'
|
|
with open(out_path, 'w', newline='', encoding='utf-8') as f:
|
|
writer = csv.DictWriter(f, fieldnames=final_cols, extrasaction='ignore')
|
|
writer.writeheader()
|
|
for row in rows:
|
|
writer.writerow(row)
|
|
|
|
total_rows += args.rows
|
|
if (file_idx + 1) % 100 == 0 or file_idx == 0 or file_idx == args.files - 1:
|
|
print(f'[{file_idx+1}/{args.files}] {total_rows} rows, {len(final_cols)} cols, file={file_idx}.csv', flush=True)
|
|
|
|
print(f'\nDone: {args.files} files x {args.rows} rows = {total_rows} total rows')
|
|
print(f'Output: {out_dir.resolve()}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|