#include <unistd.h>
#include <pcap/pcap.h>

#include "eth_utils.h"
#include "vlan_utils.h"
#include "ip4_utils.h"
#include "ip6_utils.h"
#include "tcp_utils.h"
#include "udp_utils.h"
#include "packet_def.h"
#include "packet_layer.h"
#include "packet_parse.h"
#include "packet_dump.h"
#include "packet_utils.h"

#define MAX_BUFF_SIZE 2048
#define PRINT_GREEN(fmt, ...) printf("\033[0;32m" fmt "\033[0m\n", ##__VA_ARGS__)
#define PRINT_RED(fmt, ...) printf("\033[0;31m" fmt "\033[0m\n", ##__VA_ARGS__)

struct str_buff
{
    char data[MAX_BUFF_SIZE];
    int used;
    char elimiter;
};

struct runtime
{
    char *pcap_file;
    uint64_t pcap_count;
    int tshark_format;
    int print_verbose;
};

static void str_buff_push(struct str_buff *buff, const char *str)
{
    int len = strlen(str);
    if (buff->used + len + 1 >= MAX_BUFF_SIZE)
    {
        return;
    }
    if (buff->used)
    {
        buff->data[buff->used++] = buff->elimiter;
    }
    memcpy(buff->data + buff->used, str, len);
    buff->used += len;
}

static void tshark_format(const struct runtime *rte, const struct packet *pkt)
{
    /*
        tshark -r ${pcap} -T fields \
        -e frame.number \
        -e frame.protocols \
        -e eth.src \
        -e eth.dst \
        -e vlan.id \
        -e ip.src \
        -e ip.dst \
        -e ipv6.src \
        -e ipv6.dst \
        -e tcp.srcport \
        -e tcp.dstport \
        -e udp.srcport \
        -e udp.dstport \
        >> tshark_output.txt
    */

    struct str_buff buff_proto = {.data = {0}, .used = 0, .elimiter = ':'};
    struct str_buff buff_eth_src = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_eth_dst = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_vlan_id = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_ipv4_src = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_ipv4_dst = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_ipv6_src = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_ipv6_dst = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_tcp_src = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_tcp_dst = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_udp_src = {.data = {0}, .used = 0, .elimiter = ','};
    struct str_buff buff_udp_dst = {.data = {0}, .used = 0, .elimiter = ','};

    const struct ethhdr *eth_hdr = NULL;
    const struct vlan_hdr *vlan_hdr = NULL;
    const struct ip *ip4_hdr = NULL;
    const struct ip6_hdr *ip6_hdr = NULL;
    const struct tcphdr *tcp_hdr = NULL;
    const struct udphdr *udp_hdr = NULL;

    struct in_addr src_addr_v4 = {0};
    struct in_addr dst_addr_v4 = {0};
    struct in6_addr src_addr_v6 = {0};
    struct in6_addr dst_addr_v6 = {0};

    uint16_t src_port = 0;
    uint16_t dst_port = 0;
    uint16_t vlan_id = 0;

    char tmp_src_buff[256] = {0};
    char tmp_dst_buff[256] = {0};

    int num = packet_get_layer_count(pkt);
    for (int i = 0; i < num; i++)
    {
        memset(tmp_src_buff, 0, sizeof(tmp_src_buff));
        memset(tmp_dst_buff, 0, sizeof(tmp_dst_buff));
        const struct raw_layer *layer = packet_get_raw_layer(pkt, i);
        switch (layer->proto)
        {
        case LAYER_PROTO_ETHER:
            str_buff_push(&buff_proto, "eth:ethertype");

            eth_hdr = (const struct ethhdr *)layer->hdr_ptr;
            eth_hdr_get_source(eth_hdr, tmp_src_buff, sizeof(tmp_src_buff));
            eth_hdr_get_dest(eth_hdr, tmp_dst_buff, sizeof(tmp_dst_buff));
            str_buff_push(&buff_eth_src, tmp_src_buff);
            str_buff_push(&buff_eth_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_PWETH:
            str_buff_push(&buff_proto, "pwethheuristic:pwethcw");
            break;
        case LAYER_PROTO_PPP:
            str_buff_push(&buff_proto, "ppp");
            break;
        case LAYER_PROTO_L2TP:
            str_buff_push(&buff_proto, "l2tp");
            break;
        case LAYER_PROTO_VLAN:
            str_buff_push(&buff_proto, "vlan:ethertype");

            vlan_hdr = (const struct vlan_hdr *)layer->hdr_ptr;
            vlan_id = vlan_hdr_get_vid(vlan_hdr);
            snprintf(tmp_src_buff, sizeof(tmp_src_buff), "%u", vlan_id);
            str_buff_push(&buff_vlan_id, tmp_src_buff);
            break;
        case LAYER_PROTO_PPPOE:
            str_buff_push(&buff_proto, "pppoes");
            break;
        case LAYER_PROTO_MPLS:
            str_buff_push(&buff_proto, "mpls");
            break;
        case LAYER_PROTO_IPV4:
            str_buff_push(&buff_proto, "ip");

            ip4_hdr = (const struct ip *)layer->hdr_ptr;
            src_addr_v4 = ip4_hdr_get_src_in_addr(ip4_hdr);
            dst_addr_v4 = ip4_hdr_get_dst_in_addr(ip4_hdr);
            inet_ntop(AF_INET, &src_addr_v4, tmp_src_buff, sizeof(tmp_src_buff));
            inet_ntop(AF_INET, &dst_addr_v4, tmp_dst_buff, sizeof(tmp_dst_buff));
            str_buff_push(&buff_ipv4_src, tmp_src_buff);
            str_buff_push(&buff_ipv4_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_IPV6:
            str_buff_push(&buff_proto, "ipv6");
            ip6_hdr = (const struct ip6_hdr *)layer->hdr_ptr;
            switch (ip6_hdr_get_next_header(ip6_hdr))
            {
            case IPPROTO_HOPOPTS:
                str_buff_push(&buff_proto, "ipv6.hopopts");
                break;
            case IPPROTO_ROUTING:
                str_buff_push(&buff_proto, "ipv6.routing");
                break;
            case IPPROTO_AH:
                str_buff_push(&buff_proto, "ah");
                break;
            case IPPROTO_DSTOPTS:
                str_buff_push(&buff_proto, "ipv6.dstopts");
                break;
            case IPPROTO_FRAGMENT:
                str_buff_push(&buff_proto, "ipv6.fraghdr");
                break;
            default:
                break;
            }

            src_addr_v6 = ip6_hdr_get_src_in6_addr(ip6_hdr);
            dst_addr_v6 = ip6_hdr_get_dst_in6_addr(ip6_hdr);
            inet_ntop(AF_INET6, &src_addr_v6, tmp_src_buff, sizeof(tmp_src_buff));
            inet_ntop(AF_INET6, &dst_addr_v6, tmp_dst_buff, sizeof(tmp_dst_buff));
            str_buff_push(&buff_ipv6_src, tmp_src_buff);
            str_buff_push(&buff_ipv6_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_IPAH:
            str_buff_push(&buff_proto, "ah");
            break;
        case LAYER_PROTO_GRE:
            str_buff_push(&buff_proto, "gre");
            break;
        case LAYER_PROTO_UDP:
            str_buff_push(&buff_proto, "udp");

            udp_hdr = (const struct udphdr *)layer->hdr_ptr;
            src_port = udp_hdr_get_src_port(udp_hdr);
            dst_port = udp_hdr_get_dst_port(udp_hdr);
            snprintf(tmp_src_buff, sizeof(tmp_src_buff), "%u", src_port);
            snprintf(tmp_dst_buff, sizeof(tmp_dst_buff), "%u", dst_port);
            str_buff_push(&buff_udp_src, tmp_src_buff);
            str_buff_push(&buff_udp_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_TCP:
            str_buff_push(&buff_proto, "tcp");

            tcp_hdr = (const struct tcphdr *)layer->hdr_ptr;
            src_port = tcp_hdr_get_src_port(tcp_hdr);
            dst_port = tcp_hdr_get_dst_port(tcp_hdr);
            snprintf(tmp_src_buff, sizeof(tmp_src_buff), "%u", src_port);
            snprintf(tmp_dst_buff, sizeof(tmp_dst_buff), "%u", dst_port);
            str_buff_push(&buff_tcp_src, tmp_src_buff);
            str_buff_push(&buff_tcp_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_ICMP:
            str_buff_push(&buff_proto, "icmp");
            break;
        case LAYER_PROTO_ICMP6:
            str_buff_push(&buff_proto, "icmpv6");
            break;
        case LAYER_PROTO_VXLAN:
            str_buff_push(&buff_proto, "vxlan");
            break;
        case LAYER_PROTO_GTP_C:
        case LAYER_PROTO_GTP_U:
            str_buff_push(&buff_proto, "gtp");
            break;
        default:
            str_buff_push(&buff_proto, "unknown");
            break;
        }
    }

    printf("%lu\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
           rte->pcap_count,
           buff_proto.data,
           buff_eth_src.data,
           buff_eth_dst.data,
           buff_vlan_id.data,
           buff_ipv4_src.data,
           buff_ipv4_dst.data,
           buff_ipv6_src.data,
           buff_ipv6_dst.data,
           buff_tcp_src.data,
           buff_tcp_dst.data,
           buff_udp_src.data,
           buff_udp_dst.data);
}

static void packet_handler(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes)
{
    struct runtime *rte = (struct runtime *)user;

    struct packet pkt;
    memset(&pkt, 0, sizeof(pkt));
    packet_parse(&pkt, (const char *)bytes, h->caplen);
    rte->pcap_count++;

    if (rte->print_verbose)
    {
        PRINT_GREEN("frame=%lu len=%u", rte->pcap_count, h->caplen);
        packet_dump_stdio(&pkt);
    }

    if (rte->tshark_format)
    {
        tshark_format(rte, &pkt);
    }
}

static void usage(char *cmd)
{
    printf("Usage: %s\n", cmd);
    printf("Options:\n");
    printf("  -f <pcap file>    pcap file\n");
    printf("  -t                print tshark format\n");
    printf("  -v                print verbose info\n");
    printf("  -h                print help\n");
    printf("\n");
}

int main(int argc, char **argv)
{
    int opt = 0;
    struct runtime rte = {};
    while ((opt = getopt(argc, argv, "f:tvh")) != -1)
    {
        switch (opt)
        {
        case 'f':
            rte.pcap_file = optarg;
            break;
        case 't':
            rte.tshark_format = 1;
            break;
        case 'v':
            rte.print_verbose = 1;
            break;
        case 'h':
        default:
            usage(argv[0]);
            return 0;
        }
    }

    if (rte.pcap_file == NULL)
    {
        usage(argv[0]);
        return -1;
    }

    if (rte.print_verbose)
    {
        PRINT_GREEN("pcap=%s", rte.pcap_file);
    }

    pcap_t *pcap = pcap_open_offline(rte.pcap_file, NULL);
    if (pcap == NULL)
    {
        printf("pcap_open_offline() failed\n");
        return -1;
    }
    pcap_loop(pcap, -1, packet_handler, (u_char *)&rte);
    pcap_close(pcap);

    return 0;
}