#define LOG_MODULE PacketLogModuleDnsLayer

#include "DnsResource.h"

#include "Logger.h"

#include <sstream>

#include <string.h>

#include "EndianPortable.h"

namespace pcpp

{

IDnsResource::IDnsResource(DnsLayer* dnsLayer, size_t offsetInLayer)

  : m_DnsLayer(dnsLayer), m_OffsetInLayer(offsetInLayer), m_NextResource(NULL)

{

  char decodedName[256];

  m_NameLength = decodeName((const char*)getRawData(), decodedName);

  if (m_NameLength > 0)

    m_DecodedName = decodedName;

}

IDnsResource::IDnsResource(uint8_t* emptyRawData)

  : m_DnsLayer(NULL), m_OffsetInLayer(0), m_NextResource(NULL), m_DecodedName(""), m_NameLength(0), m_ExternalRawData(emptyRawData)

{

}

uint8_t* IDnsResource::getRawData() const

{

  if (m_DnsLayer == NULL)

    return m_ExternalRawData;

  return m_DnsLayer->m_Data + m_OffsetInLayer;

}

static size_t cleanup(char* resultPtr, char* result, size_t encodedNameLength)

{

  // remove the last "."

  if (resultPtr > result)

  {

    result[resultPtr - result - 1] = 0;

  }

  if (resultPtr - result < 256)

  {

    // add the last '\0' to encodedNameLength

    resultPtr[0] = 0;

    encodedNameLength++;

  }

  return encodedNameLength;

}

size_t IDnsResource::decodeName(const char* encodedName, char* result, int iteration)

{

  size_t encodedNameLength = 0;

  size_t decodedNameLength = 0;

  char* resultPtr = result;

  resultPtr[0] = 0;

  size_t curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;

  if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)

    return encodedNameLength;

  if (iteration > 20)

  {

    return encodedNameLength;

  }

  uint8_t wordLength = encodedName[0];

  // A string to parse

  while (wordLength != 0)

  {

    // A pointer to another place in the packet

    if ((wordLength & 0xc0) == 0xc0)

    {

      if (curOffsetInLayer + 2 > m_DnsLayer->m_DataLen || encodedNameLength > 255)

        return cleanup(resultPtr, result, encodedNameLength);

      uint16_t offsetInLayer = (wordLength & 0x3f)*256 + (0xFF & encodedName[1]) + m_DnsLayer->m_OffsetAdjustment;

      if (offsetInLayer < sizeof(dnshdr) || offsetInLayer >= m_DnsLayer->m_DataLen)

      {

        PCPP_LOG_ERROR("DNS parsing error: name pointer is illegal");

        return 0;

      }

      char tempResult[256];

      memset(tempResult, 0, 256);

      int i = 0;

      decodeName((const char*)(m_DnsLayer->m_Data + offsetInLayer), tempResult, iteration+1);

      while (tempResult[i] != 0 && decodedNameLength < 255)

      {

        resultPtr[0] = tempResult[i++];

        resultPtr++;

        decodedNameLength++;

      }

      resultPtr[0] = 0;

      // in this case the length of the pointer is: 1 byte for 0xc0 + 1 byte for the offset itself

      return encodedNameLength + sizeof(uint16_t);

    }

    else

    {

      // return if next word would be outside of the DNS layer or overflow the buffer behind resultPtr

      if (curOffsetInLayer + wordLength + 1 > m_DnsLayer->m_DataLen || encodedNameLength + wordLength > 255)

      {

        // add the last '\0' to the decoded string

        if (encodedNameLength == 256)

        {

          resultPtr--;

          // cppcheck-suppress unreadVariable

          decodedNameLength--;

        }

        else

        {

          encodedNameLength++;

        }

        resultPtr[0] = 0;

        return encodedNameLength;

      }

      memcpy(resultPtr, encodedName+1, wordLength);

      resultPtr += wordLength;

      resultPtr[0] = '.';

      resultPtr++;

      decodedNameLength += wordLength + 1;

      encodedName += wordLength + 1;

      encodedNameLength += wordLength + 1;

      curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;

      if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)

      {

        // add the last '\0' to the decoded string

        if (encodedNameLength == 256)

        {

          // cppcheck-suppress unreadVariable

          decodedNameLength--;

          resultPtr--;

        }

        else

        {

          encodedNameLength++;

        }

        resultPtr[0] = 0;

        return encodedNameLength;

      }

      wordLength = encodedName[0];

    }

  }

  return cleanup(resultPtr, result, encodedNameLength);

}

void IDnsResource::encodeName(const std::string& decodedName, char* result, size_t& resultLen)

{

  resultLen = 0;

  std::stringstream strstream(decodedName);

  std::string word;

  while (getline(strstream, word, '.'))

  {

    // pointer to a different hostname in the packet

    if (word[0] == '#')

    {

      // convert the number from string to int

      std::stringstream stream(word.substr(1));

      int pointerInPacket = 0;

      stream >> pointerInPacket;

      // verify it's indeed a number and that is in the range of [0-255]

      if (stream.fail() || pointerInPacket < 0 || pointerInPacket > 0xff)

      {

        PCPP_LOG_ERROR("Error encoding the string '" << decodedName << "'");

        return;

      }

      // set the pointer to the encoded string result

      result[0] = (uint8_t)0xc0;

      result[1] = (uint8_t)pointerInPacket;

      result += 2;

      resultLen += 2;

      return; // pointer always comes last

    }

    result[0] = word.length();

    result++;

    memcpy(result, word.c_str(), word.length());

    result += word.length();

    resultLen += word.length() + 1;

  }

  result[0] = 0;

  resultLen++;

}

DnsType IDnsResource::getDnsType() const

{

  uint16_t dnsType = *(uint16_t*)(getRawData() + m_NameLength);

  return (DnsType)be16toh(dnsType);

}

void IDnsResource::setDnsType(DnsType newType)

{

  uint16_t newTypeAsInt = htobe16((uint16_t)newType);

  memcpy(getRawData() + m_NameLength, &newTypeAsInt, sizeof(uint16_t));

}

DnsClass IDnsResource::getDnsClass() const

{

  uint16_t dnsClass = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));

  return (DnsClass)be16toh(dnsClass);

}

void IDnsResource::setDnsClass(DnsClass newClass)

{

  uint16_t newClassAsInt = htobe16((uint16_t)newClass);

  memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &newClassAsInt, sizeof(uint16_t));

}

bool IDnsResource::setName(const std::string& newName)

{

  char encodedName[256];

  size_t encodedNameLen = 0;

  encodeName(newName, encodedName, encodedNameLen);

  if (m_DnsLayer != NULL)

  {

    if (encodedNameLen > m_NameLength)

    {

      if (!m_DnsLayer->extendLayer(m_OffsetInLayer, encodedNameLen-m_NameLength, this))

      {

        PCPP_LOG_ERROR("Couldn't set name for DNS query, unable to extend layer");

        return false;

      }

    }

    else if (encodedNameLen < m_NameLength)

    {

      if (!m_DnsLayer->shortenLayer(m_OffsetInLayer, m_NameLength-encodedNameLen, this))

      {

        PCPP_LOG_ERROR("Couldn't set name for DNS query, unable to shorten layer");

        return false;

      }

    }

  }

  else

  {

    size_t size = getSize();

    char* tempData = new char[size];

    memcpy(tempData, m_ExternalRawData, size);

    memcpy(m_ExternalRawData + encodedNameLen, tempData, size);

    delete[] tempData;

  }

  memcpy(getRawData(), encodedName, encodedNameLen);

  m_NameLength = encodedNameLen;

  m_DecodedName = newName;

  return true;

}

void IDnsResource::setDnsLayer(DnsLayer* dnsLayer, size_t offsetInLayer)

{

  memcpy(dnsLayer->m_Data + offsetInLayer, m_ExternalRawData, getSize());

  m_DnsLayer = dnsLayer;

  m_OffsetInLayer = offsetInLayer;

  m_ExternalRawData = NULL;

}

uint32_t DnsResource::getTTL() const

{

  uint32_t ttl = *(uint32_t*)(getRawData() + m_NameLength + 2*sizeof(uint16_t));

  return be32toh(ttl);

}

void DnsResource::setTTL(uint32_t newTTL)

{

  newTTL = htobe32(newTTL);

  memcpy(getRawData() + m_NameLength + 2*sizeof(uint16_t), &newTTL, sizeof(uint32_t));

}

size_t DnsResource::getDataLength() const

{

  size_t sizeToRead = m_NameLength + 2*sizeof(uint16_t) + sizeof(uint32_t);

  // Heap buffer overflow may occur here, check boundary of m_DnsLayer->m_Data first

  // Due to dataLength which is uint16_t, here m_DnsLayer->m_Data must have at least 2 bytes to read

  if (m_DnsLayer && m_OffsetInLayer + sizeToRead >= m_DnsLayer->m_DataLen - 1)

  {

    return 0;

  }

  uint16_t dataLength = *(uint16_t*)(getRawData() + sizeToRead);

  return be16toh(dataLength);

}

DnsResourceDataPtr DnsResource::getData() const

{

  uint8_t* resourceRawData = getRawData() + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t);

  size_t dataLength = getDataLength();

  switch (getDnsType())

  {

  case DNS_TYPE_A:

  {

    return DnsResourceDataPtr(new IPv4DnsResourceData(resourceRawData, dataLength));

  }

  case DNS_TYPE_AAAA:

  {

    return DnsResourceDataPtr(new IPv6DnsResourceData(resourceRawData, dataLength));

  }

  case DNS_TYPE_NS:

  case DNS_TYPE_CNAME:

  case DNS_TYPE_DNAM:

  case DNS_TYPE_PTR:

  {

    return DnsResourceDataPtr(new StringDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));

  }

  case DNS_TYPE_MX:

  {

    return DnsResourceDataPtr(new MxDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));

  }

  default:

  {

    return DnsResourceDataPtr(new GenericDnsResourceData(resourceRawData, dataLength));

  }

  }

}

size_t DnsResource::getDataOffset() const

{

  return (size_t)(m_OffsetInLayer + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t));

}

bool DnsResource::setData(IDnsResourceData* data)

{

  // convert data to byte array according to the DNS type

  size_t dataLength = 0;

  uint8_t dataAsByteArr[256];

  if (data == NULL)

  {

    PCPP_LOG_ERROR("Given data is NULL");

    return false;

  }

  switch (getDnsType())

  {

  case DNS_TYPE_A:

  {

    if (!data->isTypeOf<IPv4DnsResourceData>())

    {

      PCPP_LOG_ERROR("DNS record is of type A but given data isn't of type IPv4DnsResourceData");

      return false;

    }

    break;

  }

  case DNS_TYPE_AAAA:

  {

    if (!data->isTypeOf<IPv6DnsResourceData>())

    {

      PCPP_LOG_ERROR("DNS record is of type AAAA but given data isn't of type IPv6DnsResourceData");

      return false;

    }

    break;

  }

  case DNS_TYPE_NS:

  case DNS_TYPE_CNAME:

  case DNS_TYPE_DNAM:

  case DNS_TYPE_PTR:

  {

    if (!data->isTypeOf<StringDnsResourceData>())

    {

      PCPP_LOG_ERROR("DNS record is of type NS, CNAME, DNAM or PTR but given data isn't of type StringDnsResourceData");

      return false;

    }

    break;

  }

  case DNS_TYPE_MX:

  {

    if (!data->isTypeOf<MxDnsResourceData>())

    {

      PCPP_LOG_ERROR("DNS record is of type MX but given data isn't of type MxDnsResourceData");

      return false;

    }

    break;

  }

  default:

  {

    // do nothing

  }

  }

  // convert the IDnsResourceData to byte array

  if (!data->toByteArr(dataAsByteArr, dataLength, this))

  {

    PCPP_LOG_ERROR("Cannot convert DNS resource data to byte array, data is probably invalid");

    return false;

  }

  size_t dataLengthOffset = m_NameLength + (2*sizeof(uint16_t)) + sizeof(uint32_t);

  size_t dataOffset = dataLengthOffset + sizeof(uint16_t);

  if (m_DnsLayer != NULL)

  {

    size_t curLength = getDataLength();

    if (dataLength > curLength)

    {

      if (!m_DnsLayer->extendLayer(m_OffsetInLayer + dataOffset, dataLength-curLength, this))

      {

        PCPP_LOG_ERROR("Couldn't set data for DNS query, unable to extend layer");

        return false;

      }

    }

    else if (dataLength < curLength)

    {

      if (!m_DnsLayer->shortenLayer(m_OffsetInLayer + dataOffset, curLength-dataLength, this))

      {

        PCPP_LOG_ERROR("Couldn't set data for DNS query, unable to shorten layer");

        return false;

      }

    }

  }

  // write data to resource

  memcpy(getRawData() + dataOffset, dataAsByteArr, dataLength);

  //update data length in resource

  dataLength = htobe16((uint16_t)dataLength);

  memcpy(getRawData() + dataLengthOffset, &dataLength, sizeof(uint16_t));

  return true;

}

uint16_t DnsResource::getCustomDnsClass() const

{

  uint16_t value = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));

  return be16toh(value);

}

void DnsResource::setCustomDnsClass(uint16_t customValue)

{

  memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &customValue, sizeof(uint16_t));

}

}