// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/net_util.h"

#include <arpa/inet.h>

#include <netinet/in.h>

#include <sys/socket.h>

#include <sys/types.h>

#include <memory>

#include "zeek/IP.h"

#include "zeek/IPAddr.h"

#include "zeek/Reporter.h"

#include "zeek/3rdparty/doctest.h"

const char* transport_proto_string(TransportProto proto) {

    switch ( proto ) {

        case TRANSPORT_TCP: return "tcp";

        case TRANSPORT_UDP: return "udp";

        case TRANSPORT_ICMP: return "icmp";

        case TRANSPORT_UNKNOWN:

        default: return "unknown";

    }

}

namespace zeek {

uint16_t detail::ip4_in_cksum(const IPAddr& src, const IPAddr& dst, uint8_t next_proto, const uint8_t* data, int len) {

    constexpr auto nblocks = 2;

    detail::checksum_block blocks[nblocks];

    ipv4_pseudo_hdr ph;

    memset(&ph, 0, sizeof(ph));

    src.CopyIPv4(&ph.src);

    dst.CopyIPv4(&ph.dst);

    ph.len = htons(static_cast<uint16_t>(len));

    ph.next_proto = next_proto;

    blocks[0].block = reinterpret_cast<const uint8_t*>(&ph);

    blocks[0].len = sizeof(ph);

    blocks[1].block = data;

    blocks[1].len = len;

    return in_cksum(blocks, nblocks);

}

uint16_t detail::ip6_in_cksum(const IPAddr& src, const IPAddr& dst, uint8_t next_proto, const uint8_t* data, int len) {

    constexpr auto nblocks = 2;

    detail::checksum_block blocks[nblocks];

    ipv6_pseudo_hdr ph;

    memset(&ph, 0, sizeof(ph));

    src.CopyIPv6(&ph.src);

    dst.CopyIPv6(&ph.dst);

    ph.len = htonl(static_cast<uint32_t>(len));

    ph.next_proto = next_proto;

    blocks[0].block = reinterpret_cast<const uint8_t*>(&ph);

    blocks[0].len = sizeof(ph);

    blocks[1].block = data;

    blocks[1].len = len;

    return in_cksum(blocks, nblocks);

}

// Returns the ones-complement checksum of a chunk of 'b' bytes.

int ones_complement_checksum(const void* p, int b, uint32_t sum) {

    const unsigned char* sp = (unsigned char*)p;

    b /= 2; // convert to count of short's

    /* No need for endian conversions. */

    while ( --b >= 0 ) {

        sum += *sp + (*(sp + 1) << 8);

        sp += 2;

    }

    while ( sum > 0xffff )

        sum = (sum & 0xffff) + (sum >> 16);

    return sum;

}

int ones_complement_checksum(const IPAddr& a, uint32_t sum) {

    const uint32_t* bytes;

    int len = a.GetBytes(&bytes);

    return ones_complement_checksum(bytes, len * 4, sum);

}

int icmp_checksum(const struct icmp* icmpp, int len) {

    return detail::in_cksum(reinterpret_cast<const uint8_t*>(icmpp), len);

}

int mobility_header_checksum(const IP_Hdr* ip) {

    const ip6_mobility* mh = ip->MobilityHeader();

    if ( ! mh )

        return 0;

    uint32_t sum = 0;

    uint8_t mh_len = 8 + 8 * mh->ip6mob_len;

    if ( mh_len % 2 == 1 )

        reporter->Weird(ip->SrcAddr(), ip->DstAddr(), "odd_mobility_hdr_len");

    sum = ones_complement_checksum(ip->SrcAddr(), sum);

    sum = ones_complement_checksum(ip->DstAddr(), sum);

    // Note, for IPv6, strictly speaking the protocol and length fields are

    // 32 bits rather than 16 bits.  But because the upper bits are all zero,

    // we get the same checksum either way.

    sum += htons(IPPROTO_MOBILITY);

    sum += htons(mh_len);

    sum = ones_complement_checksum(mh, mh_len, sum);

    return sum;

}

int icmp6_checksum(const struct icmp* icmpp, const IP_Hdr* ip, int len) {

    // ICMP6 uses the same checksum function as ICMP4 but a different

    // pseudo-header over which it is computed.

    return detail::ip6_in_cksum(ip->SrcAddr(), ip->DstAddr(), IPPROTO_ICMPV6, reinterpret_cast<const uint8_t*>(icmpp),

                                len);

}

constexpr uint32_t CLASS_A = 0x00000000;

constexpr uint32_t CLASS_B = 0x80000000;

constexpr uint32_t CLASS_C = 0xc0000000;

constexpr uint32_t CLASS_D = 0xe0000000;

constexpr uint32_t CLASS_E = 0xf0000000;

char addr_to_class(uint32_t addr) {

    auto CHECK_CLASS = [](uint32_t addr, uint32_t cls) { return (addr & cls) == cls; };

    if ( CHECK_CLASS(addr, CLASS_E) )

        return 'E';

    else if ( CHECK_CLASS(addr, CLASS_D) )

        return 'D';

    else if ( CHECK_CLASS(addr, CLASS_C) )

        return 'C';

    else if ( CHECK_CLASS(addr, CLASS_B) )

        return 'B';

    else

        return 'A';

}

const char* fmt_conn_id(const IPAddr& src_addr, uint32_t src_port, const IPAddr& dst_addr, uint32_t dst_port) {

    static char buffer[512];

    snprintf(buffer, sizeof(buffer), "%s:%d > %s:%d", std::string(src_addr).c_str(), src_port,

             std::string(dst_addr).c_str(), dst_port);

    return buffer;

}

const char* fmt_conn_id(const uint32_t* src_addr, uint32_t src_port, const uint32_t* dst_addr, uint32_t dst_port) {

    IPAddr src(IPv6, src_addr, IPAddr::Network);

    IPAddr dst(IPv6, dst_addr, IPAddr::Network);

    return fmt_conn_id(src, src_port, dst, dst_port);

}

TEST_CASE("fmt_mac") {

    auto my_fmt_mac = [](const char* m, int len) {

        // allow working with literal strings

        return fmt_mac(reinterpret_cast<const unsigned char*>(m), len);

    };

    CHECK(my_fmt_mac("", 0) == "");

    CHECK(my_fmt_mac("\x01\x02\x03\x04\x05\x06", 4) == "");

    CHECK(my_fmt_mac("\x01\x02\x03\x04\x05\x06", 6) == "01:02:03:04:05:06");

    // NOLINTNEXTLINE(bugprone-string-literal-with-embedded-nul)

    CHECK(my_fmt_mac("\x01\x02\x03\x04\x05\x06\x00\x00", 8) == "01:02:03:04:05:06");

    CHECK(my_fmt_mac("\x01\x02\x03\x04\x05\x06\x07\x08", 8) == "01:02:03:04:05:06:07:08");

    CHECK(my_fmt_mac("\x08\x07\x06\x05\x04\x03\x02\x01", 8) == "08:07:06:05:04:03:02:01");

}

std::string fmt_mac(const unsigned char* m, int len) {

    auto [mac_bytes, slen] = fmt_mac_bytes(m, len);

    return reinterpret_cast<char*>(mac_bytes.get());

}

TEST_CASE("fmt_mac_bytes") {

    auto my_fmt_mac_bytes = [](const char* m, int len) {

        // allow working with literal strings

        return fmt_mac_bytes(reinterpret_cast<const unsigned char*>(m), len);

    };

    auto [buf1, len1] = my_fmt_mac_bytes("\x01\x02\x03\x04\x05\x06", 4);

    CHECK(len1 == 0);

    CHECK(memcmp(buf1.get(), "", 1) == 0); // still null terminated

    auto [buf2, len2] = my_fmt_mac_bytes("\x01\x02\x03\x04\x05\x06", 6);

    CHECK(len2 == 2 * 6 + 5);

    CHECK(memcmp(buf2.get(), "01:02:03:04:05:06", len2 + 1) == 0);

    // NOLINTNEXTLINE(bugprone-string-literal-with-embedded-nul)

    auto [buf3, len3] = my_fmt_mac_bytes("\x01\x02\x03\x04\x05\x06\x00\x00", 8);

    CHECK(len3 == 2 * 6 + 5);

    CHECK(memcmp(buf3.get(), "01:02:03:04:05:06", len3 + 1) == 0);

    // Check for no memory overreads

    auto [buf4, len4] = my_fmt_mac_bytes("\x01\x02\x03\x04\x05\x06\x07\x08\xff\xff", 42);

    CHECK(len4 == 2 * 8 + 7);

    CHECK(memcmp(buf4.get(), "01:02:03:04:05:06:07:08", len4 + 1) == 0);

}

std::pair<std::unique_ptr<uint8_t[]>, int> fmt_mac_bytes(const unsigned char* m, int len) {

    if ( len < 8 && len != 6 ) {

        auto buf = std::make_unique<uint8_t[]>(1);

        buf[0] = '\0';

        return {std::move(buf), 0};

    }

    int elen = (len == 6 || (m[6] == 0 && m[7] == 0)) ? 6 : 8;

    auto slen = 2 * elen + (elen - 1);

    auto buf = std::make_unique<uint8_t[]>(slen + 1);

    auto* bufp = reinterpret_cast<char*>(buf.get());

    for ( int i = 0; i < elen; i++ ) {

        zeek::util::bytetohex(m[i], &bufp[i * 2 + i]);

        if ( i < elen - 1 )

            bufp[i * 2 + i + 2] = ':';

    }

    bufp[slen] = '\0';

    return {std::move(buf), slen};

}

uint32_t extract_uint32(const u_char* data) {

    uint32_t val;

    val = data[0] << 24;

    val |= data[1] << 16;

    val |= data[2] << 8;

    val |= data[3];

    return val;

}

} // namespace zeek