// name_resolution.c

//#include "..\intrface.h"
//#include "CprDrvr.h"
#include "pch.h"

//#define USE_WSK

#ifdef USE_WSK

#ifdef __cplusplus
extern "C" {
#endif

#include <ntddk.h>
#include <wsk.h>

#ifdef __cplusplus
}
#endif

// WSK Client Dispatch table that denotes the WSK version
// that the WSK application wants to use and optionally a pointer
// to the WskClientEvent callback function
const WSK_CLIENT_DISPATCH WskAppDispatch = {
	MAKE_WSK_VERSION(1,0),	// Use WSK version 1.0
	0,						// Reserved
	NULL					// WskClientEvent callback not required for WSK version 1.0
};

void
DriverEntry_Wsk()
{
    NTSTATUS            status;
	WSK_CLIENT_NPI		wskClientNpi;

	g_Data.wskRegistred = FALSE;
	g_Data.wskProviderValid = FALSE;

	// Register the WSK application
	wskClientNpi.ClientContext = NULL;
	wskClientNpi.Dispatch = &WskAppDispatch;
	status = WskRegister(&wskClientNpi, &g_Data.wskRegistration);

	if(NT_SUCCESS(status))
		g_Data.wskRegistred = TRUE;
}

void
CprUnload_Wsk()
{
	if (g_Data.wskProviderValid)
	{
		WskReleaseProviderNPI(&g_Data.wskRegistration);
		g_Data.wskProviderValid = FALSE;
	}

	if (g_Data.wskRegistred)
	{
		WskDeregister(&g_Data.wskRegistration);
		g_Data.wskRegistred = FALSE;
	}
}

// Must call if (IS_WIN7_OR_GREATER) before calling this function
// to make sure it is being called on the correct OS.
NTSTATUS
ResolveHostName(
	PCPR_DEVICE_EXTENSION	pDevExt,
	char*					hostName
	)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	CprDebugPrint1(pDevExt, DBG_CREATECLOSE, DBG_TRACE, __FUNCTION__"++. hostName %s", hostName);

	if (g_Data.wskRegistred)
	{
		ADDRINFOEXW aiHints;
		wchar_t hostNameWide[IP_MAC_HOST_LEN];

		if (g_Data.wskProviderValid == FALSE)
		{
			status = WskCaptureProviderNPI(&g_Data.wskRegistration, 1000, &g_Data.wskProviderNpi);
			if(NT_SUCCESS(status))
				g_Data.wskProviderValid = TRUE;
		}

		if (g_Data.wskProviderValid)
		{
			int i;
			RtlZeroMemory(&aiHints, sizeof(aiHints));
			aiHints.ai_family = AF_INET;
			aiHints.ai_socktype = SOCK_STREAM;
			aiHints.ai_protocol = IPPROTO_TCP;

			for ( i = 0; i < IP_MAC_HOST_LEN; ++i )
			{
			   hostNameWide[i] = (wchar_t)hostName[i];
			}

			// Only Windows 7 on up.
			status = NameResolution(hostNameWide, NULL, &aiHints, &g_Data.wskProviderNpi, hostName);
		}

		CprDebugPrint2(pDevExt, DBG_CREATECLOSE, DBG_TRACE, __FUNCTION__"--. status %d hostName %s", status, hostName);
	}

	return status;
}

NTSTATUS
SyncIrpCompletionRoutine(
    __in PDEVICE_OBJECT Reserved,
    __in PIRP Irp,
    __in PVOID Context
    )
{    
    PKEVENT compEvent = (PKEVENT)Context;
    UNREFERENCED_PARAMETER(Reserved);
    UNREFERENCED_PARAMETER(Irp);
    KeSetEvent(compEvent, 2, FALSE);    
    return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS
NameResolution(
    __in PCWSTR				NodeName,
    __in_opt PCWSTR			ServiceName,
    __in_opt PADDRINFOEXW	Hints,
    __in PWSK_PROVIDER_NPI	WskProviderNpi,
	char *					hostName
    )
{
    NTSTATUS status;
    PIRP irp;
    KEVENT completionEvent;
    UNICODE_STRING uniNodeName, uniServiceName, *uniServiceNamePtr;
    PADDRINFOEXW results;

	CprDebugPrint1(NULL, DBG_CREATECLOSE, DBG_TRACE, __FUNCTION__"++. hostName %s", hostName);

    PAGED_CODE();
    //
    // Initialize UNICODE_STRING structures for NodeName and ServiceName 
    //
 
    RtlInitUnicodeString(&uniNodeName, NodeName);

    if(ServiceName == NULL)
	{
        uniServiceNamePtr = NULL;
    }
    else
	{
        RtlInitUnicodeString(&uniServiceName, ServiceName);
        uniServiceNamePtr = &uniServiceName;
    }

    //
    // Use an event object to synchronously wait for the 
    // WskGetAddressInfo request to be completed. 
    //
 
    KeInitializeEvent(&completionEvent, SynchronizationEvent, FALSE);

    //
    // Allocate an IRP for the WskGetAddressInfo request, and set the 
    // IRP completion routine, which will signal the completionEvent
    // when the request is completed.
    //
 
    irp = IoAllocateIrp(1, FALSE);
    if(irp == NULL)
	{
        return STATUS_INSUFFICIENT_RESOURCES;
    }        

    IoSetCompletionRoutine(irp, SyncIrpCompletionRoutine, 
		&completionEvent, TRUE, TRUE, TRUE);

    //
    // Make the WskGetAddressInfo request.
    //
 
    WskProviderNpi->Dispatch->WskGetAddressInfo (
        WskProviderNpi->Client,
        &uniNodeName,
        uniServiceNamePtr,
        NS_ALL,
        NULL, // Provider
        Hints,
        &results, 
        NULL, // OwningProcess
        NULL, // OwningThread
        irp);

    //
    // Wait for completion. Note that processing of name resolution results
    // can also be handled directly within the IRP completion routine, but
    // for simplicity, this example shows how to wait synchronously for 
    // completion.
    //
 
    KeWaitForSingleObject(&completionEvent, Executive, 
                        KernelMode, FALSE, NULL);

    status = irp->IoStatus.Status;

    IoFreeIrp(irp);

    if(!NT_SUCCESS(status))
	{
		CprDebugPrint1(NULL, DBG_CREATECLOSE, DBG_TRACE, __FUNCTION__"--. Error. status %d", status);
        return status;
    }

    //
    // Process the name resolution results by iterating through the addresses
    // within the returned ADDRINFOEXW structure.
    //
 
	CprDebugPrint1(NULL, DBG_CREATECLOSE, DBG_INFO, __FUNCTION__" results = %p", results);

	if (results != NULL)
	{
		struct sockaddr_in *si = (struct sockaddr_in *)results->ai_addr;
		CprDebugPrint1(NULL, DBG_CREATECLOSE, DBG_INFO, __FUNCTION__" results->ai_next = %p", results->ai_next);

		inet_ntoa(si->sin_addr.S_un.S_addr, hostName, IP_MAC_HOST_LEN);

		CprDebugPrint3(NULL, DBG_CREATECLOSE, DBG_INFO,
			__FUNCTION__" results->ai_addr %p, si->sin_addr.S_un.S_addr %X, hostName %s",
			results->ai_addr, si->sin_addr.S_un.S_addr, hostName);

		for( ; results->ai_next; results = results->ai_next)
		{
			si = (struct sockaddr_in *)results->ai_addr;

			CprDebugPrint3(NULL, DBG_CREATECLOSE, DBG_INFO,
				__FUNCTION__" results->ai_addr %p, si->sin_addr.S_un.S_addr %X, hostName %s",
				results->ai_addr, si->sin_addr.S_un.S_addr, hostName);

			inet_ntoa(si->sin_addr.S_un.S_addr, hostName, IP_MAC_HOST_LEN);
		}
	}

    //
    // Release the returned ADDRINFOEXW structure when no longer needed.
    //
 
    WskProviderNpi->Dispatch->WskFreeAddressInfo(
        WskProviderNpi->Client,
        results);

	CprDebugPrint2(NULL, DBG_CREATECLOSE, DBG_TRACE, __FUNCTION__"--. status %d hostName %s", status, hostName);

    return status;
} 

#else

void
DriverEntry_Wsk()
{
}

void
CprUnload_Wsk()
{
}

NTSTATUS
ResolveHostName(
	PCPR_DEVICE_EXTENSION	pDevExt,
	char*					hostName
	)
{
	return STATUS_UNSUCCESSFUL;
}

#endif