tianxuan/scripts/gen_multi2.py

"""Topology-aware test data generator — exact gen_test_data format.

Strategy:
1. Generate rows using gen_test_data.generate_row() (EXACT same format/coverages)
2. Then remap IPs: when snam/cnam has a value, assign :ipd from a fixed server pool
3. Result: same format as gen_test_data, but with realistic network topology

Usage:
    runtime\\python\\python.exe scripts\\gen_multi2.py [--files N] [--rows N] [--seed S]
"""

import sys
import random
import csv
from pathlib import Path
from datetime import datetime, timedelta
from collections import defaultdict

sys.path.insert(0, str(Path(__file__).parent))

from gen_test_data import (
    COLUMNS, COVERAGE,
    generate_row as _base_row,
    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,
    _rand_ip, _rand_hex, _maybe, _blank_or, _blank_or_plus,
)

# ════════════════════════════════
#  Server IP pools per service
# ════════════════════════════════
# Each snam/cnam maps to 1-3 specific server IPs.
# When a row has a given snam/cnam, its :ipd is replaced with one of these IPs.
# This creates realistic topology — servers have high in-degree.

# CDN: 3 edge nodes in different regions (Akamai/Cloudflare ranges)
CDN_IPS  = ['23.0.5.10', '104.16.100.10', '151.101.1.10']
API_IPS  = ['52.84.100.20', '54.230.100.20']       # Amazon CloudFront
MAIL_IPS = ['74.125.200.26', '66.102.1.26']          # Google mail
AUTH_IPS = ['40.126.32.100']                          # Microsoft Azure
SOCIAL_IPS = ['157.240.1.35']                         # Facebook
STORAGE_IPS = ['52.216.100.50']                       # Amazon S3
STREAM_IPS = ['52.84.200.50', '54.192.200.50']        # Amazon CloudFront
SEARCH_IPS = ['142.250.80.4']                         # Google

SNAM_TO_IPS = {
    'cdn.example.com':      CDN_IPS,
    '*.cloudfront.net':     CDN_IPS,
    'api.example.com':      API_IPS,
    'mail.example.com':     MAIL_IPS,
    'auth.example.net':     AUTH_IPS,
    'login.live.com':       AUTH_IPS,
    'graph.facebook.com':   SOCIAL_IPS,
    '*.s3.amazonaws.com':   STORAGE_IPS,
    'stream.example.org':   STREAM_IPS,
    'www.google.com':       SEARCH_IPS,
}

CNAM_TO_IPS = {
    'cdn.example.com':      CDN_IPS,
    '*.cloudfront.net':     CDN_IPS,
    'api.example.com':      API_IPS,
    'mail.example.com':     MAIL_IPS,
    'auth.example.net':     AUTH_IPS,
    'login.live.com':       AUTH_IPS,
    'graph.facebook.com':   SOCIAL_IPS,
    '*.s3.amazonaws.com':   STORAGE_IPS,
    'stream.example.org':   STREAM_IPS,
    'www.google.com':       SEARCH_IPS,
}

ALL_SERVER_IPS = list(set(
    ip for pool in SNAM_TO_IPS.values() for ip in pool
))

# Proxy IPs — public IPs resolvable by GeoIP MMDB
# MUST be public IPs that GeoIP can resolve (no private ranges!)
PROXY_IPS = [
    '1.1.1.1',       # Cloudflare DNS — US-West
    '8.8.8.8',       # Google DNS — US-West
    '223.5.5.5',     # AliDNS — CN
    '208.67.222.222', # OpenDNS — US-East
    '77.88.8.8',     # Yandex DNS — EU
]


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--files', type=int, default=200)
    parser.add_argument('--rows', type=int, default=10000)
    parser.add_argument('--seed', type=int, default=2026)
    parser.add_argument('--output-dir', type=str, default='data/gen_multi2')
    args = parser.parse_args()

    random.seed(args.seed)
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    _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)
    total_rows = 0
    degree_src = defaultdict(int)
    degree_dst = defaultdict(int)
    snam_ip_count = defaultdict(lambda: defaultdict(int))
    cnam_ip_count = defaultdict(lambda: defaultdict(int))

    print(f'Generating {args.files} files × {args.rows} rows → {out_dir}')

    for fi in range(args.files):
        # ── Column dropping (matching gen_multi.py exactly) ──
        dropped = set()
        for col in _droppable:
            cov = COVERAGE.get(col, 1.0)
            drop_prob = max(0.10, min(0.90, 1.0 - cov))
            if random.random() < drop_prob:
                dropped.add(col)
        if len(dropped) < 3:
            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):
            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]

        # ── Generate rows with base format ──
        file_rows = []
        file_base = base_time + timedelta(seconds=random.randint(0, 86400 * 30))

        for ri in range(args.rows):
            row = _base_row(total_rows + ri, file_base)

            # ── Fix garbled cnam/snam from gen_test_data cert_cn bug ──
            # gen_test_data line 135 generates "*.['example', 'com']" garbage
            # via f'*.{random.choice(SERVICES).split(".")[-2:]}' — missing .join()
            snam_val = row.get('snam', '').strip()
            cnam_val = row.get('cnam', '').strip()
            import re
            if snam_val and snam_val != '' and re.match(r'^\*\.[\[\(]', snam_val):
                row['snam'] = ''
                snam_val = ''
            if cnam_val and cnam_val != '' and re.match(r'^\*\.[\[\(]', cnam_val):
                row['cnam'] = ''
                cnam_val = ''

            # ── Remap IPs using topology ──
            # Determine appropriate :ipd based on snam/cnam
            server_pool = None
            if snam_val and snam_val in SNAM_TO_IPS:
                server_pool = SNAM_TO_IPS[snam_val]
            elif cnam_val and cnam_val in CNAM_TO_IPS:
                server_pool = CNAM_TO_IPS[cnam_val]

            if server_pool:
                # Assign a specific server IP for this service
                dst_ip = random.choice(server_pool)
                row[':ipd'] = dst_ip
                row['server-ip'] = dst_ip
                snam_ip_count[snam_val][dst_ip] += 1
                cnam_ip_count[cnam_val][dst_ip] += 1

                # 30% of the time, use a proxy as :ips instead of random client
                if random.random() < 0.30:
                    row[':ips'] = random.choice(PROXY_IPS)
                    row['client-ip'] = row[':ips']

                # GeoIP: use MMDB in upload flow; CSV lat/lon are ignored anyway

            file_rows.append(row)
            degree_src[row[':ips']] += 1
            degree_dst[row[':ipd']] += 1

        total_rows += args.rows

        # ── Shuffle column order per file ──
        shuffled = file_columns[:]
        random.shuffle(shuffled)

        # ── Write ──
        out_path = out_dir / f'{fi}.csv'
        with open(out_path, 'w', newline='', encoding='utf-8') as f:
            writer = csv.DictWriter(f, fieldnames=shuffled, extrasaction='ignore')
            writer.writeheader()
            for row in file_rows:
                writer.writerow(row)

        if (fi + 1) % 50 == 0:
            print(f'  [{fi + 1}/{args.files}] {total_rows} rows...')

    # ── Stats ──
    top_src = sorted(degree_src.items(), key=lambda x: -x[1])[:10]
    top_dst = sorted(degree_dst.items(), key=lambda x: -x[1])[:10]
    print(f'\nDone: {args.files} files, {total_rows} rows → {out_dir.resolve()}')
    print(f'  Avg src degree: {sum(degree_src.values())/len(degree_src):.1f}')
    print(f'  Avg dst degree: {sum(degree_dst.values())/len(degree_dst):.1f}')
    print(f'  Top source (out-degree):')
    for ip, deg in top_src:
        label = ' [proxy]' if ip in PROXY_IPS else ''
        print(f'    {ip}: {deg}{label}')
    print(f'  Top dest (in-degree):')
    for ip, deg in top_dst:
        is_server = ip in ALL_SERVER_IPS
        print(f'    {ip}: {deg}{" [server]" if is_server else ""}')
    print(f'  snam→IP mapping:')
    for snam, ips in sorted(snam_ip_count.items()):
        top = sorted(ips.items(), key=lambda x: -x[1])
        print(f'    {snam}: {dict(top)}')


if __name__ == '__main__':
    main()