seed_from_csv.py

"""
Seeds network.db with device/interface/flow/alert/metric data.

Two modes:
1. Real data: point --csv at a NetFlow-style CSV (e.g. NF-UNSW-NB15-v2 from
   https://staff.itee.uq.edu.au/marius/NIDS_datasets/) with columns similar to
   IPV4_SRC_ADDR, IPV4_DST_ADDR, L4_SRC_PORT, L4_DST_PORT, PROTOCOL, IN_BYTES,
   IN_PKTS, FLOW_DURATION_MILLISECONDS, Label, Attack. Devices/interfaces/
   alerts are derived from the distinct IPs in the flow data.
2. Synthetic fallback (default, no --csv given): generates a realistic set of
   devices, interfaces, flows, alerts, and metrics from scratch so the demo
   works with zero external downloads.

Usage:
    python seed_from_csv.py                  # synthetic data
    python seed_from_csv.py --csv path/to/nf_unsw_nb15_v2.csv --limit 20000
"""

import argparse
import ipaddress
import random
import sqlite3
from datetime import datetime, timedelta
from pathlib import Path

DB_PATH = Path(__file__).parent / "network.db"
SCHEMA_PATH = Path(__file__).parent / "schema.sql"

DEVICE_TYPES = ["router", "switch", "firewall", "host", "access_point"]
VENDORS = ["Cisco", "Juniper", "Arista", "MikroTik", "Ubiquiti", "Fortinet"]
LOCATIONS = ["HQ-Floor1", "HQ-Floor2", "DataCenter-A", "Branch-Lagos", "Branch-Abuja", "Cloud-Edge"]
PROTOCOLS = ["TCP", "UDP", "ICMP"]
ALERT_TYPES = ["port_scan", "dos", "malware", "link_down", "auth_failure", "config_drift"]
SEVERITIES = ["info", "warning", "critical"]


def reset_db():
    if DB_PATH.exists():
        DB_PATH.unlink()
    conn = sqlite3.connect(DB_PATH)
    conn.executescript(SCHEMA_PATH.read_text())
    conn.commit()
    return conn


def now_minus(seconds):
    return (datetime.utcnow() - timedelta(seconds=seconds)).isoformat(timespec="seconds")


def seed_synthetic(conn, n_devices=18, n_flows=4000, n_alerts=40, n_metrics=3000):
    cur = conn.cursor()

    device_ids = []
    for i in range(n_devices):
        dtype = random.choice(DEVICE_TYPES)
        status = random.choices(["up", "degraded", "down"], weights=[0.82, 0.13, 0.05])[0]
        cur.execute(
            "INSERT INTO devices (hostname, ip_address, device_type, vendor, location, status, last_seen) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                f"{dtype}-{i+1:02d}",
                str(ipaddress.IPv4Address(0x0A000000 + i + 1)),  # 10.0.0.x
                dtype,
                random.choice(VENDORS),
                random.choice(LOCATIONS),
                status,
                now_minus(random.randint(0, 600)),
            ),
        )
        device_ids.append(cur.lastrowid)

    interface_ids = []
    for did in device_ids:
        for j in range(random.randint(1, 4)):
            cur.execute(
                "INSERT INTO interfaces (device_id, name, status, speed_mbps) VALUES (?, ?, ?, ?)",
                (
                    did,
                    f"Gi0/{j}",
                    random.choices(["up", "down"], weights=[0.9, 0.1])[0],
                    random.choice([100, 1000, 10000]),
                ),
            )
            interface_ids.append((cur.lastrowid, did))

    for _ in range(n_flows):
        did = random.choice(device_ids)
        is_attack = random.random() < 0.06
        cur.execute(
            "INSERT INTO traffic_flows (ts, src_ip, dst_ip, src_port, dst_port, protocol, "
            "bytes, packets, duration_ms, device_id, label) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                now_minus(random.randint(0, 86400)),
                f"10.0.0.{random.randint(1, 254)}",
                f"{random.choice(['172.16', '192.168', '203.0'])}.{random.randint(0,255)}.{random.randint(1,254)}",
                random.randint(1024, 65535),
                random.choice([80, 443, 22, 53, 3389, 8080, 3306]),
                random.choice(PROTOCOLS),
                random.randint(200, 5_000_000),
                random.randint(1, 4000),
                random.randint(1, 60000),
                did,
                random.choice(ALERT_TYPES) if is_attack else "benign",
            ),
        )

    for _ in range(n_alerts):
        did = random.choice(device_ids)
        cur.execute(
            "INSERT INTO alerts (ts, device_id, severity, alert_type, description, resolved) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (
                now_minus(random.randint(0, 172800)),
                did,
                random.choices(SEVERITIES, weights=[0.5, 0.35, 0.15])[0],
                random.choice(ALERT_TYPES),
                "Auto-generated demo alert",
                random.choices([0, 1], weights=[0.4, 0.6])[0],
            ),
        )

    for _ in range(n_metrics):
        iface_id, did = random.choice(interface_ids)
        metric_name = random.choice(["latency_ms", "packet_loss_pct", "bandwidth_util_pct"])
        value = {
            "latency_ms": round(random.uniform(1, 180), 2),
            "packet_loss_pct": round(random.uniform(0, 6), 2),
            "bandwidth_util_pct": round(random.uniform(2, 98), 2),
        }[metric_name]
        cur.execute(
            "INSERT INTO metrics_timeseries (ts, device_id, interface_id, metric_name, value) "
            "VALUES (?, ?, ?, ?, ?)",
            (now_minus(random.randint(0, 3600)), did, iface_id, metric_name, value),
        )

    conn.commit()
    print(f"Synthetic seed complete: {len(device_ids)} devices, {n_flows} flows, "
          f"{n_alerts} alerts, {n_metrics} metrics.")


def seed_from_real_csv(conn, csv_path, limit):
    import csv as csv_mod

    cur = conn.cursor()
    ip_to_device = {}

    def get_or_create_device(ip):
        if ip in ip_to_device:
            return ip_to_device[ip]
        cur.execute(
            "INSERT INTO devices (hostname, ip_address, device_type, vendor, location, status, last_seen) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (f"host-{ip}", ip, "host", random.choice(VENDORS), random.choice(LOCATIONS), "up", now_minus(0)),
        )
        ip_to_device[ip] = cur.lastrowid
        return cur.lastrowid

    with open(csv_path, newline="") as f:
        reader = csv_mod.DictReader(f)
        for i, row in enumerate(reader):
            if i >= limit:
                break
            src_ip = row.get("IPV4_SRC_ADDR") or row.get("src_ip") or "0.0.0.0"
            dst_ip = row.get("IPV4_DST_ADDR") or row.get("dst_ip") or "0.0.0.0"
            device_id = get_or_create_device(src_ip)
            label = row.get("Attack") or row.get("label") or ("benign" if row.get("Label", "0") == "0" else "attack")
            cur.execute(
                "INSERT INTO traffic_flows (ts, src_ip, dst_ip, src_port, dst_port, protocol, "
                "bytes, packets, duration_ms, device_id, label) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    now_minus(random.randint(0, 86400)),
                    src_ip,
                    dst_ip,
                    int(row.get("L4_SRC_PORT", 0) or 0),
                    int(row.get("L4_DST_PORT", 0) or 0),
                    row.get("PROTOCOL", "TCP"),
                    int(row.get("IN_BYTES", 0) or 0),
                    int(row.get("IN_PKTS", 0) or 0),
                    int(float(row.get("FLOW_DURATION_MILLISECONDS", 0) or 0)),
                    device_id,
                    label,
                ),
            )

    conn.commit()
    print(f"Loaded {i+1} flows from {csv_path}, derived {len(ip_to_device)} devices.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", type=str, default=None, help="Path to a real NetFlow CSV")
    parser.add_argument("--limit", type=int, default=20000, help="Max rows to load from CSV")
    args = parser.parse_args()

    conn = reset_db()
    if args.csv:
        seed_from_real_csv(conn, args.csv, args.limit)
    else:
        seed_synthetic(conn)
    conn.close()