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

#pragma once

#include <arpa/inet.h>
#include <netinet/in.h>
#include <cstring>
#include <memory>
#include <string>

#include "zeek/threading/SerialTypes.h"

using in4_addr = in_addr;

namespace zeek {

class String;
struct ConnTuple;
class Val;

namespace detail {

class HashKey;

class ConnKey {
public:
    in6_addr ip1;
    in6_addr ip2;
    uint16_t port1 = 0;
    uint16_t port2 = 0;
    uint8_t transport;
    bool valid = true;

    ConnKey(const IPAddr& src, const IPAddr& dst, uint16_t src_port, uint16_t dst_port, uint8_t proto, bool one_way);
    ConnKey(const ConnTuple& conn);
    ConnKey(const ConnKey& rhs) { *this = rhs; }
    ConnKey(Val* v);

    bool operator<(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) < 0; }
    bool operator<=(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) <= 0; }
    bool operator==(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) == 0; }
    bool operator!=(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) != 0; }
    bool operator>=(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) >= 0; }
    bool operator>(const ConnKey& rhs) const { return memcmp(this, &rhs, sizeof(ConnKey)) > 0; }

    ConnKey& operator=(const ConnKey& rhs);

private:
    void Init(const IPAddr& src, const IPAddr& dst, uint16_t src_port, uint16_t dst_port, uint8_t proto, bool one_way);
};

} // namespace detail

/**
 * Class storing both IPv4 and IPv6 addresses.
 */
class IPAddr {
public:
    /**
     * Address family.
     */
    using Family = IPFamily;

    /**
     * Byte order.
     */
    enum ByteOrder { Host, Network };

    /**
     * Constructs the unspecified IPv6 address (all 128 bits zeroed).
     */
    IPAddr() { memset(in6.s6_addr, 0, sizeof(in6.s6_addr)); }

    /**
     * Constructs an address instance from an IPv4 address.
     *
     * @param in6 The IPv6 address.
     */
    explicit IPAddr(const in4_addr& in4) {
        memcpy(in6.s6_addr, v4_mapped_prefix, sizeof(v4_mapped_prefix));
        memcpy(&in6.s6_addr[12], &in4.s_addr, sizeof(in4.s_addr));
    }

    /**
     * Constructs an address instance from an IPv6 address.
     *
     * @param in6 The IPv6 address.
     */
    explicit IPAddr(const in6_addr& arg_in6) : in6(arg_in6) {}

    /**
     * Constructs an address instance from a string representation.
     *
     * @param s String containing an IP address as either a dotted IPv4
     * address or a hex IPv6 address.
     */
    IPAddr(const std::string& s) { Init(s.data()); }

    /**
     * Constructs an address instance from a string representation.
     *
     * @param s ASCIIZ string containing an IP address as either a
     * dotted IPv4 address or a hex IPv6 address.
     */
    IPAddr(const char* s) { Init(s); }

    /**
     * Constructs an address instance from a string representation.
     *
     * @param s String containing an IP address as either a dotted IPv4
     * address or a hex IPv6 address.
     */
    explicit IPAddr(const String& s);

    /**
     * Constructs an address instance from a raw byte representation.
     *
     * @param family The address family.
     *
     * @param bytes A pointer to the raw byte representation. This must point
     * to 4 bytes if \a family is IPv4, and to 16 bytes if \a family is
     * IPv6.
     *
     * @param order Indicates whether the raw representation pointed to
     * by \a bytes is stored in network or host order.
     */
    IPAddr(Family family, const uint32_t* bytes, ByteOrder order);

    /**
     * Copy constructor.
     */
    IPAddr(const IPAddr& other) : in6(other.in6) {};

    /**
     * Destructor.
     */
    ~IPAddr() = default;

    /**
     * Returns the address' family.
     */
    Family GetFamily() const {
        if ( memcmp(in6.s6_addr, v4_mapped_prefix, 12) == 0 )
            return IPv4;

        return IPv6;
    }

    /**
     * Returns true if the address represents a loopback device.
     */
    bool IsLoopback() const;

    /**
     * Returns true if the address represents a multicast address.
     */
    bool IsMulticast() const {
        if ( GetFamily() == IPv4 )
            return in6.s6_addr[12] == 224;

        return in6.s6_addr[0] == 0xff;
    }

    /**
     * Returns true if the address represents a broadcast address.
     */
    bool IsBroadcast() const {
        if ( GetFamily() == IPv4 )
            return ((in6.s6_addr[12] == 0xff) && (in6.s6_addr[13] == 0xff) && (in6.s6_addr[14] == 0xff) &&
                    (in6.s6_addr[15] == 0xff));

        return false;
    }

    /**
     * Retrieves the raw byte representation of the address.
     *
     * @param bytes The pointer to which \a bytes points will be set to
     * the address of the raw representation in network-byte order.
     * The return value indicates how many 32-bit words are valid starting at
     * that address. The pointer will be valid as long as the address instance
     * exists.
     *
     * @return The number of 32-bit words the raw representation uses. This
     * will be 1 for an IPv4 address and 4 for an IPv6 address.
     */
    int GetBytes(const uint32_t** bytes) const {
        if ( GetFamily() == IPv4 ) {
            *bytes = (uint32_t*)&in6.s6_addr[12];
            return 1;
        }
        else {
            *bytes = (uint32_t*)in6.s6_addr;
            return 4;
        }
    }

    /**
     * Retrieves a copy of the IPv6 raw byte representation of the address.
     * If the internal address is IPv4, then the copied bytes use the
     * IPv4 to IPv6 address mapping to return a full 16 bytes.
     *
     * @param bytes The pointer to a memory location in which the
     * raw bytes of the address are to be copied.
     *
     * @param order The byte-order in which the returned raw bytes are copied.
     * The default is network order.
     */
    void CopyIPv6(uint32_t* bytes, ByteOrder order = Network) const {
        memcpy(bytes, in6.s6_addr, sizeof(in6.s6_addr));

        if ( order == Host ) {
            for ( unsigned int i = 0; i < 4; ++i )
                bytes[i] = ntohl(bytes[i]);
        }
    }

    /**
     * Retrieves a copy of the IPv6 raw byte representation of the address.
     * @see CopyIPv6(uint32_t)
     */
    void CopyIPv6(in6_addr* arg_in6) const { memcpy(arg_in6->s6_addr, in6.s6_addr, sizeof(in6.s6_addr)); }

    /**
     * Retrieves a copy of the IPv4 raw byte representation of the address.
     * The caller should verify the address is of the IPv4 family type
     * beforehand.  @see GetFamily().
     *
     * @param in4 The pointer to a memory location in which the raw bytes
     * of the address are to be copied in network byte-order.
     */
    void CopyIPv4(in4_addr* in4) const { memcpy(&in4->s_addr, &in6.s6_addr[12], sizeof(in4->s_addr)); }

    /**
     * Returns a key that can be used to lookup the IP Address in a hash table.
     */
    std::unique_ptr<detail::HashKey> MakeHashKey() const;

    /**
     * Masks out lower bits of the address.
     *
     * @param top_bits_to_keep The number of bits \a not to mask out,
     * counting from the highest order bit. The value is always
     * interpreted relative to the IPv6 bit width, even if the address
     * is IPv4. That means if compute ``192.168.1.2/16``, you need to
     * pass in 112 (i.e., 96 + 16). The value must be in the range from
     * 0 to 128.
     */
    void Mask(int top_bits_to_keep);

    /**
     * Masks out top bits of the address.
     *
     * @param top_bits_to_chop The number of bits to mask out, counting
     * from the highest order bit.  The value is always interpreted relative
     * to the IPv6 bit width, even if the address is IPv4.  So to mask out
     * the first 16 bits of an IPv4 address, pass in 112 (i.e., 96 + 16).
     * The value must be in the range from 0 to 128.
     */
    void ReverseMask(int top_bits_to_chop);

    /**
     * Assignment operator.
     */
    IPAddr& operator=(const IPAddr& other) {
        // No self-assignment check here because it's correct without it and
        // makes the common case faster.
        in6 = other.in6;
        return *this;
    }

    /**
     * Bitwise OR operator returns the IP address resulting from the bitwise
     * OR operation on the raw bytes of this address with another.
     */
    IPAddr operator|(const IPAddr& other) {
        in6_addr result;
        for ( int i = 0; i < 16; ++i )
            result.s6_addr[i] = this->in6.s6_addr[i] | other.in6.s6_addr[i];

        return IPAddr(result);
    }

    /**
     * Returns a string representation of the address. IPv4 addresses
     * will be returned in dotted representation, IPv6 addresses in
     * compressed hex.
     */
    std::string AsString() const;

    /**
     * Returns a string representation of the address suitable for inclusion
     * in an URI.  For IPv4 addresses, this is the same as AsString(), but
     * IPv6 addresses are encased in square brackets.
     */
    std::string AsURIString() const {
        if ( GetFamily() == IPv4 )
            return AsString();

        return std::string("[") + AsString() + "]";
    }

    /**
     * Returns a host-order, plain hex string representation of the address.
     */
    std::string AsHexString() const;

    /**
     * Returns a string representation of the address. This returns the
     * same as AsString().
     */
    operator std::string() const { return AsString(); }

    /**
     * Returns a reverse pointer name associated with the IP address.
     * For example, 192.168.0.1's reverse pointer is 1.0.168.192.in-addr.arpa.
     */
    std::string PtrName() const;

    /**
     * Comparison operator for IP address.
     */
    friend bool operator==(const IPAddr& addr1, const IPAddr& addr2) {
        return memcmp(&addr1.in6, &addr2.in6, sizeof(in6_addr)) == 0;
    }

    friend bool operator!=(const IPAddr& addr1, const IPAddr& addr2) { return ! (addr1 == addr2); }

    /**
     * Comparison operator IP addresses. This defines a well-defined order for
     * IP addresses. However, the order does not necessarily correspond to
     * their numerical values.
     */
    friend bool operator<(const IPAddr& addr1, const IPAddr& addr2) {
        return memcmp(&addr1.in6, &addr2.in6, sizeof(in6_addr)) < 0;
    }

    friend bool operator<=(const IPAddr& addr1, const IPAddr& addr2) { return addr1 < addr2 || addr1 == addr2; }

    friend bool operator>=(const IPAddr& addr1, const IPAddr& addr2) { return ! (addr1 < addr2); }

    friend bool operator>(const IPAddr& addr1, const IPAddr& addr2) { return ! (addr1 <= addr2); }

    /**
     * Converts the address into the type used internally by the
     * inter-thread communication.
     */
    void ConvertToThreadingValue(threading::Value::addr_t* v) const;

    /**
     * Check if an IP prefix length would be valid against this IP address.
     *
     * @param length the IP prefix length to check
     *
     * @param len_is_v6_relative whether the length is relative to the full
     * IPv6 address length (e.g. since IPv4 addrs are internally stored
     * in v4-to-v6-mapped format, this parameter disambiguates whether
     * a the length is in the usual 32-bit space for IPv4 or the full
     * 128-bit space of IPv6 address.
     *
     * @return whether the prefix length is valid.
     */
    bool CheckPrefixLength(uint8_t length, bool len_is_v6_relative = false) const;

    /**
     * Converts an IPv4 or IPv6 string into a network address structure
     * (IPv6 or v4-to-v6-mapping in network bytes order).
     *
     * @param s the IPv4 or IPv6 string to convert (ASCII, NUL-terminated).
     *
     * @param result buffer that the caller supplies to store the result.
     *
     * @return whether the conversion was successful.
     */
    static bool ConvertString(const char* s, in6_addr* result);

    /**
     * @param s the IPv4 or IPv6 string to convert (ASCII, NUL-terminated).
     *
     * @return whether the string is a valid IP address
     */
    static bool IsValid(const char* s) {
        in6_addr tmp;
        return ConvertString(s, &tmp);
    }

    /**
     * Unspecified IPv4 addr, "0.0.0.0".
     */
    static const IPAddr v4_unspecified;

    /**
     * Unspecified IPv6 addr, "::".
     */
    static const IPAddr v6_unspecified;

private:
    friend class detail::ConnKey;
    friend class IPPrefix;

    /**
     * Initializes an address instance from a string representation.
     *
     * @param s String containing an IP address as either a dotted IPv4
     * address or a hex IPv6 address (ASCII, NUL-terminated).
     */
    void Init(const char* s);

    in6_addr in6; // IPv6 or v4-to-v6-mapped address

    // Top 96 bits of a v4-mapped-addr.
    static constexpr uint8_t v4_mapped_prefix[12] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff};
};

inline IPAddr::IPAddr(Family family, const uint32_t* bytes, ByteOrder order) {
    if ( family == IPv4 ) {
        memcpy(in6.s6_addr, v4_mapped_prefix, sizeof(v4_mapped_prefix));
        memcpy(&in6.s6_addr[12], bytes, sizeof(uint32_t));

        if ( order == Host ) {
            uint32_t* p = (uint32_t*)&in6.s6_addr[12];
            *p = htonl(*p);
        }
    }

    else {
        memcpy(in6.s6_addr, bytes, sizeof(in6.s6_addr));

        if ( order == Host ) {
            for ( unsigned int i = 0; i < 4; ++i ) {
                uint32_t* p = (uint32_t*)&in6.s6_addr[i * 4];
                *p = htonl(*p);
            }
        }
    }
}

inline bool IPAddr::IsLoopback() const {
    if ( GetFamily() == IPv4 )
        return in6.s6_addr[12] == 127;

    else
        return ((in6.s6_addr[0] == 0) && (in6.s6_addr[1] == 0) && (in6.s6_addr[2] == 0) && (in6.s6_addr[3] == 0) &&
                (in6.s6_addr[4] == 0) && (in6.s6_addr[5] == 0) && (in6.s6_addr[6] == 0) && (in6.s6_addr[7] == 0) &&
                (in6.s6_addr[8] == 0) && (in6.s6_addr[9] == 0) && (in6.s6_addr[10] == 0) && (in6.s6_addr[11] == 0) &&
                (in6.s6_addr[12] == 0) && (in6.s6_addr[13] == 0) && (in6.s6_addr[14] == 0) && (in6.s6_addr[15] == 1));
}

inline void IPAddr::ConvertToThreadingValue(threading::Value::addr_t* v) const {
    v->family = GetFamily();

    switch ( v->family ) {
        case IPv4: CopyIPv4(&v->in.in4); return;

        case IPv6: CopyIPv6(&v->in.in6); return;
    }
}

/**
 * Class storing both IPv4 and IPv6 prefixes
 * (i.e., \c 192.168.1.1/16 and \c FD00::/8.
 */
class IPPrefix {
public:
    /**
     * Constructs a prefix 0/0.
     */
    IPPrefix() = default;

    /**
     * Constructs a prefix instance from an IPv4 address and a prefix
     * length.
     *
     * @param in4 The IPv4 address.
     *
     * @param length The prefix length in the range from 0 to 32.
     */
    IPPrefix(const in4_addr& in4, uint8_t length);

    /**
     * Constructs a prefix instance from an IPv6 address and a prefix
     * length.
     *
     * @param in6 The IPv6 address.
     *
     * @param length The prefix length in the range from 0 to 128.
     */
    IPPrefix(const in6_addr& in6, uint8_t length);

    /**
     * Constructs a prefix instance from an IPAddr object and prefix length.
     *
     * @param addr The IP address.
     *
     * @param length The prefix length in the range from 0 to 128
     *
     * @param len_is_v6_relative Whether \a length is relative to the full
     * 128 bits of an IPv6 address.  If false and \a addr is an IPv4
     * address, then \a length is expected to range from 0 to 32.  If true
     * \a length is expected to range from 0 to 128 even if \a addr is IPv4,
     * meaning that the mask is to apply to the IPv4-mapped-IPv6 representation.
     */
    IPPrefix(const IPAddr& addr, uint8_t length, bool len_is_v6_relative = false);

    /**
     * Copy constructor.
     */
    IPPrefix(const IPPrefix& other) : prefix(other.prefix), length(other.length) {}

    /**
     * Destructor.
     */
    ~IPPrefix() = default;

    /**
     * Returns the prefix in the form of an IP address. The address will
     * have all bits not part of the prefixed set to zero.
     */
    const IPAddr& Prefix() const { return prefix; }

    /**
     * Returns the bit length of the prefix, relative to the 32 bits
     * of an IPv4 prefix or relative to the 128 bits of an IPv6 prefix.
     */
    uint8_t Length() const { return prefix.GetFamily() == IPv4 ? length - 96 : length; }

    /**
     * Returns the bit length of the prefix always relative to a full
     * 128 bits of an IPv6 prefix (or IPv4 mapped to IPv6).
     */
    uint8_t LengthIPv6() const { return length; }

    /**
     * Returns true if the given address is part of the prefix.
     *
     * @param addr The address to test.
     */
    bool Contains(const IPAddr& addr) const {
        IPAddr p(addr);
        p.Mask(length);
        return p == prefix;
    }
    /**
     * Assignment operator.
     */
    IPPrefix& operator=(const IPPrefix& other) {
        // No self-assignment check here because it's correct without it and
        // makes the common case faster.
        prefix = other.prefix;
        length = other.length;
        return *this;
    }

    /**
     * Returns a string representation of the prefix. IPv4 addresses
     * will be returned in dotted representation, IPv6 addresses in
     * compressed hex.
     */
    std::string AsString() const;

    operator std::string() const { return AsString(); }

    /**
     * Returns a key that can be used to lookup the IP Prefix in a hash table.
     */
    std::unique_ptr<detail::HashKey> MakeHashKey() const;

    /**
     * Converts the prefix into the type used internally by the
     * inter-thread communication.
     */
    void ConvertToThreadingValue(threading::Value::subnet_t* v) const {
        v->length = length;
        prefix.ConvertToThreadingValue(&v->prefix);
    }

    /**
     * Comparison operator for IP prefix.
     */
    friend bool operator==(const IPPrefix& net1, const IPPrefix& net2) {
        return net1.Prefix() == net2.Prefix() && net1.Length() == net2.Length();
    }

    friend bool operator!=(const IPPrefix& net1, const IPPrefix& net2) { return ! (net1 == net2); }

    /**
     * Comparison operator IP prefixes. This defines a well-defined order for
     * IP prefix. However, the order does not necessarily corresponding to their
     * numerical values.
     */
    friend bool operator<(const IPPrefix& net1, const IPPrefix& net2) {
        if ( net1.Prefix() < net2.Prefix() )
            return true;

        else if ( net1.Prefix() == net2.Prefix() )
            return net1.Length() < net2.Length();

        else
            return false;
    }

    friend bool operator<=(const IPPrefix& net1, const IPPrefix& net2) { return net1 < net2 || net1 == net2; }

    friend bool operator>=(const IPPrefix& net1, const IPPrefix& net2) { return ! (net1 < net2); }

    friend bool operator>(const IPPrefix& net1, const IPPrefix& net2) { return ! (net1 <= net2); }

    /**
     * Converts an IPv4 or IPv6 prefix string into a network address prefix structure.
     *
     * @param s the IPv4 or IPv6 prefix string to convert (ASCII, NUL-terminated).
     *
     * @param result buffer that the caller supplies to store the result.
     *
     * @return whether the conversion was successful.
     */
    static bool ConvertString(const char* s, IPPrefix* result);

    /**
     * @param s the IPv4 or IPv6 prefix string to convert (ASCII, NUL-terminated).
     *
     * @return whether the string is a valid IP address prefix
     */
    static bool IsValid(const char* s) {
        IPPrefix tmp;
        return ConvertString(s, &tmp);
    }

private:
    IPAddr prefix;      // We store it as an address with the non-prefix bits masked out via Mask().
    uint8_t length = 0; // The bit length of the prefix relative to full IPv6 addr.
};

} // namespace zeek