#include "stdafx.h"

#include <ntddk.h>

#include "WSKTest.h"
#include "lo_send.h"
#include "debug.h"



// WSK related objects
WSK_REGISTRATION WskRegistration;
WSK_PROVIDER_NPI wskProviderNpi;


#ifdef ALLOC_PRAGMA
#pragma NDIS_INIT_FUNCTION(DriverEntry)
#endif // ALLOC_PRAGMA

#if DBG
// Declare the global debug flag for this driver.
ULONG PacketDebugFlag = PACKET_DEBUG_LOUD;

#endif

// 
// Configurable parameters (addresses and ports are in host order)
//

WCHAR g_NPF_PrefixBuffer[512] = L"NPCAP" L"_";

WCHAR* bindT = NULL;

NDIS_STRING g_NPF_Prefix;
NDIS_STRING devicePrefix = NDIS_STRING_CONST("\\Device\\");
NDIS_STRING symbolicLinkPrefix = NDIS_STRING_CONST("\\DosDevices\\");


/*!
\brief Port device extension.

Structure containing some data relative to every adapter on which NPF is bound.
*/
typedef struct _DEVICE_EXTENSION
{
	NDIS_STRING	AdapterName;			///< Name of the adapter.
	PWSTR		ExportString;			///< Name of the exported device, i.e. name that the applications will use 
	///< to open this adapter through WinPcap.
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;



BOOLEAN
NPF_CreateDevice(
IN OUT PDRIVER_OBJECT adriverObjectP,
IN PUNICODE_STRING amacNameP
)
{
	NTSTATUS status;
	PDEVICE_OBJECT devObjP;
	UNICODE_STRING deviceName;
	UNICODE_STRING deviceSymLink;

	TRACE_ENTER();

	IF_LOUD(DbgPrint("\n\ncreateDevice for MAC %ws\n", amacNameP->Buffer););
	if (RtlCompareMemory(amacNameP->Buffer, devicePrefix.Buffer, devicePrefix.Length) < devicePrefix.Length)
	{
		TRACE_EXIT();
		return FALSE;
	}

	deviceName.Length = 0;
	deviceName.MaximumLength = (USHORT)(amacNameP->Length + g_NPF_Prefix.Length + sizeof(UNICODE_NULL));
	deviceName.Buffer = ExAllocatePoolWithTag(PagedPool, deviceName.MaximumLength, '3PWA');

	if (deviceName.Buffer == NULL)
	{
		TRACE_EXIT();
		return FALSE;
	}

	deviceSymLink.Length = 0;
	deviceSymLink.MaximumLength = (USHORT)(amacNameP->Length - devicePrefix.Length + symbolicLinkPrefix.Length + g_NPF_Prefix.Length + sizeof(UNICODE_NULL));

	deviceSymLink.Buffer = ExAllocatePoolWithTag(NonPagedPool, deviceSymLink.MaximumLength, '3PWA');

	if (deviceSymLink.Buffer == NULL)
	{
		ExFreePool(deviceName.Buffer);
		TRACE_EXIT();
		return FALSE;
	}

	RtlAppendUnicodeStringToString(&deviceName, &devicePrefix);
	RtlAppendUnicodeStringToString(&deviceName, &g_NPF_Prefix);
	RtlAppendUnicodeToString(&deviceName, amacNameP->Buffer + devicePrefix.Length / sizeof(WCHAR));

	RtlAppendUnicodeStringToString(&deviceSymLink, &symbolicLinkPrefix);
	RtlAppendUnicodeStringToString(&deviceSymLink, &g_NPF_Prefix);
	RtlAppendUnicodeToString(&deviceSymLink, amacNameP->Buffer + devicePrefix.Length / sizeof(WCHAR));

	IF_LOUD(DbgPrint("Creating device name: %ws\n", deviceName.Buffer);)

		status = IoCreateDevice(adriverObjectP, sizeof(DEVICE_EXTENSION), &deviceName, FILE_DEVICE_TRANSPORT,
		FILE_DEVICE_SECURE_OPEN, FALSE, &devObjP);
	// 		UNICODE_STRING sddl = RTL_CONSTANT_STRING(L"D:P(A;;GA;;;SY)(A;;GA;;;BA)");
	// 	const GUID guidClassNPF = { 0x26e0d1e0L, 0x8189, 0x12e0, { 0x99, 0x14, 0x08, 0x00, 0x22, 0x30, 0x19, 0x04 } };
	// 	status = IoCreateDeviceSecure(adriverObjectP, sizeof(DEVICE_EXTENSION), &deviceName, FILE_DEVICE_TRANSPORT,
	// 		FILE_DEVICE_SECURE_OPEN, FALSE, &sddl, (LPCGUID)&guidClassNPF, &devObjP);

	if (NT_SUCCESS(status))
	{
		PDEVICE_EXTENSION devExtP = (PDEVICE_EXTENSION)devObjP->DeviceExtension;

		IF_LOUD(DbgPrint("Device created successfully\n"););

		devObjP->Flags |= DO_DIRECT_IO;
		RtlInitUnicodeString(&devExtP->AdapterName, amacNameP->Buffer);

		IF_LOUD(DbgPrint("Trying to create SymLink %ws\n", deviceSymLink.Buffer););

		if (IoCreateSymbolicLink(&deviceSymLink, &deviceName) != STATUS_SUCCESS)
		{
			IF_LOUD(DbgPrint("\n\nError creating SymLink %ws\nn", deviceSymLink.Buffer););

			ExFreePool(deviceName.Buffer);
			ExFreePool(deviceSymLink.Buffer);

			devExtP->ExportString = NULL;

			TRACE_EXIT();
			return FALSE;
		}

		IF_LOUD(DbgPrint("SymLink %ws successfully created.\n\n", deviceSymLink.Buffer););

		devExtP->ExportString = deviceSymLink.Buffer;

		ExFreePool(deviceName.Buffer);

		TRACE_EXIT();
		return TRUE;
	}
	else
	{
		IF_LOUD(DbgPrint("\n\nIoCreateDevice status = %x\n", status););

		ExFreePool(deviceName.Buffer);
		ExFreePool(deviceSymLink.Buffer);

		TRACE_EXIT();
		return FALSE;
	}
}

_Use_decl_annotations_
VOID
NPF_Unload(
IN PDRIVER_OBJECT      DriverObject
)
{
	PDEVICE_OBJECT DeviceObject;
	PDEVICE_OBJECT OldDeviceObject;
	PDEVICE_EXTENSION DeviceExtension;
	NDIS_STRING SymLink;

	TRACE_ENTER();

	DeviceObject = DriverObject->DeviceObject;

	while (DeviceObject != NULL)
	{
		OldDeviceObject = DeviceObject;

		DeviceObject = DeviceObject->NextDevice;

		DeviceExtension = OldDeviceObject->DeviceExtension;

		TRACE_MESSAGE3(PACKET_DEBUG_LOUD, "Deleting Adapter %ws, Protocol Handle=xxx, Device Obj=%p (%p)", DeviceExtension->AdapterName.Buffer, DeviceObject, OldDeviceObject);

		if (DeviceExtension->ExportString)
		{
			RtlInitUnicodeString(&SymLink, DeviceExtension->ExportString);

			TRACE_MESSAGE1(PACKET_DEBUG_LOUD, "Deleting SymLink at %p", SymLink.Buffer);

			IoDeleteSymbolicLink(&SymLink);
			ExFreePool(DeviceExtension->ExportString);
		}

		IF_LOUD(DbgPrint("Device successfully deleted.\n\n"););
		IoDeleteDevice(OldDeviceObject);
	}

	NPF_WSKFreeSockets();

	NPF_WSKCleanup();

	// Free the adapters names
	ExFreePool(bindT);

	TRACE_EXIT();
}



_Use_decl_annotations_
NTSTATUS
DriverEntry(
DRIVER_OBJECT* DriverObject,
UNICODE_STRING* RegistryPath
)
{
	NTSTATUS Status = STATUS_SUCCESS;
	UNICODE_STRING macName;

	UNREFERENCED_PARAMETER(RegistryPath);

	// Request NX Non-Paged Pool when available
	ExInitializeDriverRuntime(DrvRtPoolNxOptIn);

	TRACE_ENTER();

	IF_LOUD(DbgPrint("\n\nThis is version [2]!!!.\n");)

		bindT = (PWCHAR)ExAllocatePoolWithTag(PagedPool, 4096, 'NPCA');
	//RtlCopyUnicodeString(bindT, L"\\Device\\{A22932C9-82CB-4080-993B-D5E82CAD06A7}"); //0006, Microsoft KM-TEST Loopback Adapter;
	wcscpy(bindT, L"\\Device\\{A22932C9-82CB-4080-993B-D5E82CAD06B8}"); //different with LBTest
	RtlInitUnicodeString(&macName, bindT);

	NdisInitUnicodeString(&g_NPF_Prefix, g_NPF_PrefixBuffer);

	DriverObject->DriverUnload = NPF_Unload;

	// 
	// Standard device driver entry points stuff.
	//
	// 	DriverObject->MajorFunction[IRP_MJ_CREATE] = NULL;
	// 	DriverObject->MajorFunction[IRP_MJ_CLOSE] = NULL;
	// 	DriverObject->MajorFunction[IRP_MJ_CLEANUP] = NULL;
	// 	DriverObject->MajorFunction[IRP_MJ_READ] = NULL;
	// 	DriverObject->MajorFunction[IRP_MJ_WRITE] = NULL;
	// 	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = NULL;

	Status = NPF_CreateDevice(DriverObject, &macName) ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
	if (Status != STATUS_SUCCESS)
	{
		IF_LOUD(DbgPrint("Failed to create WFP device.\n");)
			TRACE_EXIT();
		return Status;
	}

	Status = NPF_WSKStartup();
	if (!NT_SUCCESS(Status))
	{
		return Status;
	}

	Status = NPF_WSKInitSockets();
	if (!NT_SUCCESS(Status))
	{
		return Status;
	}

	CHAR ICMPv4PacketData[74] =	{
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,		0x00, 0x00, 0x00, 0x00, 0x08, 0x00, 0x45, 0x00,
		0x00, 0x3c, 0x4e, 0xad, 0x00, 0x00, 0x80, 0x01,		0x00, 0x00, 0x7f, 0x00, 0x00, 0x01, 0x7f, 0x00,
		0x00, 0x01, 0x08, 0x00, 0x4d, 0x58, 0x00, 0x01,		0x00, 0x03, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66,
		0x67, 0x68, 0x69, 0x6a, 0x6b, 0x6c, 0x6d, 0x6e,		0x6f, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
		0x77, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67,		0x68, 0x69
	};
	ULONG ICMPv4PacketDataSize = 74;

	CHAR ICMPv6PacketData[94] = {
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,		0x00, 0x00, 0x00, 0x00, 0x86, 0xdd, 0x60, 0x00,
		0x00, 0x00, 0x00, 0x28, 0x3a, 0x80, 0x00, 0x00,		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x80, 0x00,		0xd4, 0xf5, 0x00, 0x01, 0x00, 0x01, 0x61, 0x62,
		0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6a,		0x6b, 0x6c, 0x6d, 0x6e, 0x6f, 0x70, 0x71, 0x72,
		0x73, 0x74, 0x75, 0x76, 0x77, 0x61, 0x62, 0x63,		0x64, 0x65, 0x66, 0x67, 0x68, 0x69
	};
	ULONG ICMPv6PacketDataSize = 94;

	CHAR InvalidPacketData[74] = {
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,		0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x45, 0x00,
		0x00, 0x3c, 0x4e, 0xad, 0x00, 0x00, 0x80, 0x01,		0x00, 0x00, 0x7f, 0x00, 0x00, 0x01, 0x7f, 0x00,
		0x00, 0x01, 0x08, 0x00, 0x4d, 0x58, 0x00, 0x01,		0x00, 0x03, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66,
		0x67, 0x68, 0x69, 0x6a, 0x6b, 0x6c, 0x6d, 0x6e,		0x6f, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
		0x77, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67,		0x68, 0x69
	};
	ULONG InvalidPacketDataSize = 74;

	NPF_WSKSendPacket(ICMPv4PacketData, ICMPv4PacketDataSize);
	NPF_WSKSendPacket(ICMPv6PacketData, ICMPv6PacketDataSize);
	NPF_WSKSendPacket(InvalidPacketData, InvalidPacketDataSize);

	TRACE_EXIT();
	return STATUS_SUCCESS;
};