#pragma once

extern PNDIS_GENERIC_OBJECT g_pNdisGenericObj;
extern NDIS_HANDLE          g_hNetBufferListPool;
extern NDIS_HANDLE          g_hNetBufferPool;
extern LIST_ENTRY           g_stFreeFileList;
extern LIST_ENTRY           g_stWaitForDoneFileList;
extern LIST_ENTRY           g_stWaitForQuery2FileList;
extern LIST_ENTRY           g_stWaitForSetFilterFileList;
extern LIST_ENTRY           g_stAcceptableList;
extern LIST_ENTRY           g_stUdpInboundList;
extern LIST_ENTRY           g_stUdpALEAcceptRecvList;
extern LIST_ENTRY           g_stUdpAcceptNotifyList;
extern LIST_ENTRY           g_stUdpALEAcceptPendedList;
extern LIST_ENTRY           g_stUdpALEFlowEstablishList;
extern LIST_ENTRY           g_stFlowedFileList;
extern LIST_ENTRY           g_stPendedPacketsQueue;
extern KEVENT               g_oWorkerEvent;
extern KEVENT               g_oWorkerDoneEvent;
extern HANDLE               g_hInjectionHandleV4;
extern HANDLE               g_hInjectionHandleV6;
extern UINT32               g_nInboundTransportV4Id;
extern UINT32               g_nInboundTransportV6Id;
extern UINT32               g_nAleAcceptRecvV4Id;
extern UINT32               g_nAleAcceptRecvV6Id;
extern UINT32               g_nAleFlowEstablishedV4Id;
extern UINT32               g_nAleFlowEstablishedV6Id;
extern UINT32               g_nStreamIdV4;
extern UINT32               g_nStreamIdV6;
extern UINT32               g_nDatagramIdV4;
extern UINT32               g_nDatagramIdV6;
extern FIREWALL_CONFIG      g_stFirewallConfig;
extern BOOLEAN              g_bFiltersChangeInProgress;
extern BOOLEAN              g_bDriverUnload;
extern PVOID                g_pWorkerThreadObj;
extern WDFQUEUE             g_wdfWaitRestartQueue;
extern UINT32               g_uiRestartStatus;

WDF_DECLARE_CONTEXT_TYPE(FLOW_DATA);
WDF_DECLARE_CONTEXT_TYPE(ADDRESSES_DESCRIPTOR);

#define HOST_ADDRESS(nOctet1,nOctet2,nOctet3,nOctet4) (UINT32) ((nOctet4&0xFF) + ((nOctet3&0xFF) << 8) + ((nOctet2&0xFF) << 16) + ((nOctet1&0xFF) << 24))

FORCEINLINE
BOOLEAN IsListHeadBad(PLIST_ENTRY pListEntry)
{
    return (((pListEntry->Flink)->Blink != pListEntry) || ((pListEntry->Blink)->Flink != pListEntry));
}

BOOLEAN IsListEntryInList(
    PLIST_ENTRY pListEntry,
    PLIST_ENTRY pListHead
    );


FORCEINLINE UINT32 hostAddress(
    UINT32 nOctet1,
    UINT32 nOctet2,
    UINT32 nOctet3,
    UINT32 nOctet4
    )
{
    UINT32  nRet;
    nRet = (nOctet4 & 0xFF) + ((nOctet3 & 0xFF) << 8) + ((nOctet2 & 0xFF) << 16) + ((nOctet1 & 0xFF) << 24);
    return nRet;
}

#define NET_ADDRESS(nOctet1,nOctet2,nOctet3,nOctet4) (UINT32) ((nOctet1&0xFF) + ((nOctet2&0xFF) << 8) + ((nOctet3&0xFF) << 16) + ((nOctet4&0xFF) << 24))

FORCEINLINE UINT32 netAddress(
    UINT32 nOctet1,
    UINT32 nOctet2,
    UINT32 nOctet3,
    UINT32 nOctet4
    )
{
    UINT32  nRet;
    nRet = (nOctet1 & 0xFF) + ((nOctet2 & 0xFF) << 8) + ((nOctet3 & 0xFF) << 16) + ((nOctet4 & 0xFF) << 24);
    return nRet;
}

static const UINT32 c_ulLocalHost = HOST_ADDRESS(127, 0, 0, 1);
static const UINT32 c_ulTestIpV4 = HOST_ADDRESS(194, 85, 61, 58);

FORCEINLINE void push_css(
    context_state* arContextStateStack,
    context_state enNewContextSate
    )
{
    int     i;
    for (i = 15; i>0; i--)
    {
        arContextStateStack[i] = arContextStateStack[i - 1];
    }
    arContextStateStack[0] = enNewContextSate;
}

FORCEINLINE void push_fos(
    flow_data_owner* arFlowDataOwnerStack,
    context_state enNewFlowDataOwner
    )
{
    int     i;
    for (i = 15; i>0; i--)
    {
        arFlowDataOwnerStack[i] = arFlowDataOwnerStack[i - 1];
    }
    arFlowDataOwnerStack[0] = enNewFlowDataOwner;
}

NTSTATUS
dnsfirewallCoInitialize(PDEVICE_OBJECT deviceObject);

NTSTATUS
dnsfirewallCoRegisterCallouts(PDEVICE_OBJECT deviceObject);

void
onInboundTransportV4Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onInboundTransportV6Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onAleAcceptRecvV4Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onAleAcceptRecvV6Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onAleFlowEstablishedV4Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

NTSTATUS
onFiltersChangeNotify(
   _In_ FWPS_CALLOUT_NOTIFY_TYPE notifyType,
   _In_ const GUID* filterKey,
   _Inout_ const FWPS_FILTER* filter
   );

void
onAleFlowEstablishedV6Classify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onStreamClassify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void
onDnsFlowDeletion(
   _In_ UINT16 layerId,
   _In_ UINT32 calloutId,
   _In_ UINT64 flowContext
   );

void
onDatagramDataClassify(
    _In_ const FWPS_INCOMING_VALUES* inFixedValues,
    _In_ const FWPS_INCOMING_METADATA_VALUES* inMetaValues,
    _Inout_opt_ void* layerData,
    _In_opt_ const void* classifyContext,
    _In_ const FWPS_FILTER* filter,
    _In_ UINT64 flowContext,
    _Inout_ FWPS_CLASSIFY_OUT* classifyOut
    );

void FilterDataCleanup(
    PADDRESSES_DESCRIPTOR pDeviceAddressesDescriptor
    );

void
clonedDatagramDnsQueryCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

void
GetStreamData(
    NET_BUFFER_LIST* pNetBufferList,
    BYTE* pDataBuffer,
    UINT32 nDataBufferLength,
    UINT32* pByteCopied
    );

void
clonedStreamDnsQueryCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

void
allocatedStreamDnsQueryCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

void
clonedStreamDnsResponceCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

void
allocatedStreamDnsResponceCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

void
clonedDatagramDnsResponceCompletionFn(
    __in     VOID *context,
    __inout  NET_BUFFER_LIST *netBufferList,
    __in     BOOLEAN dispatchLevel
    );

static const UINT32 c_ulTestSubnetPart = (const UINT32) HOST_ADDRESS(85, 21, 0, 0);
static const UINT32 c_ulTestSubnetMask = (const UINT32) HOST_ADDRESS(255, 255, 0, 0);

KSTART_ROUTINE dnsfirewallWorker;