// Engineer: Tim Brockmann
// File: wlan_logger.c
// Install PCAP: apt install libpcap0.8-dev
// Compile: gcc -O2 -Wall -o wlan_logger wlan_logger.c -lpcap

#define _GNU_SOURCE
#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <getopt.h>
#include <errno.h>
#include <arpa/inet.h>

// ------------------ Simple packet info struct ------------------

typedef struct {
    uint64_t tsf;              // HW timestamp (from Radiotap TSFT) - TODO: real parsing
    int8_t   rssi_dbm;         // RSSI in dBm - TODO: real parsing
    uint8_t  mcs;              // MCS index - TODO: real parsing
    struct timespec sys_time;  // System time at receive
} wlan_pkt_info_t;

// ------------------ Global config ------------------

typedef struct {
    char iface[64];
    char log_path[256];
} logger_config_t;

static logger_config_t g_cfg;
static FILE *g_log_file = NULL;

// ------------------ Radiotap Parsing (Stub Version) ------------------

// Very simple Radiotap header structure (not a complete parser)
struct ieee80211_radiotap_header {
    uint8_t  it_version;
    uint8_t  it_pad;
    uint16_t it_len;
    uint32_t it_present; // Note: real radiotap may contain more present bitmaps (extensions)
} __attribute__((__packed__));

/*
 * TODO:
 *  - Evaluate presence bits (TSFT, DBM_ANTSIGNAL, MCS, etc.)
 *  - Handle alignment rules
 *
 * For now, we stub TSF/RSSI/MCS as "not available"
 * and focus on the program structure.
 */

static int parse_radiotap_stub(const uint8_t *data, size_t len, wlan_pkt_info_t *info)
{
    if (len < sizeof(struct ieee80211_radiotap_header)) {
        return -1;
    }

    const struct ieee80211_radiotap_header *rt =
        (const struct ieee80211_radiotap_header *)data;

    uint16_t hdr_len = rt->it_len;

    if (hdr_len > len) {
        return -1;
    }

    // Stub: placeholder values
    info->tsf      = 1;   // later: read TSFT field (present bit 0)
    info->rssi_dbm = 2;   // later: DBM_ANTSIGNAL (present bit 5)
    info->mcs      = 3;   // later: MCS extension

    return 0;
}

// Simple mapping from port to a name (extend as needed)
static const char *app_name_from_ports(uint16_t sport, uint16_t dport)
{
    uint16_t p = ntohs(dport);  // focus on destination port first

    switch (p) {
    case 319:
    case 320:
        return "PTP";
    case 22:
        return "SSH";
    case 80:
        return "HTTP";
    case 443:
        return "HTTPS";
    case 53:
        return "DNS";
    default:
        return "OTHER";
    }
}

static const char *frame_type_to_str(uint8_t type)
{
    switch (type) {
    case 0: return "MGMT";
    case 1: return "CTRL";
    case 2: return "DATA";
    case 3: return "EXT";
    default: return "UNKNOWN";
    }
}

/*
 * Very simplified 802.11 + LLC/SNAP + IPv4 + UDP/TCP parser.
 * Returns frame-type string and application name (if detectable).
 *
 * WARNING:
 *  - Only handles "simple" data frames without QoS/Addr4.
 *  - For your mesh use case, you might need a more complete parser later.
 */
static void decode_packet_info(const u_char *bytes, size_t len,
                               char *frame_type_buf, size_t ft_len,
                               char *app_name_buf, size_t app_len)
{
    // Default values
    snprintf(frame_type_buf, ft_len, "UNKNOWN");
    snprintf(app_name_buf,   app_len, "N/A");

    if (len < sizeof(struct ieee80211_radiotap_header))
        return;

    const struct ieee80211_radiotap_header *rt =
        (const struct ieee80211_radiotap_header *)bytes;
    uint16_t rt_len = rt->it_len;

    if (rt_len >= len)
        return;

    const uint8_t *mac = bytes + rt_len;
    size_t mac_len = len - rt_len;

    if (mac_len < 2)
        return;

    // Frame control
    uint16_t fc = ((uint16_t)mac[1] << 8) | mac[0];
    uint8_t type    = (fc >> 2) & 0x3;
    uint8_t subtype = (fc >> 4) & 0xF;
    (void)subtype; // not used yet

    snprintf(frame_type_buf, ft_len, "%s", frame_type_to_str(type));

    // We only care about data frames for "application" decoding
    if (type != 2) {
        return;
    }

    // Very rough assumption: data frame with 3 addresses (no QoS, no Addr4)
    // 802.11 header: 24 bytes
    if (mac_len < 24 + 8) { // need room for LLC/SNAP too
        return;
    }

    const uint8_t *llc = mac + 24;
    size_t llc_len = mac_len - 24;

    if (llc_len < 8) {
        return;
    }

    // LLC/SNAP header:
    // DSAP(1), SSAP(1), Control(1), OUI(3), EtherType(2)
    uint16_t ethertype = ((uint16_t)llc[6] << 8) | llc[7];

    if (ethertype != 0x0800 && ethertype != 0x86DD) {
        // Not IPv4/IPv6
        snprintf(app_name_buf, app_len, "NON-IP");
        return;
    }

    // Assume IPv4 for now if 0x0800
    if (ethertype == 0x0800) {
        const uint8_t *ip = llc + 8;
        size_t ip_len = llc_len - 8;
        if (ip_len < 20)
            return;

        uint8_t ip_proto = ip[9];
        const uint8_t *l4 = ip + (ip[0] & 0x0F) * 4; // IHL * 4
        if (l4 > ip + ip_len - 4)
            return;

        uint16_t sport = 0, dport = 0;
        const char *app = "IP";

        if (ip_proto == 6 || ip_proto == 17) {
            // TCP(6) or UDP(17)
            sport = ((uint16_t)l4[0] << 8) | l4[1];
            dport = ((uint16_t)l4[2] << 8) | l4[3];
            app = app_name_from_ports(sport, dport);
        }

        snprintf(app_name_buf, app_len, "%s", app);
    } else {
        // IPv6 - for now, just mark it
        snprintf(app_name_buf, app_len, "IPv6");
    }
}

// ------------------ Packet handler ------------------

static void handle_packet(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes)
{
    (void)user; // unused

    wlan_pkt_info_t info;
    memset(&info, 0, sizeof(info));

    // Get system time (for debugging / additional features)
    clock_gettime(CLOCK_REALTIME, &info.sys_time);

    if (parse_radiotap_stub(bytes, h->caplen, &info) != 0) {
        // Radiotap parsing failed -> ignore packet
        return;
    }

    // CSV output: tsf,sys_sec,sys_nsec,rssi_dbm,mcs
    fprintf(g_log_file,
            "%llu,%ld,%ld,%d,%u\n",
            (unsigned long long)info.tsf,
            (long)info.sys_time.tv_sec,
            (long)info.sys_time.tv_nsec,
            (int)info.rssi_dbm,
            (unsigned)info.mcs);

    // Optional: flush occasionally (instead of flushing on every write)
    static int counter = 0;
    if (++counter >= 1000) {
        fflush(g_log_file);
        counter = 0;
    }

    // --- NEW: decode frame type + application name for debugging ---
    char frame_type[16];
    char app_name[32];

    decode_packet_info(bytes, h->caplen, frame_type, sizeof(frame_type),
                       app_name, sizeof(app_name));

    printf("new packet: type=%s, app=%s\n", frame_type, app_name);
    fflush(stdout); // ensure immediate output
}

// ------------------ Config / Args ------------------

static void usage(const char *prog)
{
    fprintf(stderr,
            "Usage: %s --iface IFACE --log LOGFILE\n"
            "  --iface, -i   Monitor interface (e.g. mon0)\n"
            "  --log,   -l   Path to log file (CSV)\n",
            prog);
}

static int parse_args(int argc, char **argv, logger_config_t *cfg)
{
    memset(cfg, 0, sizeof(*cfg));

    static struct option long_opts[] = {
        {"iface", required_argument, 0, 'i'},
        {"log",   required_argument, 0, 'l'},
        {"help",  no_argument,       0, 'h'},
        {0,0,0,0}
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "i:l:h", long_opts, NULL)) != -1) {
        switch (opt) {
        case 'i':
            strncpy(cfg->iface, optarg, sizeof(cfg->iface)-1);
            break;
        case 'l':
            strncpy(cfg->log_path, optarg, sizeof(cfg->log_path)-1);
            break;
        case 'h':
        default:
            return -1;
        }
    }

    if (cfg->iface[0] == '\0' || cfg->log_path[0] == '\0') {
        return -1;
    }

    return 0;
}

// ------------------ Main ------------------

int main(int argc, char **argv)
{
    if (parse_args(argc, argv, &g_cfg) != 0) {
        usage(argv[0]);
        return EXIT_FAILURE;
    }

    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t *handle = NULL;

    // Open log file
    g_log_file = fopen(g_cfg.log_path, "w");
    if (!g_log_file) {
        fprintf(stderr, "Error opening log file %s: %s\n",
                g_cfg.log_path, strerror(errno));
        return EXIT_FAILURE;
    }

    // Open pcap device
    // handle = pcap_open_live(g_cfg.iface, 65535, 1, 1000, errbuf);
    handle = pcap_open_live(g_cfg.iface, 65535, 1, 1, errbuf);
    if (!handle) {
        fprintf(stderr, "pcap_open_live(%s) failed: %s\n",
                g_cfg.iface, errbuf);
        fclose(g_log_file);
        return EXIT_FAILURE;
    }

    // Optional: check link type (should be DLT_IEEE802_11_RADIO)
    int dlt = pcap_datalink(handle);
    if (dlt != DLT_IEEE802_11_RADIO) {
        fprintf(stderr,
                "Warning: interface %s has unexpected DLT %d (expected %d = IEEE802_11_RADIO)\n",
                g_cfg.iface, dlt, DLT_IEEE802_11_RADIO);
    }

    printf("wlan_logger: capturing on %s, logging to %s\n",
           g_cfg.iface, g_cfg.log_path);

    // Capture loop
    if (pcap_loop(handle, -1, handle_packet, NULL) == -1) {
        fprintf(stderr, "pcap_loop error: %s\n", pcap_geterr(handle));
    }

    pcap_close(handle);
    fclose(g_log_file);

    return EXIT_SUCCESS;
}