tianxuan/scripts/gen_complex_test.py

"""生成复杂 TLS 合成测试数据（含经纬度、MAC地址、大量空值）"""
import polars as pl
import numpy as np
import random
from pathlib import Path

RNG = np.random.RandomState(42)

CLIENTS = [f'192.168.{RNG.randint(1,255)}.{RNG.randint(1,255)}' for _ in range(80)]
SERVERS = (
    [f'10.0.{RNG.randint(1,255)}.{RNG.randint(1,255)}' for _ in range(80)]
    + [f'203.0.113.{RNG.randint(1,255)}' for _ in range(30)]
)
DOMAINS = ['mail.example.com', 'api.example.org', 'cdn.example.net', 'login.example.org',
           'storage.example.net', 'video.example.com', 'auth.example.com', 'pay.example.org']
TLS_VERSIONS = ['TLSv1.2', 'TLSv1.3', 'TLSv1.2', 'TLSv1.3', 'TLSv1.1', None]
CIPHERS = ['TLS_AES_128_GCM_SHA256', 'TLS_AES_256_GCM_SHA384',
           'TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256', None]
CERT_STATUS = ['valid', 'valid', 'valid', 'expired', 'self_signed', None]

def gen_mac():
    return ':'.join(f'{RNG.randint(0,255):02x}' for _ in range(6))

CITIES = [
    ('Beijing', 39.9, 116.4), ('Shanghai', 31.2, 121.5), ('Tokyo', 35.7, 139.7),
    ('Singapore', 1.35, 103.8), ('London', 51.5, -0.13), ('Frankfurt', 50.1, 8.68),
    ('NewYork', 40.7, -74.0), ('SanFrancisco', 37.8, -122.4),
    ('Sydney', -33.9, 151.2), ('Mumbai', 19.1, 72.9),
]

def gen_flows(n_rows=5000):
    from datetime import datetime, timedelta
    base_ts = datetime(2025, 1, 1)
    rows = []
    for _ in range(n_rows):
        src_ip = RNG.choice(CLIENTS)
        dst_ip = RNG.choice(SERVERS)
        si = RNG.randint(0, len(CITIES))
        di = RNG.randint(0, len(CITIES))
        src_city, src_lat, src_lon = CITIES[si]
        dst_city, dst_lat, dst_lon = CITIES[di]
        same_city = src_city == dst_city

        bs = int(max(0, RNG.lognormal(8, 2))) if RNG.random() > 0.05 else None
        br = int(max(0, bs * RNG.uniform(0.1, 2))) if bs and RNG.random() > 0.05 else None
        dur = round(RNG.exponential(20), 2) if RNG.random() > 0.08 else None
        hok = RNG.choice([True, True, True, False, None])
        trusted = hok and RNG.choice([True, True, False, None]) if hok else None
        tmo = int(RNG.exponential(500)) + 100 if hok else RNG.choice([None, RNG.randint(5000, 30000)])

        rows.append({
            'timestamp': (base_ts + timedelta(seconds=RNG.randint(0, 31536000))).strftime('%Y-%m-%d %H:%M:%S'),
            'src_ip': src_ip, 'dst_ip': dst_ip,
            'src_mac': gen_mac(), 'dst_mac': gen_mac(),
            'src_url': f'https://{RNG.choice(DOMAINS)}{RNG.choice(["/","/login","/api/v1/users","/api/v1/data","/favicon.ico"])}',
            'dst_url': f'https://{RNG.choice(DOMAINS)}{RNG.choice(["/","/login","/api/v1/users","/api/v1/data"])}',
            'sni': RNG.choice(DOMAINS + [None, None]),
            'src_city': src_city, 'src_latitude': src_lat, 'src_longitude': src_lon,
            'dst_city': dst_city if not same_city else None,
            'dst_latitude': dst_lat if not same_city else None,
            'dst_longitude': dst_lon if not same_city else None,
            'bytes_sent': bs, 'bytes_rev': br, 'duration': dur,
            'tls_version': RNG.choice(TLS_VERSIONS),
            'cipher': RNG.choice(CIPHERS),
            'dst_port': RNG.choice([80, 443, 8080, 8443, 53, 993, 5223, None]),
            'cert_status': RNG.choice(CERT_STATUS),
            'handshake_ok': hok, 'trusted': trusted, 'timeout_ms': tmo,
        })
    df = pl.DataFrame(rows)
    nulls = df.null_count()
    for col in df.columns:
        n = int(nulls[col].item())
        if n > 0:
            print(f'  {col}: {n}/{len(df)} 缺失 ({n*100//len(df)}%)')
    print(f'总行数: {len(df)}, 列数: {len(df.columns)}')
    return df

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--rows', type=int, default=5000)
    parser.add_argument('--output', type=str, default='data/complex_test.csv')
    parser.add_argument('--missing-cols', action='store_true',
                        help='Generate 2-3 files with 20-40%% columns randomly missing (varied per file)')
    parser.add_argument('--missing-col-pct', type=str, default='20-40',
                        help='Range of columns to drop (e.g. "20-40" or "30")')
    args = parser.parse_args()
    df = gen_flows(args.rows)
    all_columns = df.columns

    if args.missing_cols:
        # Parse missing-col-pct range
        _pct_str = args.missing_col_pct
        if '-' in _pct_str:
            _parts = _pct_str.split('-')
            _min_pct = float(_parts[0]) / 100.0
            _max_pct = float(_parts[1]) / 100.0
        else:
            _min_pct = _max_pct = float(_pct_str) / 100.0

        # Columns that must NOT be dropped (IP/port identity columns)
        _keep_cols = {'src_ip', 'dst_ip', 'dst_port'}
        _droppable = [c for c in all_columns if c not in _keep_cols]

        # Split rows into 2-3 files
        total_rows = len(df)
        num_files = random.randint(2, 3)
        _split_sizes = [total_rows // num_files] * num_files
        _split_sizes[-1] += total_rows % num_files

        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        _row_offset = 0
        for _fi in range(num_files):
            _n_rows = _split_sizes[_fi]
            if _n_rows == 0:
                continue

            # Randomly drop 20-40% of droppable columns for this file
            _drop_ratio = random.uniform(_min_pct, _max_pct)
            _n_drop = max(1, int(len(_droppable) * _drop_ratio))
            _dropped = set(random.sample(_droppable, _n_drop))
            _file_columns = [c for c in all_columns if c not in _dropped]

            # Slice the dataframe and select only kept columns
            _chunk = df.slice(_row_offset, _n_rows).select(_file_columns)

            _stem = output_path.stem
            _out = output_path.parent / f'{_stem}_part{_fi + 1}.csv'
            _chunk.write_csv(_out, null_value='')
            _row_offset += _n_rows

            print(f'Generated {_n_rows} rows → {_out} ({len(_file_columns)} columns, dropped {_n_drop}/{len(_droppable)} droppable)')
            print(f'  Kept: {sorted(_keep_cols & set(_file_columns))}')
            if len(_dropped) <= 10:
                print(f'  Dropped: {sorted(_dropped)}')
            else:
                print(f'  Dropped: {sorted(_dropped)[:10]}...')

        print(f'\nTotal: {total_rows} rows across {num_files} files with varied column sets.')
    else:
        # Standard mode: single file, all columns
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        df.write_csv(args.output, null_value='')
        print(f'已生成: {args.output} ({output_path.stat().st_size/1024/1024:.1f} MB)')