#include <unistd.h>
#include <pcap/pcap.h>
#include "ipv6_utils.h"
#include "packet_priv.h"
#include "eth_utils.h"
#include "vlan_utils.h"
#include "ipv4_utils.h"
#include "ipv6_utils.h"
#include "tcp_utils.h"
#include "udp_utils.h"

struct options
{
    char *file;
    int print_tshark_format;
    int print_readable_format;
};

static uint64_t number = 0;

#define MAX_BUFF_SIZE 2048
struct buffer
{
    char buff[MAX_BUFF_SIZE];
    int used;
    char elimiter;
};

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

static void packet_to_tshark_format(const struct packet *pkt, uint64_t idx)
{
    /*
        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 buffer buff_proto = {.buff = {0}, .used = 0, .elimiter = ':'};
    struct buffer buff_eth_src = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_eth_dst = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_vlan_id = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_ipv4_src = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_ipv4_dst = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_ipv6_src = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_ipv6_dst = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_tcp_src = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_tcp_dst = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_udp_src = {.buff = {0}, .used = 0, .elimiter = ','};
    struct buffer buff_udp_dst = {.buff = {0}, .used = 0, .elimiter = ','};
    const struct ethhdr *eth_hdr = NULL;
    const struct vlan_hdr *vlan_hdr = NULL;
    const struct ip *ipv4_hdr = NULL;
    const struct ip6_hdr *ipv6_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->type)
        {
        case LAYER_PROTO_ETHER:
            buffer_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));
            buffer_push(&buff_eth_src, tmp_src_buff);
            buffer_push(&buff_eth_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_PWETH:
            buffer_push(&buff_proto, "pwethheuristic:pwethcw");
            break;
        case LAYER_PROTO_PPP:
            buffer_push(&buff_proto, "ppp");
            break;
        case LAYER_PROTO_L2TP:
            buffer_push(&buff_proto, "l2tp");
            break;
        case LAYER_PROTO_VLAN:
            buffer_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);
            buffer_push(&buff_vlan_id, tmp_src_buff);
            break;
        case LAYER_PROTO_PPPOE:
            buffer_push(&buff_proto, "pppoes");
            break;
        case LAYER_PROTO_MPLS:
            buffer_push(&buff_proto, "mpls");
            break;
        case LAYER_PROTO_IPV4:
            buffer_push(&buff_proto, "ip");
            ipv4_hdr = (const struct ip *)layer->hdr_ptr;
            src_addr_v4 = ipv4_hdr_get_src_in_addr(ipv4_hdr);
            dst_addr_v4 = ipv4_hdr_get_dst_in_addr(ipv4_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));
            buffer_push(&buff_ipv4_src, tmp_src_buff);
            buffer_push(&buff_ipv4_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_IPV6:
            buffer_push(&buff_proto, "ipv6");
            ipv6_hdr = (const struct ip6_hdr *)layer->hdr_ptr;
            switch (ipv6_hdr_get_next_header(ipv6_hdr))
            {
            case IPPROTO_HOPOPTS:
                buffer_push(&buff_proto, "ipv6.hopopts");
                break;
            case IPPROTO_ROUTING:
                buffer_push(&buff_proto, "ipv6.routing");
                break;
            case IPPROTO_AH:
                buffer_push(&buff_proto, "ah");
                break;
            case IPPROTO_DSTOPTS:
                buffer_push(&buff_proto, "ipv6.dstopts");
                break;
            case IPPROTO_FRAGMENT:
                buffer_push(&buff_proto, "ipv6.fraghdr");
                break;
            default:
                break;
            }
            src_addr_v6 = ipv6_hdr_get_src_in6_addr(ipv6_hdr);
            dst_addr_v6 = ipv6_hdr_get_dst_in6_addr(ipv6_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));
            buffer_push(&buff_ipv6_src, tmp_src_buff);
            buffer_push(&buff_ipv6_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_IPAH:
            buffer_push(&buff_proto, "ah");
            break;
        case LAYER_PROTO_GRE:
            buffer_push(&buff_proto, "gre");
            break;
        case LAYER_PROTO_UDP:
            buffer_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);
            buffer_push(&buff_udp_src, tmp_src_buff);
            buffer_push(&buff_udp_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_TCP:
            buffer_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);
            buffer_push(&buff_tcp_src, tmp_src_buff);
            buffer_push(&buff_tcp_dst, tmp_dst_buff);
            break;
        case LAYER_PROTO_ICMP:
            buffer_push(&buff_proto, "icmp");
            break;
        case LAYER_PROTO_ICMP6:
            buffer_push(&buff_proto, "icmpv6");
            break;
        case LAYER_PROTO_VXLAN:
            buffer_push(&buff_proto, "vxlan");
            break;
        case LAYER_PROTO_GTPV1_U:
            buffer_push(&buff_proto, "gtp");
            break;
        default:
            buffer_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",
           number,
           buff_proto.buff,
           buff_eth_src.buff,
           buff_eth_dst.buff,
           buff_vlan_id.buff,
           buff_ipv4_src.buff,
           buff_ipv4_dst.buff,
           buff_ipv6_src.buff,
           buff_ipv6_dst.buff,
           buff_tcp_src.buff,
           buff_tcp_dst.buff,
           buff_udp_src.buff,
           buff_udp_dst.buff);
}

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

    struct packet pkt;
    packet_parse(&pkt, (const char *)bytes, h->caplen);
    number++;

    if (opts->print_readable_format)
    {
        printf("\033[0;32m frame=%lu len=%u \033[0m", number, h->caplen);
        packet_print_str(&pkt);
    }

    if (opts->print_tshark_format)
    {
        packet_to_tshark_format(&pkt, number);
    }
}

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("  -r                print readable format\n");
    printf("  -h                print help\n");
    printf("\n");
}

int main(int argc, char **argv)
{
    int opt = 0;
    struct options opts = {0};
    while ((opt = getopt(argc, argv, "f:trh")) != -1)
    {
        switch (opt)
        {
        case 'f':
            opts.file = optarg;
            break;
        case 't':
            opts.print_tshark_format = 1;
            break;
        case 'r':
            opts.print_readable_format = 1;
            break;
        case 'h':
        default:
            usage(argv[0]);
            return 0;
        }
    }

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

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

    return 0;
}