/*
 *
 * Copyright (c) 2007-2016 The University of Waikato, Hamilton, New Zealand.
 * All rights reserved.
 *
 * This file is part of libtrace.
 *
 * This code has been developed by the University of Waikato WAND
 * research group. For further information please see http://www.wand.net.nz/
 *
 * libtrace is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * libtrace is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *
 */

#include "libtrace_int.h"
#include "libtrace.h"
#include "protocols.h"
#include "checksum.h"
#include <stdlib.h>
#include <stdio.h> // fprintf
#include <string.h>

/* This file contains all the protocol decoding functions for transport layer
 * protocols. This includes functions for access port numbers.
 *
 * Supported protocols include (but are not limited to):
 * 	TCP
 * 	UDP
 * 	ICMP
 */

/* Get the size of the payload as it was in the original packet, i.e. prior
 * to any truncation.
 *
 * Basically, wire length minus the packet headers.
 *
 * Currently only supports IP (v4 and v6) and TCP, UDP and ICMP. Will return
 * 0 if an unsupported protocol header is encountered, or if one of the
 * headers is truncated.
 */
DLLEXPORT size_t trace_get_payload_length(const libtrace_packet_t *packet)
{

    void *layer;
    uint16_t ethertype;
    uint8_t proto;
    uint32_t rem;
    libtrace_ip_t *ip;
    libtrace_ip6_t *ip6;
    libtrace_tcp_t *tcp;
    size_t len = 0;
    uint8_t iplenzero = 0;

    /* Just use the cached length if we can */
    if (packet->cached.payload_length != -1)
        return packet->cached.payload_length;

    /* Set to zero so that we can return early without having to
     * worry about forgetting to update the cached value */
    ((libtrace_packet_t *)packet)->cached.payload_length = 0;
    layer = trace_get_layer3(packet, &ethertype, &rem);
    if (!layer)
        return 0;
    switch (ethertype) {
    case TRACE_ETHERTYPE_IP:
        ip = (libtrace_ip_t *)layer;
        if (rem < sizeof(libtrace_ip_t))
            return 0;
        if (ntohs(ip->ip_len) == 0) {
            iplenzero = 1;
            break;
        }
        len = ntohs(ip->ip_len) - (4 * ip->ip_hl);

        /* Deal with v6 within v4 */
        if (ip->ip_p == TRACE_IPPROTO_IPV6)
            len -= sizeof(libtrace_ip6_t);

        break;
    case TRACE_ETHERTYPE_IPV6:
        ip6 = (libtrace_ip6_t *)layer;
        if (rem < sizeof(libtrace_ip6_t))
            return 0;
        len = ntohs(ip6->plen);
        if (len == 0) {
            iplenzero = 1;
        }
        break;
    default:
        return 0;
    }

    if (iplenzero) {
        /* deal with cases where IP length is zero due to
         * hardware segmentation offload */
        uint8_t *iplayer, *pktstart;
        libtrace_linktype_t linktype;
        uint32_t rem;

        iplayer = (uint8_t *)layer;
        pktstart = (uint8_t *)trace_get_packet_buffer(packet, &linktype, &rem);

        len = rem - (iplayer - pktstart);
        if (ethertype == TRACE_ETHERTYPE_IP) {
            ip = (libtrace_ip_t *)layer;
            len -= (4 * ip->ip_hl);
        } else {
            len -= sizeof(libtrace_ip6_t);
        }
    }

    layer = trace_get_transport(packet, &proto, &rem);
    if (!layer)
        return 0;

    switch (proto) {
    case TRACE_IPPROTO_TCP:
        if (rem < sizeof(libtrace_tcp_t))
            return 0;
        tcp = (libtrace_tcp_t *)layer;

        if (len < (size_t)(4 * tcp->doff))
            return 0;

        len -= (4 * tcp->doff);
        break;
    case TRACE_IPPROTO_UDP:
        if (rem < sizeof(libtrace_udp_t))
            return 0;
        if (len < sizeof(libtrace_udp_t))
            return 0;
        len -= sizeof(libtrace_udp_t);
        break;
    case TRACE_IPPROTO_ICMP:
        if (rem < sizeof(libtrace_icmp_t))
            return 0;
        if (len < sizeof(libtrace_icmp_t))
            return 0;
        len -= sizeof(libtrace_icmp_t);
        break;
    case TRACE_IPPROTO_ICMPV6:
        if (rem < sizeof(libtrace_icmp6_t))
            return 0;
        if (len < sizeof(libtrace_icmp6_t))
            return 0;
        len -= sizeof(libtrace_icmp6_t);
        break;

    default:
        return 0;
    }

    ((libtrace_packet_t *)packet)->cached.payload_length = len;
    return len;
}

DLLEXPORT void *trace_get_transport(const libtrace_packet_t *packet,
                                    uint8_t *proto, uint32_t *remaining)
{
    uint8_t dummy_proto;
    uint16_t ethertype;
    uint32_t dummy_remaining;
    void *transport;

    if (!proto)
        proto = &dummy_proto;

    if (!remaining)
        remaining = &dummy_remaining;

    if (packet->cached.l4_header) {
        /*
        void *link;
        libtrace_linktype_t linktype;
        link = trace_get_packet_buffer(packet, &linktype, remaining);
        if (!link)
                return NULL;
        */
        *proto = packet->cached.transport_proto;
        *remaining = packet->cached.l4_remaining;
        return packet->cached.l4_header;
    }

    transport = trace_get_layer3(packet, &ethertype, remaining);

    if (!transport || *remaining == 0)
        return NULL;

    switch (ethertype) {
    case TRACE_ETHERTYPE_IP: /* IPv4 */
        transport = trace_get_payload_from_ip((libtrace_ip_t *)transport, proto,
                                              remaining);
        /* IPv6 */
        if (transport && *proto == TRACE_IPPROTO_IPV6) {
            transport = trace_get_payload_from_ip6((libtrace_ip6_t *)transport,
                                                   proto, remaining);
        }
        break;
    case TRACE_ETHERTYPE_IPV6: /* IPv6 */
        transport = trace_get_payload_from_ip6((libtrace_ip6_t *)transport,
                                               proto, remaining);
        break;
    default:
        *proto = 0;
        transport = NULL;
        break;
    }

    ((libtrace_packet_t *)packet)->cached.transport_proto = *proto;
    ((libtrace_packet_t *)packet)->cached.l4_header = transport;
    ((libtrace_packet_t *)packet)->cached.l4_remaining = *remaining;

    return transport;
}

DLLEXPORT libtrace_tcp_t *trace_get_tcp(libtrace_packet_t *packet)
{
    uint8_t proto;
    uint32_t rem = 0;
    libtrace_tcp_t *tcp;

    tcp = (libtrace_tcp_t *)trace_get_transport(packet, &proto, &rem);

    if (!tcp || proto != TRACE_IPPROTO_TCP)
        return NULL;

    /* We should return NULL if there isn't a full TCP header, because the
     * caller has no way of telling how much of a TCP header we have
     * returned - use trace_get_transport() if you want to deal with
     * partial headers
     *
     * NOTE: We're not going to insist that all the TCP options are present
     * as well, because lots of traces are snapped after 20 bytes of TCP
     * header and I don't really want to break libtrace programs that
     * use this function to process those traces */

    if (rem < sizeof(libtrace_tcp_t))
        return NULL;

    return (libtrace_tcp_t *)tcp;
}

DLLEXPORT libtrace_tcp_t *trace_get_tcp_from_ip(libtrace_ip_t *ip,
                                                uint32_t *remaining)
{
    libtrace_tcp_t *tcpptr = 0;

    if (ip->ip_p == TRACE_IPPROTO_TCP) {
        tcpptr =
            (libtrace_tcp_t *)trace_get_payload_from_ip(ip, NULL, remaining);
    }

    return tcpptr;
}

DLLEXPORT libtrace_udp_t *trace_get_udp(libtrace_packet_t *packet)
{
    uint8_t proto;
    uint32_t rem = 0;
    libtrace_udp_t *udp;

    udp = (libtrace_udp_t *)trace_get_transport(packet, &proto, &rem);

    if (!udp || proto != TRACE_IPPROTO_UDP)
        return NULL;

    /* Make sure we return a full UDP header as the caller has no way of
     * telling how much of the packet is remaining */
    if (rem < sizeof(libtrace_udp_t))
        return NULL;

    return udp;
}

DLLEXPORT libtrace_udp_t *trace_get_udp_from_ip(libtrace_ip_t *ip,
                                                uint32_t *remaining)
{
    libtrace_udp_t *udpptr = 0;

    if (ip->ip_p == TRACE_IPPROTO_UDP) {
        udpptr =
            (libtrace_udp_t *)trace_get_payload_from_ip(ip, NULL, remaining);
    }

    return udpptr;
}

DLLEXPORT libtrace_icmp_t *trace_get_icmp(libtrace_packet_t *packet)
{
    uint8_t proto;
    uint32_t rem = 0;
    libtrace_icmp_t *icmp;

    icmp = (libtrace_icmp_t *)trace_get_transport(packet, &proto, &rem);

    if (!icmp || proto != TRACE_IPPROTO_ICMP)
        return NULL;

    /* Make sure we return a full ICMP header as the caller has no way of
     * telling how much of the packet is remaining */
    if (rem < sizeof(libtrace_icmp_t))
        return NULL;

    return icmp;
}

DLLEXPORT libtrace_icmp6_t *trace_get_icmp6(libtrace_packet_t *packet)
{
    uint8_t proto;
    uint32_t rem = 0;
    libtrace_icmp6_t *icmp;

    icmp = (libtrace_icmp6_t *)trace_get_transport(packet, &proto, &rem);

    if (!icmp || proto != TRACE_IPPROTO_ICMPV6)
        return NULL;

    /* Make sure we return a full ICMP header as the caller has no way of
     * telling how much of the packet is remaining */
    if (rem < sizeof(libtrace_icmp6_t))
        return NULL;

    return icmp;
}

DLLEXPORT libtrace_icmp_t *trace_get_icmp_from_ip(libtrace_ip_t *ip,
                                                  uint32_t *remaining)
{
    libtrace_icmp_t *icmpptr = 0;

    if (ip->ip_p == TRACE_IPPROTO_ICMP) {
        icmpptr =
            (libtrace_icmp_t *)trace_get_payload_from_ip(ip, NULL, remaining);
    }

    return icmpptr;
}

DLLEXPORT void *trace_get_payload_from_udp(libtrace_udp_t *udp,
                                           uint32_t *remaining)
{
    if (remaining) {
        if (*remaining < sizeof(libtrace_udp_t)) {
            *remaining = 0;
            return NULL;
        }
        *remaining -= sizeof(libtrace_udp_t);
    }
    return (void *)((char *)udp + sizeof(libtrace_udp_t));
}

DLLEXPORT void *trace_get_payload_from_tcp(libtrace_tcp_t *tcp,
                                           uint32_t *remaining)
{
    unsigned int dlen = tcp->doff * 4;
    if (remaining) {
        if (*remaining < dlen) {
            *remaining = 0;
            return NULL;
        }
        *remaining -= dlen;
    }
    return (void *)((char *)tcp + dlen);
}

DLLEXPORT void *trace_get_payload_from_icmp(libtrace_icmp_t *icmp,
                                            uint32_t *remaining)
{
    if (remaining) {
        if (*remaining < sizeof(libtrace_icmp_t)) {
            *remaining = 0;
            return NULL;
        }
        *remaining -= sizeof(libtrace_icmp_t);
    }
    return (char *)icmp + sizeof(libtrace_icmp_t);
}

DLLEXPORT void *trace_get_payload_from_icmp6(libtrace_icmp6_t *icmp,
                                             uint32_t *remaining)
{
    if (remaining) {
        if (*remaining < sizeof(libtrace_icmp6_t)) {
            *remaining = 0;
            return NULL;
        }
        *remaining -= sizeof(libtrace_icmp6_t);
    }
    return (char *)icmp + sizeof(libtrace_icmp6_t);
}

/* Return the source port
 */
DLLEXPORT uint16_t trace_get_source_port(const libtrace_packet_t *packet)
{
    uint32_t remaining;
    uint8_t proto;
    struct ports_t *port;
    uint16_t fragoff;
    uint8_t more;

    fragoff = trace_get_fragment_offset(packet, &more);

    /* If we're not the first fragment, we're unlikely to be able
     * to get any useful port numbers from this packet.
     */
    if (fragoff != 0)
        return 0;

    port = (struct ports_t *)trace_get_transport((libtrace_packet_t *)packet,
                                                 &proto, &remaining);

    /* Snapped too early */
    if (remaining < 2)
        return 0;

    /* ICMP *technically* doesn't have ports */
    if (proto == TRACE_IPPROTO_ICMP || proto == TRACE_IPPROTO_ICMPV6)
        return 0;

    if (port)
        return ntohs(port->src);
    else
        return 0;
}

/* Same as get_source_port except use the destination port */
DLLEXPORT uint16_t trace_get_destination_port(const libtrace_packet_t *packet)
{
    uint32_t remaining;
    uint8_t proto;
    struct ports_t *port;
    uint16_t fragoff;
    uint8_t more;

    fragoff = trace_get_fragment_offset(packet, &more);

    /* If we're not the first fragment, we're unlikely to be able
     * to get any useful port numbers from this packet.
     */
    if (fragoff != 0)
        return 0;

    port = (struct ports_t *)trace_get_transport((libtrace_packet_t *)packet,
                                                 &proto, &remaining);
    /* Snapped too early */
    if (remaining < 4)
        return 0;

    /* ICMP *technically* doesn't have ports */
    if (proto == TRACE_IPPROTO_ICMP || proto == TRACE_IPPROTO_ICMPV6)
        return 0;

    if (port)
        return ntohs(port->dst);
    else
        return 0;
}

DLLEXPORT uint16_t *trace_checksum_transport(libtrace_packet_t *packet,
                                             uint16_t *csum)
{

    void *header = NULL;
    uint16_t ethertype;
    uint32_t remaining;
    uint32_t sum = 0;
    uint8_t proto = 0;
    char *csum_ptr = NULL;
    int plen = 0;

    uint8_t safety[65536];
    uint8_t *ptr = safety;

    header = trace_get_layer3(packet, &ethertype, &remaining);

    if (header == NULL)
        return NULL;

    if (ethertype == TRACE_ETHERTYPE_IP) {
        libtrace_ip_t *ip = (libtrace_ip_t *)header;

        if (remaining < sizeof(libtrace_ip_t))
            return NULL;

        sum = ipv4_pseudo_checksum(ip);

    } else if (ethertype == TRACE_ETHERTYPE_IPV6) {
        libtrace_ip6_t *ip = (libtrace_ip6_t *)header;

        if (remaining < sizeof(libtrace_ip6_t))
            return 0;

        sum = ipv6_pseudo_checksum(ip);
    }

    header = trace_get_transport(packet, &proto, &remaining);

    if (proto == TRACE_IPPROTO_TCP) {
        libtrace_tcp_t *tcp = (libtrace_tcp_t *)header;
        header = trace_get_payload_from_tcp(tcp, &remaining);

        csum_ptr = (char *)(&tcp->check);

        memcpy(ptr, tcp, tcp->doff * 4);

        tcp = (libtrace_tcp_t *)ptr;
        tcp->check = 0;

        ptr += (tcp->doff * 4);
    }

    else if (proto == TRACE_IPPROTO_UDP) {

        libtrace_udp_t *udp = (libtrace_udp_t *)header;
        header = trace_get_payload_from_udp(udp, &remaining);

        csum_ptr = (char *)(&udp->check);
        memcpy(ptr, udp, sizeof(libtrace_udp_t));

        udp = (libtrace_udp_t *)ptr;
        udp->check = 0;

        ptr += sizeof(libtrace_udp_t);
    }

    else if (proto == TRACE_IPPROTO_ICMP) {
        /* ICMP doesn't use the pseudo header */
        sum = 0;

        libtrace_icmp_t *icmp = (libtrace_icmp_t *)header;
        header = trace_get_payload_from_icmp(icmp, &remaining);

        csum_ptr = (char *)(&icmp->checksum);
        memcpy(ptr, icmp, sizeof(libtrace_icmp_t));

        icmp = (libtrace_icmp_t *)ptr;
        icmp->checksum = 0;

        ptr += sizeof(libtrace_icmp_t);

    } else {
        return NULL;
    }

    sum += add_checksum(safety, (uint16_t)(ptr - safety));

    plen = trace_get_payload_length(packet);
    if (plen < 0)
        return NULL;

    if (remaining < (uint32_t)plen)
        return NULL;

    if (header == NULL)
        return NULL;

    sum += add_checksum(header, (uint16_t)plen);
    *csum = ntohs(finish_checksum(sum));

    return (uint16_t *)csum_ptr;
}

DLLEXPORT void *trace_get_payload_from_gre(libtrace_gre_t *gre,
                                           uint32_t *remaining)
{
    uint8_t flags = ntohs(gre->flags);
    uint32_t size = 4; /* GRE is 4 bytes long by default */
    if (remaining && *remaining < size) {
        *remaining = 0;
        return NULL;
    }

    if ((flags & LIBTRACE_GRE_FLAG_VERMASK) == LIBTRACE_GRE_PPTP_VERSION) {
        size += 4;

        if ((flags & LIBTRACE_GRE_FLAG_SEQ) != 0) {
            size += 4;
        }
        if ((flags & LIBTRACE_GRE_FLAG_ACK) != 0) {
            size += 4;
        }
    } else {

        if ((ntohs(gre->flags) & LIBTRACE_GRE_FLAG_CHECKSUM) != 0) {
            size += 4; /* An extra 4 bytes. */
        }

        if ((ntohs(gre->flags) & LIBTRACE_GRE_FLAG_KEY) != 0) {
            size += 4; /* An extra 4 bytes. */
        }

        if ((ntohs(gre->flags) & LIBTRACE_GRE_FLAG_SEQ) != 0) {
            size += 4; /* An extra 4 bytes. */
        }
    }

    if (remaining) {
        if (*remaining < size) {
            *remaining = 0;
            return NULL;
        }
        *remaining -= size;
    }
    return (char *)gre + size;
}
