#include "Driver.h"


VOID DriverUnload(PDRIVER_OBJECT driver_object) {
	UNREFERENCED_PARAMETER(driver_object);
	DbgPrint("DriverUnload\n");


}

WSK_CLIENT_DISPATCH kClientDispatch = {
	MAKE_WSK_VERSION(1,0),	// Use WSK version 1.0
	0,						// Reserved
	NULL					// WskClientEvent callback not required for WSK version 1.0
};

NTSTATUS DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath) {

	NTSTATUS nStatus = STATUS_UNSUCCESSFUL;
	DbgPrint("DriverEntry\n");

	DriverObject->DriverUnload = DriverUnload;


	WSK_REGISTRATION registration_;
	WSK_CLIENT_NPI client_npi_;
	WSK_PROVIDER_NPI provider_npi_;


	client_npi_.Dispatch = &kClientDispatch;

	nStatus = WskRegister(&client_npi_, &registration_);

	if (!NT_SUCCESS(nStatus)) {
		return nStatus;
	}

	nStatus = WskCaptureProviderNPI(&registration_,WSK_INFINITE_WAIT, &provider_npi_);

	if (!NT_SUCCESS(nStatus)) {
		WskDeregister(&registration_);
		return nStatus;
	}


	ADDRINFOEXW hints = { 0 };
	hints.ai_flags |= AI_CANONNAME;
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;


	PIRP irp = NULL;
	KEVENT event;


	PADDRINFOEXW remote_addr_info = NULL;
	PSOCKADDR_IN remote_addr = NULL;

	KeInitializeEvent(&event, SynchronizationEvent, FALSE);

	UNICODE_STRING host_unicode, uniServiceName;

	RtlInitUnicodeString(&host_unicode, "baidu.com");

	RtlInitUnicodeString(&uniServiceName, "80");

	irp = IoAllocateIrp(1, FALSE);
	if (!irp) {
		nStatus = STATUS_INSUFFICIENT_RESOURCES;
		goto ret;
	}
	

	IoSetCompletionRoutine(irp, CompletionRoutine, &event, TRUE, TRUE, TRUE);

	nStatus = provider_npi_.Dispatch->WskGetAddressInfo(provider_npi_.Client, &host_unicode, &uniServiceName, NS_ALL, NULL, &hints, &remote_addr_info, NULL, NULL, irp);

	//DbgBreakPoint();

	KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);//&timeout


	if (nStatus == STATUS_TIMEOUT) {
		IoCancelIrp(irp);
		KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
		goto ret;
	}

	if (!NT_SUCCESS(irp->IoStatus.Status)) {
		nStatus = irp->IoStatus.Status;

		KdPrint(("windbg>>>%llx", nStatus));

		goto ret;
	}

	remote_addr = (PSOCKADDR_IN)remote_addr_info->ai_addr;

	KdPrint(("%s", remote_addr_info->ai_addr->sa_data));

	goto ret;


ret:

	if (irp) {
		IoFreeIrp(irp);
	}

	if (remote_addr_info) {
		provider_npi_.Dispatch->WskFreeAddressInfo(
			provider_npi_.Client,
			remote_addr_info
		);
	}


	WskReleaseProviderNPI(&registration_);
	WskDeregister(&registration_);


	return STATUS_SUCCESS;

}

NTSTATUS CompletionRoutine(PDEVICE_OBJECT reserved, PIRP irp, PVOID ctx) {
	UNREFERENCED_PARAMETER(reserved);
	UNREFERENCED_PARAMETER(irp);
	//_Analysis_assume_(ctx != NULL);

	KeSetEvent((PKEVENT)ctx, IO_NO_INCREMENT, FALSE);
	KdPrint(("windbg>>>ctx%p\n", ctx));

	return STATUS_MORE_PROCESSING_REQUIRED;
}