#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <signal.h>

#include <sys/socket.h>

#include <pcap.h>

#include "osfp_common.h"
#include "MESA_osfp.h"
#include "osfp_log.h"
#include "osfp_fingerprint.h"
#include "osfp_score_db.h"

#define DEFAULT_FP_FILE_PATH "./fp.json"

#define ETHERNET_HEADER_LEN     14
#define VLAN_HEADER_LEN         4
#define IPV4_HEADER_LEN         20
#define TCP_HEADER_LEN          20
#define IPV6_HEADER_LEN         40

#define VLAN_MAX_LAYER  2


/* Port is just a uint16_t */
typedef uint16_t Port;
#define SET_PORT(v, p) ((p) = (v))
#define COPY_PORT(a,b) ((b) = (a))

/* Address */
typedef struct Address_ {
    char family;
    union {
        uint32_t        address_un_data32[4]; /* type-specific field */
        uint16_t        address_un_data16[8]; /* type-specific field */
        uint8_t         address_un_data8[16]; /* type-specific field */
        struct in6_addr address_un_in6;
    } address;
} Address;

#define addr_data32 address.address_un_data32
#define addr_data16 address.address_un_data16
#define addr_data8  address.address_un_data8
#define addr_in6addr    address.address_un_in6

/* Set the IPv4 addresses into the Addrs of the Packet.
 * Make sure p->ip4h is initialized and validated.
 *
 * We set the rest of the struct to 0 so we can
 * prevent using memset. */
#define SET_IPV4_SRC_ADDR(p, a) do {                              \
        (a)->family = AF_INET;                                    \
        (a)->addr_data32[0] = (uint32_t)(p)->iph->saddr; \
        (a)->addr_data32[1] = 0;                                  \
        (a)->addr_data32[2] = 0;                                  \
        (a)->addr_data32[3] = 0;                                  \
    } while (0)

#define SET_IPV4_DST_ADDR(p, a) do {                              \
        (a)->family = AF_INET;                                    \
        (a)->addr_data32[0] = (uint32_t)(p)->iph->daddr; \
        (a)->addr_data32[1] = 0;                                  \
        (a)->addr_data32[2] = 0;                                  \
        (a)->addr_data32[3] = 0;                                  \
    } while (0)

/* Set the IPv6 addresses into the Addrs of the Packet.
 * Make sure p->ip6h is initialized and validated. */
#define SET_IPV6_SRC_ADDR(p, a) do {                    \
        (a)->family = AF_INET6;                         \
        (a)->addr_data32[0] = (p)->ip6h->ip6_src.__in6_u.__u6_addr32[0];  \
        (a)->addr_data32[1] = (p)->ip6h->ip6_src.__in6_u.__u6_addr32[1];  \
        (a)->addr_data32[2] = (p)->ip6h->ip6_src.__in6_u.__u6_addr32[2];  \
        (a)->addr_data32[3] = (p)->ip6h->ip6_src.__in6_u.__u6_addr32[3];  \
    } while (0)

#define SET_IPV6_DST_ADDR(p, a) do {                    \
        (a)->family = AF_INET6;                         \
        (a)->addr_data32[0] = (p)->ip6h->ip6_dst.__in6_u.__u6_addr32[0];  \
        (a)->addr_data32[1] = (p)->ip6h->ip6_dst.__in6_u.__u6_addr32[1];  \
        (a)->addr_data32[2] = (p)->ip6h->ip6_dst.__in6_u.__u6_addr32[2];  \
        (a)->addr_data32[3] = (p)->ip6h->ip6_dst.__in6_u.__u6_addr32[3];  \
    } while (0)

#define TCP_GET_RAW_SRC_PORT(tcph)           ntohs((tcph)->source)
#define TCP_GET_RAW_DST_PORT(tcph)           ntohs((tcph)->dest)

#define TCP_GET_SRC_PORT(p)                  TCP_GET_RAW_SRC_PORT((p)->tcph)
#define TCP_GET_DST_PORT(p)                  TCP_GET_RAW_DST_PORT((p)->tcph)

/* Set the TCP ports into the Ports of the Packet.
 * Make sure p->tcph is initialized and validated. */
#define SET_TCP_SRC_PORT(pkt, prt) do {            \
        SET_PORT(TCP_GET_SRC_PORT((pkt)), *(prt)); \
    } while (0)

#define SET_TCP_DST_PORT(pkt, prt) do {            \
        SET_PORT(TCP_GET_DST_PORT((pkt)), *(prt)); \
    } while (0)

#define GET_IPV4_SRC_ADDR_U32(p) ((p)->src.addr_data32[0])
#define GET_IPV4_DST_ADDR_U32(p) ((p)->dst.addr_data32[0])
#define GET_IPV4_SRC_ADDR_PTR(p) ((p)->src.addr_data32)
#define GET_IPV4_DST_ADDR_PTR(p) ((p)->dst.addr_data32)

#define GET_IPV6_SRC_IN6ADDR(p) ((p)->src.addr_in6addr)
#define GET_IPV6_DST_IN6ADDR(p) ((p)->dst.addr_in6addr)
#define GET_IPV6_SRC_ADDR(p) ((p)->src.addr_data32)
#define GET_IPV6_DST_ADDR(p) ((p)->dst.addr_data32)
#define GET_TCP_SRC_PORT(p)  ((p)->sp)
#define GET_TCP_DST_PORT(p)  ((p)->dp)


typedef struct Packet_ {
    struct ethhdr *ethh;
    struct iphdr *iph;
    struct ip6_hdr *ip6h;
    struct tcphdr *tcph;

    char srcip[46];
    char dstip[46];

    Address src;
    Address dst;
    union {
        Port sp;
        // icmp type and code of this packet
        struct {
            uint8_t type;
            uint8_t code;
        } icmp_s;
    };
    union {
        Port dp;
        // icmp type and code of the expected counterpart (for flows)
        struct {
            uint8_t type;
            uint8_t code;
        } icmp_d;
    };

    int vlan_layer;
} Packet;

typedef struct EthernetHdr_ {
    uint8_t eth_dst[6];
    uint8_t eth_src[6];
    uint16_t eth_type;
} __attribute__((__packed__)) EthernetHdr;

unsigned char *fp_file_path;
unsigned char *fp_output_file_path;
FILE *fingerprinting_output_fp;

unsigned int debug_enable;
unsigned char *if_name;
unsigned char *pcap_file_name;
unsigned char *bpf_string;
pcap_t *pcap_handle;

int processed_packet;
int link_type;

struct osfp_profile_counter identify_profile;

unsigned int identify_failed_count = 0;
unsigned int identify_count = 0;
unsigned int result_os_count[OSFP_OS_CLASS_MAX];

void usage(void) {
    fprintf(stderr,
        "Usage: osfp_match [ ...options... ] [ 'filter rule' ]\n"
        "\n"
        "Network interface options:\n"
        "\n"
        "  -i iface  - listen on the specified network interface\n"
        "  -r file   - read offline pcap data from a given file\n"
        "  -f file   - read fingerprint database from 'file' (%s)\n",
        DEFAULT_FP_FILE_PATH);
    exit(1);
}

int packet_decode_tcp(Packet *p, const unsigned char *data, unsigned int len)
{
    int ret = -1;
    int tcp_hdr_len;
    struct tcphdr *tcph;

    if (len < TCP_HEADER_LEN) {
        goto exit;
    }

    tcph = (struct tcphdr *)data;
    tcp_hdr_len = tcph->doff << 2;

    if (len < tcp_hdr_len) {
        goto exit;
    }

    p->tcph = tcph;
    SET_TCP_SRC_PORT(p,&p->sp);
    SET_TCP_DST_PORT(p,&p->dp);

    ret = 0;
exit:
    return ret;
}

int packet_decode_ipv4(Packet *p,        const unsigned char *data, unsigned int len)
{
    int ret = -1;
    int ip_total_len, ip_hdr_len;
    struct iphdr *iph;

    if (len < IPV4_HEADER_LEN) {
        goto exit;
    }

    iph = (struct iphdr *)data;
    ip_total_len = ntohs(iph->tot_len);
    ip_hdr_len = iph->ihl << 2;

    if (ip_hdr_len < IPV4_HEADER_LEN) {
        goto exit;
    }

    if (ip_total_len < ip_hdr_len) {
        goto exit;
    }

    if (len < ip_total_len) {
        goto exit;
    }

    p->iph = iph;
    /* set the address struct */
    SET_IPV4_SRC_ADDR(p,&p->src);
    SET_IPV4_DST_ADDR(p,&p->dst);

    switch (p->iph->protocol) {
        case IPPROTO_TCP:
            packet_decode_tcp(p, data + ip_hdr_len, ip_total_len - ip_hdr_len);
            break;
        default:
            goto exit;
    }

    ret = 0;
exit:
    return ret;
}

int packet_decode_ipv6(Packet *p, const unsigned char *data, unsigned int len)
{
    int ret = -1;
    unsigned short ip6_payload_len;
    unsigned char ip6_nexthdr;
    struct ip6_hdr *ip6h;

    if (len < IPV6_HEADER_LEN) {
        goto exit;
    }

    ip6h = (struct ip6_hdr *)data;
    ip6_payload_len = ntohs(ip6h->ip6_ctlun.ip6_un1.ip6_un1_plen);
    ip6_nexthdr = ip6h->ip6_ctlun.ip6_un1.ip6_un1_nxt;

    if (len < IPV6_HEADER_LEN + ip6_payload_len) {
        goto exit;
    }

    p->ip6h = ip6h;
    SET_IPV6_SRC_ADDR(p,&p->src);
    SET_IPV6_DST_ADDR(p,&p->dst);

    switch (ip6_nexthdr) {
        case IPPROTO_TCP:
            packet_decode_tcp(p, data + IPV6_HEADER_LEN, ip6_payload_len);
            break;
        default:
            goto exit;
    }

    ret = 0;
exit:
    return ret;
}

int pakcet_decode_network_layer(Packet *p,        const unsigned char *data, unsigned int len, unsigned short proto)
{
    switch (proto) {
        case ETH_P_IP: {
            packet_decode_ipv4(p, data, len);
            break;
        }
        case ETH_P_IPV6: {
            packet_decode_ipv6(p, data, len);
            break;
        }
        case ETH_P_8021Q:
            if (p->vlan_layer > VLAN_MAX_LAYER || len < VLAN_HEADER_LEN) {
                return -1;
            }
            unsigned short vlan_proto = ntohs(*(unsigned short*)(data + 2));
            pakcet_decode_network_layer(p, data + VLAN_HEADER_LEN, len - VLAN_HEADER_LEN, vlan_proto);
            break;
        default:
            printf("L3 proto type: %02x not yet supported\n", proto);
            break;
    }

    return 0;
}

int packet_decode_ethernet(Packet *p,        const unsigned char *data, unsigned int len)
{
    int ret = -1;
    unsigned short proto;
    struct ethhdr *ethh;

    if (len < ETHERNET_HEADER_LEN) {
        goto exit;
    }

    ethh = (struct ethhdr *)data;
    proto = ntohs(ethh->h_proto);

    p->ethh = ethh;

    ret = pakcet_decode_network_layer(p, data + ETHERNET_HEADER_LEN, len - ETHERNET_HEADER_LEN, proto);
    if (ret != 0) {
        goto exit;
    }

    return 0;
exit:
    return ret;
}

void packet_decode_link_layer(Packet *p, const unsigned char *data, unsigned int len, int datalink)
{
    switch (datalink) {
        case DLT_EN10MB:
            packet_decode_ethernet(p, data, len);
            break;
        default:
         printf("Datalink type: %02x not yet supported\n", link_type);
         pcap_breakloop(pcap_handle);
         break;
    }
}

void packet_decode(Packet *p, const unsigned char *data, unsigned int len, int datalink)
{
    packet_decode_link_layer(p, data, len, datalink);
}

size_t strlcat(char *dst, const char *src, size_t siz)
{
    register char *d = dst;
    register const char *s = src;
    register size_t n = siz;
    size_t dlen;

    /* Find the end of dst and adjust bytes left but don't go past end */
    while (n-- != 0 && *d != '\0')
        d++;
    dlen = d - dst;
    n = siz - dlen;

    if (n == 0)
        return(dlen + strlen(s));
    while (*s != '\0') {
        if (n != 1) {
            *d++ = *s;
            n--;
        }
        s++;
    }
    *d = '\0';

    return(dlen + (s - src)); /* count does not include NUL */
}

static const char *PrintInetIPv6(const void *src, char *dst, socklen_t size)
{
    int i;
    char s_part[6];
    uint16_t x[8];
    memcpy(&x, src, 16);

    /* current IPv6 format is fixed size */
    if (size < 8 * 5) {
        printf("Too small buffer to write IPv6 address");
        return NULL;
    }
    memset(dst, 0, size);
    for(i = 0; i < 8; i++) {
        snprintf(s_part, sizeof(s_part), "%04x:", htons(x[i]));
        strlcat(dst, s_part, size);
    }
    /* suppress last ':' */
    dst[strlen(dst) - 1] = 0;

    return dst;
}

const char *PrintInet(int af, const void *src, char *dst, socklen_t size)
{
    switch (af) {
        case AF_INET:
            snprintf(dst, size, "%u.%u.%u.%u",
                            ((unsigned char *)src)[0],
                            ((unsigned char *)src)[1],
                            ((unsigned char *)src)[2],
                            ((unsigned char *)src)[3]);
            return dst;
        case AF_INET6:
            /* Format IPv6 without deleting zeroes */
            return PrintInetIPv6(src, dst, size);
        default:
            printf("Unsupported protocol: %d", af);
    }
    return NULL;
}

void example_detect(struct osfp_db *osfp_db, Packet *p)
{
    int ret;
    char str_buf[1024];
    //unsigned char *iph = (unsigned char *)(p->iph != NULL ? (void *)p->iph : (void *)p->ip6h);
    struct iphdr *iph;
    struct ip6_hdr *ip6h;
    struct tcphdr *tcph;
    unsigned int tcph_len;
    struct osfp_result *result = NULL;

    iph = (struct iphdr *)p->iph;
    ip6h = (struct ip6_hdr *)p->ip6h;
    tcph = (struct tcphdr *)p->tcph;
    tcph_len = tcph->doff << 2;

    osfp_profile_cycle(c1);
    osfp_profile_cycle(c2);


    osfp_profile_get_cycle(c1);
    if (iph) {
        result = MESA_osfp_ipv4_identify(osfp_db, iph, tcph, tcph_len);
    } else if (ip6h) {
        result = MESA_osfp_ipv6_identify(osfp_db, ip6h, tcph, tcph_len);
    } else {
        goto exit;
    }
    osfp_profile_get_cycle(c2);
    osfp_profile_counter_update(&identify_profile, c2 - c1);

    identify_count++;

    if (result == NULL) {
        identify_failed_count++;
        printf("osfp header match failed, erro: %s\n", "?");
        goto exit;
    }

    result_os_count[result->likely_os_class]++;

    char *json = MESA_osfp_result_score_detail_export(result);

    if (1) {
        printf("Example ipv4 header detect: --------------------------\n");
        printf("Connection info: %s:%d -> %s:%d\n", p->srcip, p->sp, p->dstip, p->dp);
        printf("Most likely os class: %s\n", MESA_osfp_result_os_name_get(result));
        printf("Details:\n");
        printf("%s\n", json);
    }

exit:
    if (result) {
        MESA_osfp_result_free(result);
    }
    return;
}

void process_packet(char *user, struct pcap_pkthdr *h, u_char *pkt)
{
    int ret;
    struct osfp_db *osfp_db = (struct osfp_db *)user;
    Packet  packet = {0}, *p = &packet;

    // decode packet
    packet_decode(p, pkt, h->len, link_type);
    if (p->tcph == NULL || (p->iph == NULL && p->ip6h == NULL)) {
        goto exit;
    }

    // only for tcp syn request packet
    if (!p->tcph->syn) {
        goto exit;
    }

    if (p->tcph->ack) {
        printf("---------------------------  SYN/ACK\n");
    } else {
        printf("---------------------------  SYN\n");
    }

    if (p->iph) {
        PrintInet(AF_INET, (const void *)&(p->src.addr_data32[0]), p->srcip, sizeof(p->srcip));
        PrintInet(AF_INET, (const void *)&(p->dst.addr_data32[0]), p->dstip, sizeof(p->dstip));
    } else if (p->ip6h) {
        PrintInet(AF_INET6, (const void *)&(p->src.address), p->srcip, sizeof(p->srcip));
        PrintInet(AF_INET6, (const void *)&(p->dst.address), p->dstip, sizeof(p->dstip));
    }

    // tcp/ip header detect example for user
    int i;
    for (i = 0; i < 1; i++) {
        example_detect(osfp_db, p);
    }

    printf("---------------------------    processed packet count %d\n", ++processed_packet);

exit:
    return;
}

static void signal_handler(int signum)
{
    printf("profile identify: avg: %lu max: %lu min: %lu curr: %lu total: %lu count: %lu\n",
            identify_profile.total_cycle / identify_profile.count,
            identify_profile.max_cycle,
            identify_profile.min_cycle,
            identify_profile.curr_cycle,
            identify_profile.total_cycle,
            identify_profile.count);

    osfp_profile_print_stats();

    printf("total %u, failed %u\n",
                identify_count, identify_failed_count);
    int i;
    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        printf("%s: %u\n", osfp_os_class_id_to_name(i), result_os_count[i]);
    }

    exit(0);
}

int main(int argc, char *argv[])
{
    int r;

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    while ((r = getopt(argc, argv, "+f:i:r:o:d")) != -1) {
        switch(r) {
            case 'f':
                if (fp_file_path) {
                    printf("Multiple -f options not supported.\n");
                    exit(1);
                }
                fp_file_path = (unsigned char*)optarg;
                break;
            case 'i':
                if (if_name) {
                    printf("Multiple -i options not supported.\n");
                    exit(1);
                }
                if_name = (unsigned char*)optarg;
                break;
            case 'r':
                if (pcap_file_name) {
                    printf("Multiple -r options not supported.\n");
                    exit(1);
                }
                pcap_file_name = (unsigned char*)optarg;
                break;
            case 'o':
                if (fp_output_file_path) {
                    printf("Multiple -o options not supported.\n");
                    exit(1);
                }
                fp_output_file_path = (unsigned char*)optarg;
                break;
            case 'd':
                debug_enable = 1;
                break;
            default:
                usage();
                break;
        }
    }

    if (optind < argc) {
        if (optind + 1 == argc) {
            bpf_string = argv[optind];
        } else {
            printf("Filter rule must be a single parameter (use quotes).\n");
            exit(1);
        }
    }

    // fingerprinting out file create
    if (fp_output_file_path) {
        fingerprinting_output_fp = fopen(fp_output_file_path, "a+");
        if (!fingerprinting_output_fp) {
            printf("No such file: %s\n", fp_output_file_path);
            exit(1);
        }
    }

    // prepare pcap handle

    char pcap_err[PCAP_ERRBUF_SIZE];

    if (pcap_file_name) {
        if (access((char*)pcap_file_name, R_OK)) {
            printf("No such file: %s\n", pcap_file_name);
            exit(1);
        }
        pcap_handle = pcap_open_offline((char*)pcap_file_name, pcap_err);
        if (pcap_handle == NULL ) {
            printf("Pcap file open failed. File name: %s, Err: %s\n", pcap_file_name, pcap_err);
            exit(1);
        }
    } else if (if_name) {
        pcap_handle = pcap_open_live((char*)if_name, 65535, 1, 5, pcap_err);
        if (pcap_handle == NULL) {
            printf("Pcap live open failed. Interface name: %s, Err: %s\n", if_name, pcap_err);
            exit(1);
        }
    } else {
        usage();
    }

    // setup bpf filter
    if (bpf_string) {
        struct bpf_program bpf_filter;

        if (pcap_compile(pcap_handle, &bpf_filter, bpf_string, 1, 0) < 0) {
            printf("bpf compilation error %s", pcap_geterr(pcap_handle));
            exit(1);
        }
        
        if (pcap_setfilter(pcap_handle, &bpf_filter) < 0) {
            printf("could not set bpf filter %s", pcap_geterr(pcap_handle));
            pcap_freecode(&bpf_filter);
            exit(1);
        }
        pcap_freecode(&bpf_filter);
    }

    // get link type
    link_type = pcap_datalink(pcap_handle);

    // create osfp db
    if (fp_file_path == NULL) {
        fp_file_path = DEFAULT_FP_FILE_PATH;
    }

    if (debug_enable) {
        osfp_log_level_set(OSFP_LOG_LEVEL_DEBUG);
    }

    osfp_profile_set(1);

    struct osfp_db *osfp_db = MESA_osfp_db_new(fp_file_path);
    if (osfp_db == NULL) {
        printf("could not create osfp context. fingerprints file: %s\n", fp_file_path);
        exit(1);
    }

    // loop
    while (1) {
        int r = pcap_dispatch(pcap_handle, 0, (pcap_handler)process_packet, (void*)osfp_db);    
        if (r < 0) {
            printf("error code: %d, error: %s\n", r, pcap_geterr(pcap_handle));
            break;
        }
    }

    // destroy osfp db
    MESA_osfp_db_free(osfp_db);

    return 0;
}