#pragma once

#include "DnsLayerEnums.h"

#include "DnsResource.h"

#include "DnsResourceData.h"

#include "Layer.h"

/// @file

/// @namespace pcpp

/// @brief The main namespace for the PcapPlusPlus lib

namespace pcpp

{

  /// @struct dnshdr

  /// Represents the fixed part of the DNS header, meaning the part that doesn't include the DNS data (queries,

  /// answers, authorities and additional records)

#pragma pack(push, 1)

  struct dnshdr

  {

    /// DNS query identification

    uint16_t transactionID;

#if (BYTE_ORDER == LITTLE_ENDIAN)

    uint16_t

        /// Recursion desired flag

        recursionDesired : 1,

        /// Truncated flag

        truncation : 1,

        /// Authoritative answer flag

        authoritativeAnswer : 1,

        /// Operation Code

        opcode : 4,

        /// Query/Response flag

        queryOrResponse : 1,

        /// Return Code

        responseCode : 4,

        /// Checking disabled flag

        checkingDisabled : 1,

        /// Authenticated data flag

        authenticData : 1,

        /// Zero flag (Reserved)

        zero : 1,

        /// Recursion available flag

        recursionAvailable : 1;

#elif (BYTE_ORDER == BIG_ENDIAN)

    uint16_t

        /// Query/Response flag

        queryOrResponse : 1,

        /// Operation Code

        opcode : 4,

        /// Authoritative answer flag

        authoritativeAnswer : 1,

        /// Truncated flag

        truncation : 1,

        /// Recursion desired flag

        recursionDesired : 1,

        /// Recursion available flag

        recursionAvailable : 1,

        /// Zero flag (Reserved)

        zero : 1,

        /// Authenticated data flag

        authenticData : 1,

        /// Checking disabled flag

        checkingDisabled : 1,

        /// Return Code

        responseCode : 4;

#endif

    /// Number of DNS query records in packet

    uint16_t numberOfQuestions;

    /// Number of DNS answer records in packet

    uint16_t numberOfAnswers;

    /// Number of authority records in packet

    uint16_t numberOfAuthority;

    /// Number of additional records in packet

    uint16_t numberOfAdditional;

  };

#pragma pack(pop)

  static_assert(sizeof(dnshdr) == 12, "dnshdr size is not 12 bytes");

  // forward declarations

  class DnsQuery;

  class IDnsResource;

  class DnsResource;

  class IDnsResourceData;

  /// @class DnsLayer

  /// Represents the DNS protocol layer

  class DnsLayer : public Layer

  {

    friend class IDnsResource;

    friend class DnsQuery;

    friend class DnsResource;

  public:

    /// A constructor that creates the layer from an existing packet raw data

    /// @param[in] data A pointer to the raw data

    /// @param[in] dataLen Size of the data in bytes

    /// @param[in] prevLayer A pointer to the previous layer

    /// @param[in] packet A pointer to the Packet instance where layer will be stored in

    DnsLayer(uint8_t* data, size_t dataLen, Layer* prevLayer, Packet* packet);

    /// A constructor that creates an empty DNS layer: all members of dnshdr are set to 0 and layer will contain no

    /// records

    DnsLayer();

    /// A copy constructor for this layer

    /// @param[in] other The DNS layer to copy from

    DnsLayer(const DnsLayer& other);

    /// An assignment operator for this layer

    /// @param[in] other The DNS layer to assign

    /// @return A reference to the assignee

    DnsLayer& operator=(const DnsLayer& other);

    ~DnsLayer() override;

    /// Get a pointer to the DNS header (as opposed to the DNS data which is the queries, answers, etc. Data can be

    /// retrieved through the other methods of this layer. Notice the return value points directly to the data, so

    /// every change will change the actual packet data

    /// @return A pointer to the @ref dnshdr

    dnshdr* getDnsHeader() const;

    /// Searches for a DNS query by its name field. Notice this method returns only a query which its name equals to

    /// the requested name. If several queries match the requested name, the first one will be returned. If no

    /// queries match the requested name, nullptr will be returned

    /// @param[in] name The name of the query to search

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return The first matching DNS query or nullptr if no queries were found

    DnsQuery* getQuery(const std::string& name, bool exactMatch) const;

    /// @return The first DNS query in the packet or nullptr if packet doesn't contain any queries

    DnsQuery* getFirstQuery() const;

    /// Get the DNS query following a certain query

    /// @param[in] query A pointer to a DNS query that exist in the packet

    /// @return The DNS query following 'query'. If 'query' is nullptr or 'query' is the last query in the packet

    /// nullptr will be returned

    DnsQuery* getNextQuery(DnsQuery* query) const;

    /// @return The number of DNS queries in the packet

    size_t getQueryCount() const;

    /// Add a new DNS query to the layer

    /// @param[in] name The value that shall be set in the name field of the query

    /// @param[in] dnsType The value that shall be set in the DNS type field of the query

    /// @param[in] dnsClass The value that shall be set in the DNS class field of the query

    /// @return A pointer to the newly created DNS query or nullptr if query could not be created (an appropriate

    /// error log message will be printed in this case)

    DnsQuery* addQuery(const std::string& name, DnsType dnsType, DnsClass dnsClass);

    /// Add a new DNS query similar to an already existing DNS query. All query fields will be copied from the

    /// existing query

    /// @param[in] copyQuery The record to create the new record from. copyQuery won't be changed in any way

    /// @return A pointer to the newly created DNS query or nullptr if query could not be created (an appropriate

    /// error log message will be printed in this case)

    DnsQuery* addQuery(DnsQuery* const copyQuery);

    /// Remove an existing query by name. If several queries matches the name, the first match will be removed

    /// @param[in] queryNameToRemove The name of the query to remove

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return True if query was found and successfully removed or false if query was not found or couldn't be

    /// removed

    bool removeQuery(const std::string& queryNameToRemove, bool exactMatch);

    /// Remove an existing query

    /// @param[in] queryToRemove A pointer to the query to remove

    /// @return True if query was found and successfully removed or false if query was not found or couldn't be

    /// removed

    bool removeQuery(DnsQuery* queryToRemove);

    /// Searches for a DNS answer by its name field. Notice this method returns only an answer which its name equals

    /// to the requested name. If several answers match the requested name, the first one will be returned. If no

    /// answers match the requested name, nullptr will be returned

    /// @param[in] name The name of the answer to search

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return The first matching DNS answer or nullptr if no answers were found

    DnsResource* getAnswer(const std::string& name, bool exactMatch) const;

    /// @return The first DNS answer in the packet or nullptr if packet doesn't contain any answers

    DnsResource* getFirstAnswer() const;

    /// Get the DNS answer following a certain answer

    /// @param[in] answer A pointer to a DNS answer that exist in the packet

    /// @return The DNS answer following 'answer'. If 'answer' is nullptr or 'answer' is the last answer in the

    /// packet nullptr will be returned

    DnsResource* getNextAnswer(DnsResource* answer) const;

    /// @return The number of DNS answers in the packet

    size_t getAnswerCount() const;

    /// Add a new DNS answer to the layer

    /// @param[in] name The value that shall be set in the name field of the answer

    /// @param[in] dnsType The value that shall be set in the DNS type field of the answer

    /// @param[in] dnsClass The value that shall be set in the DNS class field of the answer

    /// @param[in] ttl The value that shall be set in the 'time-to-leave' field of the answer

    /// @param[in] data The answer data to be set. The type of the data should match the type of the DNS record

    /// (for example: DNS record of type A should have data of type IPv4DnsResourceData. Please see

    /// DnsResource#setData() for more info on this

    /// @return A pointer to the newly created DNS answer or nullptr if answer could not be created (an appropriate

    /// error log message will be printed in this case)

    DnsResource* addAnswer(const std::string& name, DnsType dnsType, DnsClass dnsClass, uint32_t ttl,

                           IDnsResourceData* data);

    /// Add a new DNS answer similar to an already existing DNS answer. All answer fields will be copied from the

    /// existing answer

    /// @param[in] copyAnswer The record to create the new record from. copyAnswer won't be changed in any way

    /// @return A pointer to the newly created DNS answer or nullptr if query could not be created (an appropriate

    /// error log message will be printed in this case)

    DnsResource* addAnswer(DnsResource* const copyAnswer);

    /// Remove an existing answer by name. If several answers matches the name, the first match will be removed

    /// @param[in] answerNameToRemove The name of the answer to remove

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return True if answer was found and successfully removed or false if answer was not found or couldn't be

    /// removed

    bool removeAnswer(const std::string& answerNameToRemove, bool exactMatch);

    /// Remove an existing answer

    /// @param[in] answerToRemove A pointer to the answer to remove

    /// @return True if answer was found and successfully removed or false if answer was not found or couldn't be

    /// removed

    bool removeAnswer(DnsResource* answerToRemove);

    /// Searches for a DNS authority by its name field. Notice this method returns only an authority which its name

    /// equals to the requested name. If several authorities match the requested name, the first one will be

    /// returned. If no authorities match the requested name, nullptr will be returned

    /// @param[in] name The name of the authority to search

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return The first matching DNS authority or nullptr if no authorities were found

    DnsResource* getAuthority(const std::string& name, bool exactMatch) const;

    /// @return The first DNS authority in the packet or nullptr if packet doesn't contain any authorities

    DnsResource* getFirstAuthority() const;

    /// Get the DNS authority following a certain authority

    /// @param[in] authority A pointer to a DNS authority that exist in the packet

    /// @return The DNS authority following 'authority'. If 'authority' is nullptr or 'authority' is the last

    /// authority in the packet nullptr will be returned

    DnsResource* getNextAuthority(DnsResource* authority) const;

    /// @return The number of DNS authorities in the packet

    size_t getAuthorityCount() const;

    /// Add a new DNS authority to the layer

    /// @param[in] name The value that shall be set in the name field of the authority

    /// @param[in] dnsType The value that shall be set in the DNS type field of the authority

    /// @param[in] dnsClass The value that shall be set in the DNS class field of the authority

    /// @param[in] ttl The value that shall be set in the 'time-to-leave' field of the authority

    /// @param[in] data The authority data to be set. The type of the data should match the type of the DNS record

    /// (for example: DNS record of type A should have data of type IPv4DnsResourceData. Please see

    /// DnsResource#setData() for more info on this

    /// @return A pointer to the newly created DNS authority or nullptr if authority could not be created (an

    /// appropriate error log message will be printed in this case)

    DnsResource* addAuthority(const std::string& name, DnsType dnsType, DnsClass dnsClass, uint32_t ttl,

                              IDnsResourceData* data);

    /// Add a new DNS authority similar to an already existing DNS authority. All authority fields will be copied

    /// from the existing authority

    /// @param[in] copyAuthority The record to create the new record from. copyAuthority won't be changed in any way

    /// @return A pointer to the newly created DNS authority or nullptr if query could not be created (an

    /// appropriate error log message will be printed in this case)

    DnsResource* addAuthority(DnsResource* const copyAuthority);

    /// Remove an existing authority by name. If several authorities matches the name, the first match will be

    /// removed

    /// @param[in] authorityNameToRemove The name of the authority to remove

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return True if authority was found and successfully removed or false if authority was not found or couldn't

    /// be removed

    bool removeAuthority(const std::string& authorityNameToRemove, bool exactMatch);

    /// Remove an existing authority

    /// @param[in] authorityToRemove A pointer to the authority to remove

    /// @return True if authority was found and successfully removed or false if authority was not found or couldn't

    /// be removed

    bool removeAuthority(DnsResource* authorityToRemove);

    /// Searches for a DNS additional record by its name field. Notice this method returns only an additional record

    /// which its name equals to the requested name. If several additional records match the requested name, the

    /// first one will be returned. If no additional records match the requested name, nullptr will be returned

    /// @param[in] name The name of the additional record to search

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return The first matching DNS additional record or nullptr if no additional records were found

    DnsResource* getAdditionalRecord(const std::string& name, bool exactMatch) const;

    /// @return The first DNS additional record in the packet or nullptr if packet doesn't contain any additional

    /// records

    DnsResource* getFirstAdditionalRecord() const;

    /// Get the DNS additional record following a certain additional record

    /// @param[in] additionalRecord A pointer to a DNS additional record that exist in the packet

    /// @return The DNS additional record following 'additionalRecord'. If 'additionalRecord' is nullptr or

    /// 'additionalRecord' is the last additional record in the packet nullptr will be returned

    DnsResource* getNextAdditionalRecord(DnsResource* additionalRecord) const;

    /// @return The number of DNS additional records in the packet

    size_t getAdditionalRecordCount() const;

    /// Add a new DNS additional record to the layer

    /// @param[in] name The value that shall be set in the name field of the additional record

    /// @param[in] dnsType The value that shall be set in the DNS type field of the additional record

    /// @param[in] dnsClass The value that shall be set in the DNS class field of the additional record

    /// @param[in] ttl The value that shall be set in the 'time-to-leave' field of the additional record

    /// @param[in] data The additional record data to be set. The type of the data should match the type of the DNS

    /// record (for example: DNS record of type A should have data of type IPv4DnsResourceData. Please see

    /// DnsResource#setData() for more info on this

    /// @return A pointer to the newly created DNS additional record or nullptr if additional record could not be

    /// created (an appropriate error log message will be printed in this case)

    DnsResource* addAdditionalRecord(const std::string& name, DnsType dnsType, DnsClass dnsClass, uint32_t ttl,

                                     IDnsResourceData* data);

    /// Add a new DNS additional record to the layer that doesn't have DNS class and TTL. Instead these bytes may

    /// contains some arbitrary data. In the future I may add support for these kinds of additional data records.

    /// For now, these bytes are set as raw

    /// @param[in] name The value that shall be set in the name field of the additional record

    /// @param[in] dnsType The value that shall be set in the DNS type field of the additional record

    /// @param[in] customData1 Two bytes of the arbitrary data that will be set in the offset usually used for the

    /// DNS class

    /// @param[in] customData2 Four bytes of the arbitrary data that will be set in the offset usually used for the

    /// TTL

    /// @param[in] data The additional record data to be set. The type of the data should match the type of the DNS

    /// record. (for example: DNS record of type A should have data of type IPv4DnsResourceData. Please see

    /// DnsResource#setData() for more info on this

    /// @return A pointer to the newly created DNS additional record or nullptr if additional record could not be

    /// created (an appropriate error log message will be printed in this case)

    DnsResource* addAdditionalRecord(const std::string& name, DnsType dnsType, uint16_t customData1,

                                     uint32_t customData2, IDnsResourceData* data);

    /// Add a new DNS additional record similar to an already existing DNS additional record. All additional record

    /// fields will be copied from the existing additional record

    /// @param[in] copyAdditionalRecord The record to create the new record from. copyAdditionalRecord won't be

    /// changed in any way

    /// @return A pointer to the newly created DNS additional record or nullptr if query could not be created (an

    /// appropriate error log message will be printed in this case)

    DnsResource* addAdditionalRecord(DnsResource* const copyAdditionalRecord);

    /// Remove an existing additional record by name. If several additional records matches the name, the first

    /// match will be removed

    /// @param[in] additionalRecordNameToRemove The name of the additional record to remove

    /// @param[in] exactMatch Indicate whether to match the whole name or just a part of it

    /// @return True if additional record was found and successfully removed or false if additional record was not

    /// found or couldn't be removed

    bool removeAdditionalRecord(const std::string& additionalRecordNameToRemove, bool exactMatch);

    /// Remove an existing additional record

    /// @param[in] additionalRecordToRemove A pointer to the additional record to remove

    /// @return True if additional record was found and successfully removed or false if additional record was not

    /// found or couldn't be removed

    bool removeAdditionalRecord(DnsResource* additionalRecordToRemove);

    // implement abstract methods

    /// Does nothing for this layer (DnsLayer is always last)

    void parseNextLayer() override

    {}

    /// @return The size of the DNS data in the packet including he DNS header and size of all queries, answers,

    /// authorities and additional records

    size_t getHeaderLen() const override

    {

      return m_DataLen;

    }  // No layer above DNS

    /// Does nothing for this layer

    void computeCalculateFields() override

    {}

    std::string toString() const override;

    OsiModelLayer getOsiModelLayer() const override

    {

      return OsiModelApplicationLayer;

    }

    /// A static method that checks whether the port is considered as DNS

    /// @param[in] port The port number to be checked

    /// @return True if the port is associated with the DNS protocol

    static inline bool isDnsPort(uint16_t port);

    /// A static method that validates the input data

    /// @param[in] data The pointer to the beginning of a byte stream of a DNS packet

    /// @param[in] dataLen The length of the byte stream

    /// @param[in] dnsOverTcp Should be set to "true" if this is DNS is over TCP, otherwise set to "false"

    /// (which is also the default value)

    /// @return True if the data is valid and can represent a DNS packet

    static inline bool isDataValid(const uint8_t* data, size_t dataLen, bool dnsOverTcp = false);

  protected:

    DnsLayer(uint8_t* data, size_t dataLen, Layer* prevLayer, Packet* packet, size_t offsetAdjustment);

    explicit DnsLayer(size_t offsetAdjustment);

  private:

    IDnsResource* m_ResourceList;

    DnsQuery* m_FirstQuery;

    DnsResource* m_FirstAnswer;

    DnsResource* m_FirstAuthority;

    DnsResource* m_FirstAdditional;

    uint16_t m_OffsetAdjustment;

    size_t getBasicHeaderSize();

    void init(size_t offsetAdjustment, bool callParseResource);

    void initNewLayer(size_t offsetAdjustment);

    IDnsResource* getFirstResource(DnsResourceType resType) const;

    void setFirstResource(DnsResourceType resType, IDnsResource* resource);

    using Layer::extendLayer;

    bool extendLayer(int offsetInLayer, size_t numOfBytesToExtend, IDnsResource* resource);

    using Layer::shortenLayer;

    bool shortenLayer(int offsetInLayer, size_t numOfBytesToShorten, IDnsResource* resource);

    IDnsResource* getResourceByName(IDnsResource* startFrom, size_t resourceCount, const std::string& name,

                                    bool exactMatch) const;

    void parseResources();

    DnsResource* addResource(DnsResourceType resType, const std::string& name, DnsType dnsType, DnsClass dnsClass,

                             uint32_t ttl, IDnsResourceData* data);

    bool removeResource(IDnsResource* resourceToRemove);

  };

  /// @class DnsOverTcpLayer

  /// Represents the DNS over TCP layer.

  /// DNS over TCP is described here: https://tools.ietf.org/html/rfc7766 .

  /// It is very similar to DNS over UDP, except for one field: TCP message length which is added in the beginning of

  /// the message before the other DNS data properties. The rest of the data is similar.

  /// Note: DNS over TCP can spread over more than one packet, but this implementation doesn't support this use-case

  /// and assumes the whole message fits in a single packet.

  class DnsOverTcpLayer : public DnsLayer

  {

  public:

    /// A constructor that creates the layer from an existing packet raw data

    /// @param[in] data A pointer to the raw data

    /// @param[in] dataLen Size of the data in bytes

    /// @param[in] prevLayer A pointer to the previous layer

    /// @param[in] packet A pointer to the Packet instance where layer will be stored in

    DnsOverTcpLayer(uint8_t* data, size_t dataLen, Layer* prevLayer, Packet* packet)

        : DnsLayer(data, dataLen, prevLayer, packet, sizeof(uint16_t))

    {}

    /// A constructor that creates an empty DNS layer: all members of dnshdr are set to 0 and layer will contain no

    /// records

    DnsOverTcpLayer() : DnsLayer(sizeof(uint16_t))

    {}

    /// A copy constructor for this layer

    /// @param[in] other The DNS over TCP layer to copy from

    DnsOverTcpLayer(const DnsOverTcpLayer& other) : DnsLayer(other)

    {}

    /// @return The value of the TCP message length as described in https://tools.ietf.org/html/rfc7766#section-8

    uint16_t getTcpMessageLength();

    /// Set the TCP message length value as described in https://tools.ietf.org/html/rfc7766#section-8

    /// @param[in] value The value to set

    void setTcpMessageLength(uint16_t value);

    // overridden methods

    /// Calculate the TCP message length field

    void computeCalculateFields() override;

  };

  // implementation of inline methods

  bool DnsLayer::isDnsPort(uint16_t port)

  {

    switch (port)

    {

    case 53:

    case 5353:

    case 5355:

      return true;

    default:

      return false;

    }

  }

  bool DnsLayer::isDataValid(const uint8_t* data, size_t dataLen, bool dnsOverTcp)

  {

    size_t minSize = sizeof(dnshdr) + (dnsOverTcp ? sizeof(uint16_t) : 0);

    return data && dataLen >= minSize;

  }

}  // namespace pcpp