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

#pragma once

#include "zeek/zeek-config.h"

// clang-format off
#include <sys/types.h> // for u_char
#include <netinet/in.h>
#include <netinet/ip.h>
// clang-format on

#ifdef HAVE_NETINET_IP6_H
#include <netinet/ip6.h>
#else
#include "net_util.h" // for struct ip6_hdr
#endif

#include <vector>

#include "zeek/IntrusivePtr.h"

namespace zeek
	{

class IPAddr;
class RecordVal;
class VectorVal;
using RecordValPtr = IntrusivePtr<RecordVal>;
using VectorValPtr = IntrusivePtr<VectorVal>;

namespace detail
	{
class FragReassembler;
	}

#ifndef IPPROTO_MOBILITY
#define IPPROTO_MOBILITY 135
#endif

struct ip6_mobility
	{
	uint8_t ip6mob_payload;
	uint8_t ip6mob_len;
	uint8_t ip6mob_type;
	uint8_t ip6mob_rsv;
	uint16_t ip6mob_chksum;
	};

/**
 * Base class for IPv6 header/extensions.
 */
class IPv6_Hdr
	{
public:
	/**
	 * Construct an IPv6 header or extension header from assigned type number.
	 */
	IPv6_Hdr(uint8_t t, const u_char* d) : type(t), data(d) { }

	/**
	 * Replace the value of the next protocol field.
	 */
	void ChangeNext(uint8_t next_type)
		{
		switch ( type )
			{
			case IPPROTO_IPV6:
				((ip6_hdr*)data)->ip6_nxt = next_type;
				break;
			case IPPROTO_HOPOPTS:
			case IPPROTO_DSTOPTS:
			case IPPROTO_ROUTING:
			case IPPROTO_FRAGMENT:
			case IPPROTO_AH:
			case IPPROTO_MOBILITY:
				((ip6_ext*)data)->ip6e_nxt = next_type;
				break;
			case IPPROTO_ESP:
			default:
				break;
			}
		}

	~IPv6_Hdr() { }

	/**
	 * Returns the assigned IPv6 extension header type number of the header
	 * that immediately follows this one.
	 */
	uint8_t NextHdr() const
		{
		switch ( type )
			{
			case IPPROTO_IPV6:
				return ((ip6_hdr*)data)->ip6_nxt;
			case IPPROTO_HOPOPTS:
			case IPPROTO_DSTOPTS:
			case IPPROTO_ROUTING:
			case IPPROTO_FRAGMENT:
			case IPPROTO_AH:
			case IPPROTO_MOBILITY:
				return ((ip6_ext*)data)->ip6e_nxt;
			case IPPROTO_ESP:
			default:
				return IPPROTO_NONE;
			}
		}

	/**
	 * Returns the length of the header in bytes.
	 */
	uint16_t Length() const
		{
		switch ( type )
			{
			case IPPROTO_IPV6:
				return 40;
			case IPPROTO_HOPOPTS:
			case IPPROTO_DSTOPTS:
			case IPPROTO_ROUTING:
			case IPPROTO_MOBILITY:
				return 8 + 8 * ((ip6_ext*)data)->ip6e_len;
			case IPPROTO_FRAGMENT:
				return 8;
			case IPPROTO_AH:
				return 8 + 4 * ((ip6_ext*)data)->ip6e_len;
			case IPPROTO_ESP:
				return 8; // encrypted payload begins after 8 bytes
			default:
				return 0;
			}
		}

	/**
	 * Returns the RFC 1700 et seq. IANA assigned number for the header.
	 */
	uint8_t Type() const { return type; }

	/**
	 * Returns pointer to the start of where header structure resides in memory.
	 */
	const u_char* Data() const { return data; }

	/**
	 * Returns the script-layer record representation of the header.
	 */
	RecordValPtr ToVal(VectorValPtr chain) const;
	RecordValPtr ToVal() const;

protected:
	uint8_t type;
	const u_char* data;

private:
	bool IsOptionTruncated(uint16_t off) const;
	};

class IPv6_Hdr_Chain
	{
public:
	/**
	 * Initializes the header chain from an IPv6 header structure.
	 */
	IPv6_Hdr_Chain(const struct ip6_hdr* ip6, uint64_t len) { Init(ip6, len, false); }

	~IPv6_Hdr_Chain();

	/**
	 * @return a copy of the header chain, but with pointers to individual
	 * IPv6 headers now pointing within \a new_hdr.
	 */
	IPv6_Hdr_Chain* Copy(const struct ip6_hdr* new_hdr) const;

	/**
	 * Returns the number of headers in the chain.
	 */
	size_t Size() const { return chain.size(); }

	/**
	 * Returns the sum of the length of all headers in the chain in bytes.
	 */
	uint16_t TotalLength() const { return length; }

	/**
	 * Accesses the header at the given location in the chain.
	 */
	const IPv6_Hdr* operator[](const size_t i) const { return chain[i]; }

	/**
	 * Returns whether the header chain indicates a fragmented packet.
	 */
	bool IsFragment() const;

	/**
	 * Returns pointer to fragment header structure if the chain contains one.
	 */
	const struct ip6_frag* GetFragHdr() const
		{
		return IsFragment() ? (const struct ip6_frag*)chain[chain.size() - 1]->Data() : nullptr;
		}

	/**
	 * If the header chain is a fragment, returns the offset in number of bytes
	 * relative to the start of the Fragmentable Part of the original packet.
	 */
	uint16_t FragOffset() const
		{
		return IsFragment() ? (ntohs(GetFragHdr()->ip6f_offlg) & 0xfff8) : 0;
		}

	/**
	 * If the header chain is a fragment, returns the identification field.
	 */
	uint32_t ID() const { return IsFragment() ? ntohl(GetFragHdr()->ip6f_ident) : 0; }

	/**
	 * If the header chain is a fragment, returns the M (more fragments) flag.
	 */
	int MF() const { return IsFragment() ? (ntohs(GetFragHdr()->ip6f_offlg) & 0x0001) != 0 : 0; }

	/**
	 * If the chain contains a Destination Options header with a Home Address
	 * option as defined by Mobile IPv6 (RFC 6275), then return it, else
	 * return the source address in the main IPv6 header.
	 */
	IPAddr SrcAddr() const;

	/**
	 * If the chain contains a Routing header with non-zero segments left,
	 * then return the last address of the first such header, else return
	 * the destination address of the main IPv6 header.
	 */
	IPAddr DstAddr() const;

	/**
	 * Returns a vector of ip6_ext_hdr RecordVals that includes script-layer
	 * representation of all extension headers in the chain.
	 */
	VectorValPtr ToVal() const;

protected:
	// for access to protected ctor that changes next header values that
	// point to a fragment
	friend class detail::FragReassembler;

	IPv6_Hdr_Chain() = default;

	/**
	 * Initializes the header chain from an IPv6 header structure, and replaces
	 * the first next protocol pointer field that points to a fragment header.
	 */
	IPv6_Hdr_Chain(const struct ip6_hdr* ip6, uint16_t next, uint64_t len)
		{
		Init(ip6, len, true, next);
		}

	/**
	 * Initializes the header chain from an IPv6 header structure of a given
	 * length, possibly setting the first next protocol pointer field that
	 * points to a fragment header.
	 */
	void Init(const struct ip6_hdr* ip6, uint64_t total_len, bool set_next, uint16_t next = 0);

	/**
	 * Process a routing header and allocate/remember the final destination
	 * address if it has segments left and is a valid routing header.
	 */
	void ProcessRoutingHeader(const struct ip6_rthdr* r, uint16_t len);

	/**
	 * Inspect a Destination Option header's options for things we need to
	 * remember, such as the Home Address option from Mobile IPv6.
	 */
	void ProcessDstOpts(const struct ip6_dest* d, uint16_t len);

	std::vector<IPv6_Hdr*> chain;

	/**
	 * The summation of all header lengths in the chain in bytes.
	 */
	uint16_t length = 0;

	/**
	 * Home Address of the packet's source as defined by Mobile IPv6 (RFC 6275).
	 */
	IPAddr* homeAddr = nullptr;

	/**
	 * The final destination address in chain's first Routing header that has
	 * non-zero segments left.
	 */
	IPAddr* finalDst = nullptr;
	};

/**
 * A class that wraps either an IPv4 or IPv6 packet and abstracts methods
 * for inquiring about common features between the two.
 */
class IP_Hdr
	{
public:
	/**
	 * Construct the header wrapper from an IPv4 packet.  Caller must have
	 * already checked that the header is not truncated.
	 * @param arg_ip4 pointer to memory containing an IPv4 packet.
	 * @param arg_del whether to take ownership of \a arg_ip4 pointer's memory.
	 * @param reassembled whether this header is for a reassembled packet.
	 */
	IP_Hdr(const struct ip* arg_ip4, bool arg_del, bool reassembled = false)
		: ip4(arg_ip4), del(arg_del), reassembled(reassembled)
		{
		}

	/**
	 * Construct the header wrapper from an IPv6 packet.  Caller must have
	 * already checked that the static IPv6 header is not truncated.  If
	 * the packet contains extension headers and they are truncated, that can
	 * be checked afterwards by comparing \a len with \a TotalLen.  E.g.
	 * The IP packet analyzer does this to skip truncated packets.
	 * @param arg_ip6 pointer to memory containing an IPv6 packet.
	 * @param arg_del whether to take ownership of \a arg_ip6 pointer's memory.
	 * @param len the packet's length in bytes.
	 * @param c an already-constructed header chain to take ownership of.
	 * @param reassembled whether this header is for a reassembled packet.
	 */
	IP_Hdr(const struct ip6_hdr* arg_ip6, bool arg_del, uint64_t len,
	       const IPv6_Hdr_Chain* c = nullptr, bool reassembled = false)
		: ip6(arg_ip6), ip6_hdrs(c ? c : new IPv6_Hdr_Chain(ip6, len)), del(arg_del),
		  reassembled(reassembled)
		{
		}

	/**
	 * Copy a header.  The internal buffer which contains the header data
	 * must not be truncated.  Also note that if that buffer points to a full
	 * packet payload, only the IP header portion is copied.
	 */
	IP_Hdr* Copy() const;

	/**
	 * Destructor.
	 */
	~IP_Hdr()
		{
		delete ip6_hdrs;

		if ( del )
			{
			delete[](struct ip*) ip4;
			delete[](struct ip6_hdr*) ip6;
			}
		}

	/**
	 * If an IPv4 packet is wrapped, return a pointer to it, else null.
	 */
	const struct ip* IP4_Hdr() const { return ip4; }

	/**
	 * If an IPv6 packet is wrapped, return a pointer to it, else null.
	 */
	const struct ip6_hdr* IP6_Hdr() const { return ip6; }

	/**
	 * Returns the source address held in the IP header.
	 */
	IPAddr IPHeaderSrcAddr() const;

	/**
	 * Returns the destination address held in the IP header.
	 */
	IPAddr IPHeaderDstAddr() const;

	/**
	 * For IPv4 or IPv6 headers that don't contain a Home Address option
	 * (Mobile IPv6, RFC 6275), return source address held in the IP header.
	 * For IPv6 headers that contain a Home Address option, return that address.
	 */
	IPAddr SrcAddr() const;

	/**
	 * For IPv4 or IPv6 headers that don't contain a Routing header with
	 * non-zero segments left, return destination address held in the IP header.
	 * For IPv6 headers with a Routing header that has non-zero segments left,
	 * return the last address in the first such Routing header.
	 */
	IPAddr DstAddr() const;

	/**
	 * Returns a pointer to the payload of the IP packet, usually an
	 * upper-layer protocol.
	 */
	const u_char* Payload() const
		{
		if ( ip4 )
			return ((const u_char*)ip4) + ip4->ip_hl * 4;

		return ((const u_char*)ip6) + ip6_hdrs->TotalLength();
		}

	/**
	 * Returns a pointer to the mobility header of the IP packet, if present,
	 * else a null pointer.
	 */
	const ip6_mobility* MobilityHeader() const
		{
		if ( ip4 )
			return nullptr;
		else if ( (*ip6_hdrs)[ip6_hdrs->Size() - 1]->Type() != IPPROTO_MOBILITY )
			return nullptr;
		else
			return (const ip6_mobility*)(*ip6_hdrs)[ip6_hdrs->Size() - 1]->Data();
		}

	/**
	 * Returns the length of the IP packet's payload (length of packet minus
	 * header length or, for IPv6, also minus length of all extension headers).
	 *
	 * Also returns 0 if the IPv4 length field is set to zero - which is, e.g.,
	 * the case when TCP segment offloading is enabled.
	 */
	uint16_t PayloadLen() const
		{
		if ( ip4 )
			{
			// prevent overflow in case of segment offloading/zeroed header length.
			auto total_len = ntohs(ip4->ip_len);
			return total_len ? total_len - ip4->ip_hl * 4 : 0;
			}

		return ntohs(ip6->ip6_plen) + 40 - ip6_hdrs->TotalLength();
		}

	/**
	 * Returns the length of the IP packet (length of headers and payload).
	 */
	uint32_t TotalLen() const
		{
		if ( ip4 )
			return ntohs(ip4->ip_len);

		return ntohs(ip6->ip6_plen) + 40;
		}

	/**
	 * Returns length of IP packet header (includes extension headers for IPv6).
	 */
	uint16_t HdrLen() const { return ip4 ? ip4->ip_hl * 4 : ip6_hdrs->TotalLength(); }

	/**
	 * For IPv6 header chains, returns the type of the last header in the chain.
	 */
	uint8_t LastHeader() const
		{
		if ( ip4 )
			return IPPROTO_RAW;

		size_t i = ip6_hdrs->Size();
		if ( i > 0 )
			return (*ip6_hdrs)[i - 1]->Type();

		return IPPROTO_NONE;
		}

	/**
	 * Returns the protocol type of the IP packet's payload, usually an
	 * upper-layer protocol.  For IPv6, this returns the last (extension)
	 * header's Next Header value.
	 */
	unsigned char NextProto() const
		{
		if ( ip4 )
			return ip4->ip_p;

		size_t i = ip6_hdrs->Size();
		if ( i > 0 )
			return (*ip6_hdrs)[i - 1]->NextHdr();

		return IPPROTO_NONE;
		}

	/**
	 * Returns the IPv4 Time to Live or IPv6 Hop Limit field.
	 */
	unsigned char TTL() const { return ip4 ? ip4->ip_ttl : ip6->ip6_hlim; }

	/**
	 * Returns whether the IP header indicates this packet is a fragment.
	 */
	bool IsFragment() const
		{
		return ip4 ? (ntohs(ip4->ip_off) & 0x3fff) != 0 : ip6_hdrs->IsFragment();
		}

	/**
	 * Returns the fragment packet's offset in relation to the original
	 * packet in bytes.
	 */
	uint16_t FragOffset() const
		{
		return ip4 ? (ntohs(ip4->ip_off) & 0x1fff) * 8 : ip6_hdrs->FragOffset();
		}

	/**
	 * Returns the fragment packet's identification field.
	 */
	uint32_t ID() const { return ip4 ? ntohs(ip4->ip_id) : ip6_hdrs->ID(); }

	/**
	 * Returns whether a fragment packet's "More Fragments" field is set.
	 */
	int MF() const { return ip4 ? (ntohs(ip4->ip_off) & 0x2000) != 0 : ip6_hdrs->MF(); }

	/**
	 * Returns whether a fragment packet's "Don't Fragment" field is set.
	 * Note that IPv6 has no such field.
	 */
	int DF() const { return ip4 ? ((ntohs(ip4->ip_off) & 0x4000) != 0) : 0; }

	/**
	 * Returns value of an IPv6 header's flow label field or 0 if it's IPv4.
	 */
	uint32_t FlowLabel() const { return ip4 ? 0 : (ntohl(ip6->ip6_flow) & 0x000fffff); }

	/**
	 * Returns number of IP headers in packet (includes IPv6 extension headers).
	 */
	size_t NumHeaders() const { return ip4 ? 1 : ip6_hdrs->Size(); }

	/**
	 * Returns an ip_hdr or ip6_hdr_chain RecordVal.
	 */
	RecordValPtr ToIPHdrVal() const;

	/**
	 * Returns a pkt_hdr RecordVal, which includes not only the IP header, but
	 * also upper-layer (tcp/udp/icmp) headers.
	 */
	RecordValPtr ToPktHdrVal() const;

	/**
	 * Same as above, but simply add our values into the record at the
	 * specified starting index.
	 */
	RecordValPtr ToPktHdrVal(RecordValPtr pkt_hdr, int sindex) const;

	bool Reassembled() const { return reassembled; }

private:
	const struct ip* ip4 = nullptr;
	const struct ip6_hdr* ip6 = nullptr;
	const IPv6_Hdr_Chain* ip6_hdrs = nullptr;
	bool del = false;
	bool reassembled = false;
	};

	} // namespace zeek
