#!/usr/bin/env python3
import sqlite3
import csv
import argparse
from datetime import datetime

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", required=True, help="Input speedtest CSV")
    parser.add_argument("--kismet", required=True, help="Path to .kismet log file")
    parser.add_argument("--output", required=True, help="Output enriched CSV")
    return parser.parse_args()

def convert_timestamp_to_epoch(ts_string):
    try:
        return int(datetime.strptime(ts_string.split(".")[0], "%Y-%m-%dT%H:%M:%S").timestamp())
    except Exception as e:
        print(f"[!] Invalid timestamp format: {ts_string}")
        return None

def get_rf_metrics(cursor, bssid, channel, start_time, end_time):
    cursor.execute("""
        SELECT COUNT(*) FROM devices
        WHERE type = 'Wi-Fi Client'
        AND last_time BETWEEN ? AND ?
        AND LOWER(json_extract(device, '$.kismet.device.base.bssid')) = ?
    """, (start_time, end_time, bssid.lower()))
    clients_on_ap = cursor.fetchone()[0]

    cursor.execute("""
        SELECT COUNT(*) FROM devices
        WHERE type = 'Wi-Fi Client'
        AND last_time BETWEEN ? AND ?
        AND (
            json_extract(device, '$.kismet.device.base.channel') = ? OR
            json_type(json_extract(device, '$.kismet.device.base.freq_khz_map')) = 'object'
        )
        AND json_extract(device, '$.kismet.device.base.packets.total') > 0
    """, (start_time, end_time, channel))
    clients_on_channel = cursor.fetchone()[0]

    cursor.execute("""
        SELECT COUNT(*) FROM devices
        WHERE type = 'Wi-Fi AP'
            AND last_time BETWEEN ? AND ?
            AND (
                json_extract(device, '$.kismet.device.base.channel') = ? OR
                json_type(json_extract(device, '$.kismet.device.base.freq_khz_map')) = 'object'
            )
            AND strongest_signal < 0
    """, (start_time, end_time, channel))
    aps_on_channel = cursor.fetchone()[0]

    cursor.execute("""
        SELECT strongest_signal FROM devices
        WHERE type = 'Wi-Fi AP'
            AND last_time BETWEEN ? AND ?
            AND (
                json_extract(device, '$.kismet.device.base.channel') = ? OR
                json_type(json_extract(device, '$.kismet.device.base.freq_khz_map')) = 'object'
            )
            AND strongest_signal < 0
    """, (start_time, end_time, channel))
    signals = cursor.fetchall()
    avg_signal = sum([s[0] for s in signals]) / len(signals) if signals else None
    strongest_signal = max([s[0] for s in signals]) if signals else None

    cursor.execute("""
        SELECT COUNT(*) FROM devices
        WHERE type = 'Wi-Fi Client'
        AND last_time BETWEEN ? AND ?
        AND json_extract(device, '$.kismet.device.base.channel') IS NULL
        AND json_type(json_extract(device, '$.kismet.device.base.freq_khz_map')) != 'object'
        AND json_extract(device, '$.kismet.device.base.packets.total') = 0
    """, (start_time, end_time))
    unlinked_devices = cursor.fetchone()[0]

    congestion_score = round(clients_on_channel / aps_on_channel, 2) if aps_on_channel else None

    return clients_on_ap, clients_on_channel, aps_on_channel, congestion_score, avg_signal, strongest_signal, unlinked_devices

def main():
    args = parse_args()
    conn = sqlite3.connect(args.kismet)
    cursor = conn.cursor()

    with open(args.csv, newline='') as infile, open(args.output, 'w', newline='') as outfile:
        reader = csv.DictReader(infile)
        fieldnames = reader.fieldnames + [
            "ClientsOnAP", "ClientsOnChannel", "APsOnChannel", "CongestionScore"
        ]
        writer = csv.DictWriter(outfile, fieldnames=fieldnames)
        writer.writeheader()

        for row in reader:
            tstart = convert_timestamp_to_epoch(row["StartTimestamp"])
            tend = convert_timestamp_to_epoch(row["EndTimestamp"])
            try:
                bssid = row["BSSID"].strip().lower()
                channel = int(row["Channel"])
            except Exception as e:
                print(f"[!] Failed to extract BSSID/Channel: {e}")
                writer.writerow(row)
                continue

            if tstart is None or tend is None:
                writer.writerow(row)
                continue

            clients_ap, clients_chan, aps_chan, congestion = get_rf_metrics(cursor, bssid, channel, tstart, tend)

            row["ClientsOnAP"] = clients_ap
            row["ClientsOnChannel"] = clients_chan
            row["APsOnChannel"] = aps_chan
            row["CongestionScore"] = congestion

            writer.writerow(row)

    conn.close()
    print(f"[+] Enrichment complete: {args.output}")

if __name__ == "__main__":
    main()