#!/usr/bin/env python3
"""
RF Attack Classifier — Synthetic Data Generator
Generates 50k samples across 6 classes for training an ESP32-S3 attack detector.

Feature vector (64 elements):
  [0-5]   Frame type histogram (mgmt, ctrl, data, beacon, deauth, eapol) — normalized counts
  [6-11]  Frame subtype histogram (probe_req, probe_resp, assoc, disassoc, auth, action)
  [12-15] Rate statistics (mean, std, min, max) of frames per second over sliding window
  [16-19] BSSID diversity (unique count, new_per_sec, churn_rate, concentration_index)
  [20-23] SSID statistics (unique count, avg_length, max_length, entropy)
  [24-27] RSSI statistics (mean, std, min, max) across observed BSSIDs
  [28-31] RSSI per-BSSID variance (mean_var, max_var, sudden_jump_count, jump_magnitude)
  [32-35] Timing features (mean_iat, std_iat, min_iat, burst_ratio)
  [36-39] Sequence number features (gap_count, dup_count, out_of_order_ratio, reset_count)
  [40-43] Source address diversity (unique_src, src_churn, broadcast_ratio, multicast_ratio)
  [44-47] ARP features (arp_rate, gratuitous_arp_ratio, dup_ip_mac_count, arp_reply_no_request)
  [48-51] EAPOL features (eapol_rate, pmkid_pattern_score, incomplete_handshake, msg1_ratio)
  [52-55] Channel features (active_channels, channel_switches, dwell_time_std, primary_channel_ratio)
  [56-59] Payload features (mean_size, std_size, empty_payload_ratio, anomalous_size_ratio)
  [60-63] Derived features (entropy_score, periodicity_score, attack_confidence_raw, noise_floor_delta)

Classes: {0: normal, 1: deauth, 2: beacon_flood, 3: evil_twin, 4: pmkid, 5: arp_spoof}
"""

import numpy as np
import os
import argparse

SEED = 42
NUM_FEATURES = 64
CLASSES = ['normal', 'deauth', 'beacon_flood', 'evil_twin', 'pmkid', 'arp_spoof']
NUM_CLASSES = len(CLASSES)


def generate_normal(n: int, rng: np.random.Generator) -> np.ndarray:
    """Normal WiFi traffic — balanced frame types, stable RSSI, regular timing."""
    X = np.zeros((n, NUM_FEATURES), dtype=np.float32)

    # Frame type histogram — mostly data and beacons
    X[:, 0] = rng.normal(0.25, 0.05, n).clip(0.05, 0.5)   # mgmt
    X[:, 1] = rng.normal(0.15, 0.04, n).clip(0.02, 0.3)    # ctrl
    X[:, 2] = rng.normal(0.45, 0.08, n).clip(0.2, 0.7)     # data
    X[:, 3] = rng.normal(0.12, 0.03, n).clip(0.02, 0.25)   # beacon
    X[:, 4] = rng.normal(0.005, 0.003, n).clip(0, 0.02)    # deauth (very low)
    X[:, 5] = rng.normal(0.01, 0.005, n).clip(0, 0.03)     # eapol

    # Frame subtype histogram
    X[:, 6] = rng.normal(0.05, 0.02, n).clip(0, 0.15)      # probe_req
    X[:, 7] = rng.normal(0.04, 0.02, n).clip(0, 0.12)      # probe_resp
    X[:, 8] = rng.normal(0.02, 0.01, n).clip(0, 0.05)      # assoc
    X[:, 9] = rng.normal(0.005, 0.003, n).clip(0, 0.02)    # disassoc
    X[:, 10] = rng.normal(0.03, 0.015, n).clip(0, 0.08)    # auth
    X[:, 11] = rng.normal(0.01, 0.005, n).clip(0, 0.03)    # action

    # Rate statistics (frames per second)
    X[:, 12] = rng.normal(50, 15, n).clip(10, 120)          # mean fps
    X[:, 13] = rng.normal(10, 4, n).clip(2, 30)             # std fps
    X[:, 14] = rng.normal(15, 8, n).clip(1, 40)             # min fps
    X[:, 15] = rng.normal(90, 25, n).clip(40, 200)          # max fps

    # BSSID diversity
    X[:, 16] = rng.normal(8, 3, n).clip(2, 20)              # unique BSSIDs
    X[:, 17] = rng.normal(0.1, 0.05, n).clip(0, 0.3)       # new per sec
    X[:, 18] = rng.normal(0.05, 0.03, n).clip(0, 0.15)     # churn rate
    X[:, 19] = rng.normal(0.3, 0.1, n).clip(0.1, 0.6)      # concentration index

    # SSID statistics
    X[:, 20] = rng.normal(6, 2.5, n).clip(1, 15)            # unique SSIDs
    X[:, 21] = rng.normal(8, 2, n).clip(3, 15)              # avg length
    X[:, 22] = rng.normal(15, 4, n).clip(5, 32)             # max length
    X[:, 23] = rng.normal(2.5, 0.5, n).clip(1, 4)           # entropy

    # RSSI statistics
    X[:, 24] = rng.normal(-55, 8, n).clip(-85, -30)         # mean RSSI
    X[:, 25] = rng.normal(5, 2, n).clip(1, 12)              # std RSSI
    X[:, 26] = rng.normal(-75, 6, n).clip(-90, -55)         # min RSSI
    X[:, 27] = rng.normal(-35, 5, n).clip(-50, -20)         # max RSSI

    # RSSI per-BSSID variance
    X[:, 28] = rng.normal(3, 1.5, n).clip(0.5, 8)           # mean var
    X[:, 29] = rng.normal(6, 3, n).clip(1, 15)              # max var
    X[:, 30] = rng.normal(1, 0.8, n).clip(0, 4)             # sudden jump count
    X[:, 31] = rng.normal(5, 2, n).clip(1, 12)              # jump magnitude

    # Timing features
    X[:, 32] = rng.normal(0.02, 0.008, n).clip(0.005, 0.05) # mean IAT
    X[:, 33] = rng.normal(0.01, 0.004, n).clip(0.002, 0.03) # std IAT
    X[:, 34] = rng.normal(0.001, 0.0005, n).clip(0, 0.005)  # min IAT
    X[:, 35] = rng.normal(0.1, 0.04, n).clip(0.01, 0.3)     # burst ratio

    # Sequence number features
    X[:, 36] = rng.normal(2, 1.5, n).clip(0, 8)             # gap count
    X[:, 37] = rng.normal(0.5, 0.4, n).clip(0, 2)           # dup count
    X[:, 38] = rng.normal(0.01, 0.008, n).clip(0, 0.04)     # out of order ratio
    X[:, 39] = rng.normal(0.5, 0.3, n).clip(0, 2)           # reset count

    # Source address diversity
    X[:, 40] = rng.normal(10, 4, n).clip(2, 25)             # unique src
    X[:, 41] = rng.normal(0.05, 0.03, n).clip(0, 0.15)      # src churn
    X[:, 42] = rng.normal(0.08, 0.04, n).clip(0, 0.2)       # broadcast ratio
    X[:, 43] = rng.normal(0.02, 0.01, n).clip(0, 0.05)      # multicast ratio

    # ARP features (low in normal)
    X[:, 44] = rng.normal(0.5, 0.3, n).clip(0, 1.5)         # arp rate
    X[:, 45] = rng.normal(0.05, 0.03, n).clip(0, 0.15)      # gratuitous arp ratio
    X[:, 46] = rng.normal(0, 0.01, n).clip(0, 0.05)         # dup ip-mac count
    X[:, 47] = rng.normal(0.1, 0.06, n).clip(0, 0.3)        # arp reply no request

    # EAPOL features (low in normal)
    X[:, 48] = rng.normal(0.01, 0.008, n).clip(0, 0.04)     # eapol rate
    X[:, 49] = rng.normal(0.0, 0.01, n).clip(0, 0.05)       # pmkid pattern score
    X[:, 50] = rng.normal(0.02, 0.015, n).clip(0, 0.06)     # incomplete handshake
    X[:, 51] = rng.normal(0.25, 0.1, n).clip(0, 0.5)        # msg1 ratio

    # Channel features
    X[:, 52] = rng.normal(3, 1.5, n).clip(1, 8)             # active channels
    X[:, 53] = rng.normal(2, 1.5, n).clip(0, 6)             # channel switches
    X[:, 54] = rng.normal(5, 3, n).clip(1, 15)              # dwell time std
    X[:, 55] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)      # primary channel ratio

    # Payload features
    X[:, 56] = rng.normal(300, 100, n).clip(50, 800)         # mean size
    X[:, 57] = rng.normal(200, 80, n).clip(30, 500)          # std size
    X[:, 58] = rng.normal(0.05, 0.03, n).clip(0, 0.15)      # empty payload ratio
    X[:, 59] = rng.normal(0.02, 0.015, n).clip(0, 0.06)     # anomalous size ratio

    # Derived features
    X[:, 60] = rng.normal(3.5, 0.5, n).clip(2, 5)           # entropy score
    X[:, 61] = rng.normal(0.3, 0.15, n).clip(0, 0.7)        # periodicity score
    X[:, 62] = rng.normal(0.05, 0.03, n).clip(0, 0.15)      # attack confidence raw
    X[:, 63] = rng.normal(2, 1, n).clip(0, 5)               # noise floor delta

    return X


def generate_deauth(n: int, rng: np.random.Generator) -> np.ndarray:
    """Deauth flood — high deauth frame rate, burst timing, spoofed BSSIDs."""
    X = generate_normal(n, rng)

    # Spike deauth frames dramatically
    X[:, 4] = rng.normal(0.45, 0.12, n).clip(0.2, 0.8)     # deauth fraction
    X[:, 0] = rng.normal(0.55, 0.1, n).clip(0.3, 0.85)     # mgmt dominates
    X[:, 2] = rng.normal(0.15, 0.08, n).clip(0.02, 0.35)   # data drops
    X[:, 9] = rng.normal(0.15, 0.06, n).clip(0.05, 0.35)   # disassoc rises

    # Rate spikes
    X[:, 12] = rng.normal(350, 100, n).clip(150, 800)       # mean fps much higher
    X[:, 15] = rng.normal(600, 150, n).clip(300, 1200)      # max fps

    # Timing — very bursty
    X[:, 32] = rng.normal(0.003, 0.001, n).clip(0.0005, 0.008)  # mean IAT (fast)
    X[:, 33] = rng.normal(0.002, 0.001, n).clip(0.0003, 0.005)  # std IAT
    X[:, 35] = rng.normal(0.7, 0.15, n).clip(0.4, 0.95)         # burst ratio high

    # Sequence anomalies
    X[:, 36] = rng.normal(15, 6, n).clip(5, 40)             # gap count high
    X[:, 37] = rng.normal(8, 4, n).clip(2, 20)              # dup count high
    X[:, 38] = rng.normal(0.15, 0.06, n).clip(0.05, 0.35)   # out of order

    # Source churn (spoofed MACs)
    X[:, 40] = rng.normal(3, 1.5, n).clip(1, 8)             # few unique src (targeted)
    X[:, 41] = rng.normal(0.02, 0.015, n).clip(0, 0.06)     # low churn

    # Derived
    X[:, 62] = rng.normal(0.75, 0.12, n).clip(0.45, 0.95)   # attack confidence high

    return X


def generate_beacon_flood(n: int, rng: np.random.Generator) -> np.ndarray:
    """Beacon flood — hundreds of unique SSIDs, beacon frame dominance."""
    X = generate_normal(n, rng)

    # Beacon frames dominate
    X[:, 3] = rng.normal(0.65, 0.12, n).clip(0.4, 0.9)     # beacon ratio
    X[:, 0] = rng.normal(0.7, 0.1, n).clip(0.5, 0.92)      # mgmt dominates
    X[:, 2] = rng.normal(0.1, 0.05, n).clip(0.02, 0.25)    # data drops

    # Massive BSSID and SSID diversity
    X[:, 16] = rng.normal(150, 50, n).clip(50, 400)         # unique BSSIDs
    X[:, 17] = rng.normal(15, 5, n).clip(5, 35)             # new per sec
    X[:, 18] = rng.normal(0.8, 0.1, n).clip(0.5, 0.98)     # churn rate extreme
    X[:, 19] = rng.normal(0.01, 0.005, n).clip(0, 0.03)     # concentration very low

    X[:, 20] = rng.normal(120, 40, n).clip(40, 300)          # unique SSIDs
    X[:, 21] = rng.normal(12, 3, n).clip(5, 20)              # avg SSID length (often random)
    X[:, 23] = rng.normal(4.5, 0.3, n).clip(3.5, 5.0)       # high entropy SSIDs

    # Rate high
    X[:, 12] = rng.normal(250, 80, n).clip(100, 500)
    X[:, 15] = rng.normal(450, 120, n).clip(200, 800)

    # RSSI — often similar (single transmitter)
    X[:, 25] = rng.normal(2, 1, n).clip(0.5, 5)             # low RSSI std
    X[:, 28] = rng.normal(1, 0.5, n).clip(0.2, 3)           # low per-BSSID var

    # Derived
    X[:, 60] = rng.normal(4.8, 0.15, n).clip(4.2, 5.0)      # very high entropy
    X[:, 62] = rng.normal(0.7, 0.12, n).clip(0.4, 0.92)

    return X


def generate_evil_twin(n: int, rng: np.random.Generator) -> np.ndarray:
    """Evil twin — duplicate BSSID with different signal strength."""
    X = generate_normal(n, rng)

    # RSSI variance spikes (same BSSID, two locations)
    X[:, 25] = rng.normal(15, 4, n).clip(8, 28)             # high overall RSSI std
    X[:, 28] = rng.normal(18, 5, n).clip(10, 35)            # high per-BSSID variance
    X[:, 29] = rng.normal(25, 6, n).clip(15, 45)            # max variance huge
    X[:, 30] = rng.normal(8, 3, n).clip(3, 18)              # many sudden jumps
    X[:, 31] = rng.normal(20, 5, n).clip(10, 35)            # jump magnitude

    # BSSID diversity slightly elevated but not extreme
    X[:, 16] = rng.normal(10, 3, n).clip(4, 20)
    X[:, 19] = rng.normal(0.45, 0.1, n).clip(0.2, 0.7)     # higher concentration (dup BSSID)

    # Beacon rate slightly elevated
    X[:, 3] = rng.normal(0.2, 0.05, n).clip(0.1, 0.35)

    # Auth frames elevated (twin tries to capture clients)
    X[:, 10] = rng.normal(0.08, 0.03, n).clip(0.03, 0.18)

    # Derived
    X[:, 62] = rng.normal(0.6, 0.15, n).clip(0.3, 0.85)

    return X


def generate_pmkid(n: int, rng: np.random.Generator) -> np.ndarray:
    """PMKID capture attempt — EAPOL patterns, incomplete handshakes."""
    X = generate_normal(n, rng)

    # EAPOL features spike
    X[:, 5] = rng.normal(0.12, 0.04, n).clip(0.05, 0.25)   # eapol frame ratio
    X[:, 48] = rng.normal(0.15, 0.05, n).clip(0.05, 0.3)   # eapol rate
    X[:, 49] = rng.normal(0.75, 0.12, n).clip(0.45, 0.95)  # pmkid pattern score HIGH
    X[:, 50] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)     # incomplete handshake HIGH
    X[:, 51] = rng.normal(0.8, 0.1, n).clip(0.55, 0.98)    # msg1 ratio HIGH (mostly msg1)

    # Auth frames elevated
    X[:, 10] = rng.normal(0.1, 0.04, n).clip(0.04, 0.22)

    # Moderate rate increase
    X[:, 12] = rng.normal(80, 25, n).clip(30, 180)

    # Source churn — attacker probing multiple APs
    X[:, 40] = rng.normal(4, 2, n).clip(1, 10)
    X[:, 41] = rng.normal(0.15, 0.06, n).clip(0.03, 0.3)

    # Derived
    X[:, 62] = rng.normal(0.65, 0.12, n).clip(0.35, 0.88)

    return X


def generate_arp_spoof(n: int, rng: np.random.Generator) -> np.ndarray:
    """ARP spoofing — gratuitous ARPs, duplicate IP-MAC mappings."""
    X = generate_normal(n, rng)

    # ARP features spike
    X[:, 44] = rng.normal(8, 3, n).clip(3, 20)              # arp rate high
    X[:, 45] = rng.normal(0.6, 0.15, n).clip(0.3, 0.9)     # gratuitous arp ratio HIGH
    X[:, 46] = rng.normal(5, 2, n).clip(2, 12)              # dup ip-mac count HIGH
    X[:, 47] = rng.normal(0.55, 0.15, n).clip(0.25, 0.85)  # arp reply no request HIGH

    # Data frames slightly elevated (ARP is layer 2)
    X[:, 2] = rng.normal(0.5, 0.08, n).clip(0.3, 0.7)

    # Payload — ARP packets are small
    X[:, 56] = rng.normal(80, 30, n).clip(28, 150)          # mean size small
    X[:, 57] = rng.normal(40, 20, n).clip(10, 100)          # std size small

    # Periodicity (ARP poison is often periodic)
    X[:, 61] = rng.normal(0.7, 0.12, n).clip(0.4, 0.92)

    # Derived
    X[:, 62] = rng.normal(0.6, 0.13, n).clip(0.3, 0.82)

    return X


GENERATORS = [
    generate_normal,
    generate_deauth,
    generate_beacon_flood,
    generate_evil_twin,
    generate_pmkid,
    generate_arp_spoof,
]


def main():
    parser = argparse.ArgumentParser(description='Generate synthetic RF attack data')
    parser.add_argument('--total', type=int, default=50000, help='Total samples')
    parser.add_argument('--test-frac', type=float, default=0.2, help='Test fraction')
    parser.add_argument('--seed', type=int, default=SEED)
    parser.add_argument('--output-dir', type=str, default=os.path.dirname(os.path.abspath(__file__)))
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    samples_per_class = args.total // NUM_CLASSES
    remainder = args.total - samples_per_class * NUM_CLASSES

    all_X = []
    all_y = []

    for cls_idx, gen_fn in enumerate(GENERATORS):
        n = samples_per_class + (1 if cls_idx < remainder else 0)
        X_cls = gen_fn(n, rng)
        y_cls = np.full(n, cls_idx, dtype=np.int32)
        all_X.append(X_cls)
        all_y.append(y_cls)
        print(f"  Class {cls_idx} ({CLASSES[cls_idx]}): {n} samples generated")

    X = np.concatenate(all_X, axis=0)
    y = np.concatenate(all_y, axis=0)

    # Add small Gaussian noise for realism
    X += rng.normal(0, 0.01, X.shape).astype(np.float32)

    # Shuffle
    idx = rng.permutation(len(X))
    X, y = X[idx], y[idx]

    # Split
    n_test = int(len(X) * args.test_frac)
    X_train, X_test = X[n_test:], X[:n_test]
    y_train, y_test = y[n_test:], y[:n_test]

    train_path = os.path.join(args.output_dir, 'train.npz')
    test_path = os.path.join(args.output_dir, 'test.npz')
    np.savez_compressed(train_path, X=X_train, y=y_train)
    np.savez_compressed(test_path, X=X_test, y=y_test)

    print(f"\nDataset summary:")
    print(f"  Total samples: {len(X)}")
    print(f"  Train: {len(X_train)} | Test: {len(X_test)}")
    print(f"  Features: {NUM_FEATURES}")
    print(f"  Classes: {NUM_CLASSES} — {CLASSES}")
    print(f"  Saved: {train_path}, {test_path}")

    # Quick class balance check
    for split_name, split_y in [('Train', y_train), ('Test', y_test)]:
        unique, counts = np.unique(split_y, return_counts=True)
        print(f"  {split_name} distribution: { {CLASSES[u]: c for u, c in zip(unique, counts)} }")


if __name__ == '__main__':
    main()
