#include <avm_coredump.h>
#include <linux/nmi.h>

#include "conf.h"

struct ip4_hdr {
    u8  ver_hdr_len;
    u8  diff;
    u16 len;
    u16 id;
    u16 flags;
    u8  ttl;
    u8  proto;
    u16 checksum;
    u8  addr_from[4];
    u8  addr_to[4];
};

struct ip6_hdr {
    u8 ver_cls_1;
    u8 cls_2_flw_1;
    u16 flw_2;

    u16 len;
    u8 next_hdr;
    u8 hop_limit;
    u8 addr_from[16];
    u8 addr_to[16];
};

struct eth_hdr {
    u8  mac_to[6];
    u8  mac_from[6];
    u16 eth_type;
};

struct udp_hdr {
    u16 port_from;
    u16 port_to;
    u16 udp_len;
    u16 udp_sum;
};

static void avm_coredump_periodic(void) {
    touch_nmi_watchdog();
    avm_coredump_net_poll_eth();
}

static void ip4_hdr_fill_checksum(struct ip4_hdr *hdr) {
    u32 checksum = 0;
    u16 *p = (u16*)hdr;
    unsigned int i;

    for (i = 0; i < sizeof(struct ip4_hdr)/sizeof(u16); i++, p++) {
        checksum += *p;
    }

    while (checksum & 0xFF0000) {
        checksum = (checksum & 0xFFFF) + (checksum >> 16);
    }

    hdr->checksum = ~checksum;
}

static void ip6_udp_hdr_fill_checksum(void *buf) {
    struct eth_hdr *eth;
    struct ip6_hdr *ip6;
    struct udp_hdr *udp;
    u32 checksum = 0;
    u16 *p;
    u32 temp;
    unsigned int i;

    eth = (struct eth_hdr*)buf;
    ip6 = (struct ip6_hdr*)(eth + 1);
    udp = (struct udp_hdr*)(ip6 + 1);

#define ADD(buf, size) \
    for (p = (void*)(buf), i = 0; i < (size) / 2; p++, i++) { \
        checksum += *p; \
        checksum = (checksum & 0xFFFF) + (checksum >> 16); \
    } \
    if ((size) & 1) { \
        checksum += (*((u8*)p) << 8); \
        checksum = (checksum & 0xFFFF) + (checksum >> 16); \
    } \

    ADD(ip6->addr_from, 16);
    ADD(ip6->addr_to, 16);
    temp = ip6->len;
    ADD(&temp, 4);

    temp = ip6->next_hdr;
    ADD(&temp, 4);

    ADD(udp, ntohs(udp->udp_len));
#undef ADD

    udp->udp_sum = ~checksum;
}

void avm_coredump_net_prepare_pkg(u16 port, void **tbuf, u32 *available_len) {
    struct eth_hdr *eth;
    struct ip4_hdr *ip4;
    struct ip6_hdr *ip6;
    struct udp_hdr *udp;
    u8 *buf = conf->packet;

    eth = (struct eth_hdr*)buf;
    if (conf->family == AF_INET6) {
        ip6 = (struct ip6_hdr*)(eth + 1);
        udp = (struct udp_hdr*)(ip6 + 1);
    }
    else {
        ip4 = (struct ip4_hdr*)(eth + 1);
        udp = (struct udp_hdr*)(ip4 + 1);
    }

    // Build ethernet header
    memcpy(eth->mac_from, conf->source_mac, 6);
    memcpy(eth->mac_to, conf->target_mac, 6);
    if (conf->family == AF_INET6) {
        eth->eth_type = htons(0x86DD);

        ip6->ver_cls_1 = 0x60;
        ip6->cls_2_flw_1 = 0x00;
        ip6->flw_2 = htons(0x0000);
        ip6->next_hdr = 17;
        ip6->hop_limit = 64;
        memcpy(ip6->addr_from, conf->source_ip, 16);
        memcpy(ip6->addr_to, conf->target_ip, 16);
    }
    else {
        eth->eth_type = htons(0x0800);

        // Build ipv4 header
        ip4->ver_hdr_len = 0x45;
        ip4->diff = 0x0;
        ip4->id = 0xCAFE;
        ip4->flags = htons(0x4000);
        ip4->ttl = 64;
        ip4->proto = 17;
        ip4->checksum = htons(0x0000);
        memcpy(ip4->addr_from, conf->source_ip, 4);
        memcpy(ip4->addr_to, conf->target_ip, 4);
    }

    udp->port_from = htons(conf->source_port);
    udp->port_to = htons(port);
    udp->udp_sum = htons(0x0000); // 0x0 means checksum disabled

    buf = (u8*)(udp + 1);
    *tbuf = buf;
    *available_len = MAX_PACKET_SIZE - (buf - conf->packet);
}

int avm_coredump_net_send_pkg(u32 len) {
    struct eth_hdr *eth;
    struct ip4_hdr *ip4;
    struct ip6_hdr *ip6;
    struct udp_hdr *udp;
    int ret;
    u32 total_len;
    u8 *buf = conf->packet;

    eth = (struct eth_hdr*)buf;
    if (conf->family == AF_INET6) {
        ip6 = (struct ip6_hdr*)(eth + 1);
        udp = (struct udp_hdr*)(ip6 + 1);
    }
    else {
        ip4 = (struct ip4_hdr*)(eth + 1);
        udp = (struct udp_hdr*)(ip4 + 1);
    }

    udp->udp_len = htons(len + sizeof(struct udp_hdr));

    if (conf->family == AF_INET6) {
        total_len = sizeof(struct eth_hdr) + sizeof(struct ip6_hdr) + sizeof(struct udp_hdr) + len;
        ip6->len = htons(sizeof(struct udp_hdr) + len);
        ip6_udp_hdr_fill_checksum(conf->packet);
    }
    else {
        total_len = sizeof(struct eth_hdr) + sizeof(struct ip4_hdr) + sizeof(struct udp_hdr) + len;
        ip4->len = htons(sizeof(struct ip4_hdr) + sizeof(struct udp_hdr) + len);
        ip4_hdr_fill_checksum(ip4);
    }

    do {
        ret = avm_coredump_net_send_eth(conf->packet, total_len);
        avm_coredump_periodic();
    } while (ret == -EAGAIN);

    return ret;
}

int avm_coredump_net_init(void) {
   avm_coredump_net_prepare_setup();

   return 0;
}

void avm_coredump_net_exit(void) {
}

int avm_coredump_net_setup(void) {
    return avm_coredump_net_force_setup(conf->source_mac);
}

void avm_coredump_net_teardown(void) {
    avm_coredump_net_shutdown();
}