#pragma once

struct dns_question_record
{
    dns_question_record()
    : sQNAME(NULL)
    {
        RtlZeroMemory(&stDnsWireQuestion, sizeof(stDnsWireQuestion));
    }
    ~dns_question_record()
    {
        if (NULL != sQNAME)
        {
            free((void*) sQNAME);
            sQNAME = NULL;
        }
        RtlZeroMemory(&stDnsWireQuestion, sizeof(stDnsWireQuestion));
    }
    LPCTSTR SetQNAME(LPCTSTR sName)
    {
        if (NULL != sQNAME)
        {
            free((void*) sQNAME);
        }
        sQNAME = LPCTSTR(_wcsdup(sName));
        return sQNAME;
    }
    void ClearName()
    {
        if (NULL != sQNAME)
        {
            free((void*) sQNAME);
            sQNAME = NULL;
        }
    }
    LPCTSTR sQNAME;
    DNS_WIRE_QUESTION   stDnsWireQuestion;
};

inline bool IsPrivate(int nA4Count, int nA4PrivateCount, int nA6Count, int nA6PrivaiteCount)
{
    return (((nA4Count > 0) && (nA4Count == nA4PrivateCount) && (nA6Count == 0))
        || ((nA4Count == 0) && (nA6Count > 0) && (nA6Count == nA6PrivaiteCount))
        || ((nA4Count > 0) && (nA4Count == nA4PrivateCount) && (nA6Count > 0) && (nA6Count == nA6PrivaiteCount)));
}


// CDnsMessageWB command target
class CDnsResorceRecord;
class CDnsMessageWB
{
public:
	CDnsMessageWB();
	virtual ~CDnsMessageWB();
    DNS_HEADER m_stDnsHeader;
    CDnsResorceRecord* m_pFirstDnsRRecord;
    CDnsResorceRecord* m_pLastDnsRRecord;
    bool ParseDnsMessageBuffer(PDNS_MESSAGE_BUFFER pDnsMessageBufer, WORD wMessageLength);
    unsigned short m_nAnswerCount;
    unsigned short m_nNameServerCount;
    unsigned short m_nAdditionalCount;
    dns_question_record m_stDnsQuestionRecord;
    CHAR* m_pParsedMessage;
    CString& GetDomainName(CString& sCurrentQNAME, PCHAR& pCurrentBufferPoint);
    int m_nParsedMessageLength;
    CString GetCharSTRING(CString& sCurrentCharSTRING, PCHAR& pCurrentBufferPoint);
    CDnsResorceRecord* RemoveFromTail(void);
    unsigned short m_nQuestionCount;
    void AddRRToTail(CDnsResorceRecord* pDnsResorceRecord);

    int GetRestBytes(PCHAR pCurrentBufferPoint, WORD wByteCount)
    {
        return m_nParsedMessageLength - (int)((pCurrentBufferPoint + wByteCount) - m_pParsedMessage);
    }
    CDnsResorceRecord* RemoveFromHeader(void);
    void RemoveRRAt(CDnsResorceRecord* pRemovedRR);
    void PerformQuerySection(PBYTE& rCurrentBufferPosition);
    CDnsMessageWB(
        CDnsMessageWB* pSourceDnsMessage,
        CDnsResorceRecord* pDnsResorceRecord
        );
    UINT32 PerformDnsMessageBuffer(
        PBYTE pMessageBuffer,
        UINT32 nBufferLength
        );
    void DnsReParseResponceMessage(
        DNS_APPROVAL_CONFIG& rApprovalConfig,
        PBYTE pOutputBuffer,
        size_t nOutputBufferLength
        );
    bool DnsQueryParse(
        PDNS_MESSAGE_BUFFER pDnsBuffer,
        WORD wMessageLength,
        DNS_APPROVAL_CONFIG& rApprovalConfig,
        PBYTE pOutputBuffer,
        size_t nOutputBufferLength
        );
    void ParseQuestionSection(
        PCHAR &rCurrentBuffer,
        WORD wBufferLengt
        );
    void PerformDomanName(
        CString sDomainName,
        PBYTE pMessageBuffer,
        CDnsResorceRecord* pCurrentRR,
        int& rCurrentMessagePosition
        );
    int FindDomanName(
        CString sDomainName,
        PBYTE pMessageBuffer,
        CDnsResorceRecord* pCurrentRR,
        int nCurrentMessagePosition
        );
    void PerformCharString(
        CString sWCharString,
        PBYTE pMessageBuffer,
        int& rCurrentMessagePosition
        );
    DNS_MESSAGE_BUFFER* m_pSavedQueryMessage;
	void PerformUnpackDomanName(
		CString sDomainName,
		PBYTE pMessageBuffer,
		CDnsResorceRecord* pCurrentRR,
		int& rCurrentMessagePosition
		);
    static CString GetRRTypeName(WORD wRRType);
    CString m_sDnsQuestionTypeName;
    void DumpDnsHeader();
    void Clear();
    int m_nCookiesValue;
};