#define LOG_MODULE PacketLogModuleDhcpLayer

#include "DhcpLayer.h"
#include "Logger.h"

namespace pcpp
{

#define DHCP_MAGIC_NUMBER 0x63538263


DhcpOption DhcpOptionBuilder::build() const
{
	size_t recSize = 2 * sizeof(uint8_t) + m_RecValueLen;
	uint8_t recType = static_cast<uint8_t>(m_RecType);

	if ((recType == DHCPOPT_END || recType == DHCPOPT_PAD))
	{
		if (m_RecValueLen != 0)
		{
			PCPP_LOG_ERROR("Can't set DHCP END option or DHCP PAD option with size different than 0, tried to set size " << (int)m_RecValueLen);
			return DhcpOption(nullptr);
		}

		recSize = sizeof(uint8_t);
	}

	uint8_t* recordBuffer = new uint8_t[recSize];
	memset(recordBuffer, 0, recSize);
	recordBuffer[0] = recType;
	if (recSize > 1)
	{
		recordBuffer[1] = static_cast<uint8_t>(m_RecValueLen);
		if (m_RecValue != nullptr)
			memcpy(recordBuffer+2, m_RecValue, m_RecValueLen);
		else
			memset(recordBuffer+2, 0, m_RecValueLen);
	}

	return DhcpOption(recordBuffer);
}

DhcpLayer::DhcpLayer(uint8_t* data, size_t dataLen, Layer* prevLayer, Packet* packet) : Layer(data, dataLen, prevLayer, packet)
{
	 m_Protocol = DHCP;
}

void DhcpLayer::initDhcpLayer(size_t numOfBytesToAllocate)
{
	m_DataLen = numOfBytesToAllocate;
	m_Data = new uint8_t[numOfBytesToAllocate];
	memset(m_Data, 0, numOfBytesToAllocate);
	m_Protocol = DHCP;
}

DhcpLayer::DhcpLayer() : Layer()
{
	initDhcpLayer(sizeof(dhcp_header));
}

DhcpLayer::DhcpLayer(DhcpMessageType msgType, const MacAddress& clientMacAddr) : Layer()
{
	initDhcpLayer(sizeof(dhcp_header) + 4*sizeof(uint8_t));

	setClientHardwareAddress(clientMacAddr);

	uint8_t* msgTypeOptionPtr = m_Data + sizeof(dhcp_header);
	msgTypeOptionPtr[0] = (uint8_t)DHCPOPT_DHCP_MESSAGE_TYPE; // option code
	msgTypeOptionPtr[1] = 1; // option len
	msgTypeOptionPtr[2] = (uint8_t)msgType; // option data - message type

	msgTypeOptionPtr[3] =  (uint8_t)DHCPOPT_END;
}

MacAddress DhcpLayer::getClientHardwareAddress() const
{
	dhcp_header* hdr = getDhcpHeader();
	if (hdr != nullptr && hdr->hardwareType == 1 && hdr->hardwareAddressLength == 6)
		return MacAddress(hdr->clientHardwareAddress);

	PCPP_LOG_DEBUG("Hardware type isn't Ethernet or hardware addr len != 6, returning MacAddress:Zero");

	return MacAddress::Zero;
}

void DhcpLayer::setClientHardwareAddress(const MacAddress& addr)
{
	dhcp_header* hdr = getDhcpHeader();
	hdr->hardwareType = 1; // Ethernet
	hdr->hardwareAddressLength = 6; // MAC address length
	addr.copyTo(hdr->clientHardwareAddress);
}

void DhcpLayer::computeCalculateFields()
{
	dhcp_header* hdr = getDhcpHeader();

	hdr->magicNumber = DHCP_MAGIC_NUMBER;

	DhcpMessageType msgType = getMessageType();
	switch(msgType)
	{
	case DHCP_DISCOVER:
	case DHCP_REQUEST:
	case DHCP_DECLINE:
	case DHCP_RELEASE:
	case DHCP_INFORM:
	case DHCP_UNKNOWN_MSG_TYPE:
		hdr->opCode = DHCP_BOOTREQUEST;
		break;
	case DHCP_OFFER:
	case DHCP_ACK:
	case DHCP_NAK:
		hdr->opCode = DHCP_BOOTREPLY;
		break;
	default:
		break;
	}

	hdr->hardwareType = 1; //Ethernet
	hdr->hardwareAddressLength = 6; // MAC address length
}

std::string DhcpLayer::toString() const
{
	std::string msgType = "Unknown";
	switch (getMessageType())
	{
	case DHCP_DISCOVER:
	{
		msgType = "Discover";
		break;
	}
	case DHCP_OFFER:
	{
		msgType = "Offer";
		break;
	}
	case DHCP_REQUEST:
	{
		msgType = "Request";
		break;
	}
	case DHCP_DECLINE:
	{
		msgType = "Decline";
		break;
	}
	case DHCP_ACK:
	{
		msgType = "Acknowledge";
		break;
	}
	case DHCP_NAK:
	{
		msgType = "Negative Acknowledge";
		break;
	}
	case DHCP_RELEASE:
	{
		msgType = "Release";
		break;
	}
	case DHCP_INFORM:
	{
		msgType = "Inform";
		break;
	}
	default:
		break;

	}

	return "DHCP layer (" + msgType + ")";
}

DhcpMessageType DhcpLayer::getMessageType() const
{
	DhcpOption opt = getOptionData(DHCPOPT_DHCP_MESSAGE_TYPE);
	if (opt.isNull())
		return DHCP_UNKNOWN_MSG_TYPE;

	return (DhcpMessageType)opt.getValueAs<uint8_t>();
}

bool DhcpLayer::setMessageType(DhcpMessageType msgType)
{
	if (msgType == DHCP_UNKNOWN_MSG_TYPE)
		return false;

	DhcpOption opt = getOptionData(DHCPOPT_DHCP_MESSAGE_TYPE);
	if (opt.isNull())
	{
		opt = addOptionAfter(DhcpOptionBuilder(DHCPOPT_DHCP_MESSAGE_TYPE, (uint8_t)msgType), DHCPOPT_UNKNOWN);
		if (opt.isNull())
			return false;
	}

	opt.setValue<uint8_t>((uint8_t)msgType);
	return true;
}

DhcpOption DhcpLayer::getOptionData(DhcpOptionTypes option) const
{
	return m_OptionReader.getTLVRecord((uint8_t)option, getOptionsBasePtr(), getHeaderLen() - sizeof(dhcp_header));
}

DhcpOption DhcpLayer::getFirstOptionData() const
{
	return m_OptionReader.getFirstTLVRecord(getOptionsBasePtr(), getHeaderLen() - sizeof(dhcp_header));
}

DhcpOption DhcpLayer::getNextOptionData(DhcpOption dhcpOption) const
{
	return m_OptionReader.getNextTLVRecord(dhcpOption, getOptionsBasePtr(), getHeaderLen() - sizeof(dhcp_header));
}

size_t DhcpLayer::getOptionsCount() const
{
	return m_OptionReader.getTLVRecordCount(getOptionsBasePtr(), getHeaderLen() - sizeof(dhcp_header));
}

DhcpOption DhcpLayer::addOptionAt(const DhcpOptionBuilder& optionBuilder, int offset)
{
	DhcpOption newOpt = optionBuilder.build();

	if (newOpt.isNull())
	{
		PCPP_LOG_ERROR("Cannot build new option of type " << (int)newOpt.getType());
		return DhcpOption(nullptr);
	}

	size_t sizeToExtend = newOpt.getTotalSize();

	if (!extendLayer(offset, sizeToExtend))
	{
		PCPP_LOG_ERROR("Could not extend DhcpLayer in [" << newOpt.getTotalSize() << "] bytes");
		newOpt.purgeRecordData();
		return DhcpOption(nullptr);
	}

	memcpy(m_Data + offset, newOpt.getRecordBasePtr(), newOpt.getTotalSize());

	uint8_t* newOptPtr = m_Data + offset;

	m_OptionReader.changeTLVRecordCount(1);

	newOpt.purgeRecordData();

	return DhcpOption(newOptPtr);
}

DhcpOption DhcpLayer::addOption(const DhcpOptionBuilder& optionBuilder)
{
	int offset = 0;
	DhcpOption endOpt = getOptionData(DHCPOPT_END);
	if (!endOpt.isNull())
		offset = endOpt.getRecordBasePtr() - m_Data;
	else
		offset = getHeaderLen();

	return addOptionAt(optionBuilder, offset);
}

DhcpOption DhcpLayer::addOptionAfter(const DhcpOptionBuilder& optionBuilder, DhcpOptionTypes prevOption)
{
	int offset = 0;

	DhcpOption prevOpt = getOptionData(prevOption);

	if (prevOpt.isNull())
	{
		offset = sizeof(dhcp_header);
	}
	else
	{
		offset = prevOpt.getRecordBasePtr() + prevOpt.getTotalSize() - m_Data;
	}

	return addOptionAt(optionBuilder, offset);
}

bool DhcpLayer::removeOption(DhcpOptionTypes optionType)
{
	DhcpOption optToRemove = getOptionData(optionType);
	if (optToRemove.isNull())
	{
		return false;
	}

	int offset = optToRemove.getRecordBasePtr() - m_Data;

	if (!shortenLayer(offset, optToRemove.getTotalSize()))
	{
		return false;
	}

	m_OptionReader.changeTLVRecordCount(-1);
	return true;
}

bool DhcpLayer::removeAllOptions()
{
	int offset = sizeof(dhcp_header);

	if (!shortenLayer(offset, getHeaderLen()-offset))
		return false;

	m_OptionReader.changeTLVRecordCount(0-getOptionsCount());
	return true;
}


}