#include "ntddk.h"
#include "tcphook.h"
#include "dispatch.h"
#include "tdikrnl.h"

PDEVICE_OBJECT g_pFltTcpDevice;

NTSTATUS GetDeviceObject(CCHAR * name, PDEVICE_OBJECT * devObj);
NTSTATUS CreateTcpDevice(PDRIVER_OBJECT pDriverObject,
        PDEVICE_OBJECT * pfltObj,
        CCHAR * name);

/*
* 注册TCP驱动钩子
*/
NTSTATUS HookTcpIrp(IN PDRIVER_OBJECT pDriverObject)
{
        NTSTATUS ntStatus;
       
    DbgPrint("Entering HookTcpIp Routine...\n");

        ntStatus = CreateTcpDevice(pDriverObject,
                &g_pFltTcpDevice,
                "\\Device\\Tcp");
       
        if (!NT_SUCCESS(ntStatus)) {
                return ntStatus;
        }
        ntStatus = GetDeviceObject("\\Device\\Tcp", &g_pFltTcpDevice);
        return ntStatus;
}

NTSTATUS CreateTcpDevice(PDRIVER_OBJECT pDriverObject,
        PDEVICE_OBJECT * pfltObj,
        CCHAR * name)
{
        NTSTATUS ntStatus;
        STRING tcpNameString;
        UNICODE_STRING uitcpDeviceName;
        PDEVICE_EXTENSION pDeviceExtension = NULL;
       
    DbgPrint("Entering CreateTcpDevice Routine...\n");
       
    //创建过滤TCP驱动
        ntStatus = IoCreateDevice(pDriverObject,
        sizeof(DEVICE_EXTENSION),
        NULL,
        FILE_DEVICE_NETWORK,
        0,
        TRUE,
        &g_pFltTcpDevice);
    if (!NT_SUCCESS(ntStatus)) {
        return ntStatus;
    }
        DbgPrint("Create network device success...\n");
        g_pFltTcpDevice->Flags |= DO_BUFFERED_IO;
        g_pFltTcpDevice->Flags &= ~DO_DEVICE_INITIALIZING;
        DbgPrint("Initializing tcpextension...\n");
        RtlZeroMemory(g_pFltTcpDevice->DeviceExtension, sizeof(DEVICE_EXTENSION));
        pDeviceExtension = (PDEVICE_EXTENSION)g_pFltTcpDevice->DeviceExtension;
        //向TCPIP驱动上层中插入自己的TCPIP过滤驱动
        RtlInitAnsiString(&tcpNameString, name);
        RtlAnsiStringToUnicodeString(&uitcpDeviceName, &tcpNameString, TRUE);
        ntStatus = IoAttachDevice(g_pFltTcpDevice, &uitcpDeviceName, &pDeviceExtension->pTcpDevice);
        RtlFreeUnicodeString(&uitcpDeviceName);
        return ntStatus;
}

NTSTATUS GetDeviceObject(CCHAR * name, PDEVICE_OBJECT * devObj)
{
        STRING szName;
        UNICODE_STRING str;
        NTSTATUS status;
        PFILE_OBJECT fileobj;

        RtlInitAnsiString(&szName, name);
        RtlAnsiStringToUnicodeString(&str, &szName, TRUE);

        status = IoGetDeviceObjectPointer(&str, FILE_ALL_ACCESS, &fileobj, devObj);
        if (status == STATUS_SUCCESS) {
                ObDereferenceObject(fileobj);
        }
        RtlFreeUnicodeString(&str);
       
        return status;
}

//处理创建IRP消息,设置CONTROL回调函数
NTSTATUS DispatchControl(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
        PIO_STACK_LOCATION currentIrpStack;
        NTSTATUS ntStatus;
       
        DbgPrint("Entering DispatchControl routine..\n");
       
        currentIrpStack = IoGetCurrentIrpStackLocation(pIrp);

    //如何是DEVICE_CONTROL 则转化到INTERNAL_DEVICE_CONTROL
    if (KeGetCurrentIrql() == PASSIVE_LEVEL) {
                DbgPrint("DispatchControl: this is PASSIVE_LEVEL and MapUserRequest\n");
        ntStatus = TdiMapUserRequest(pDeviceObject, pIrp, currentIrpStack);
    } else {
        ntStatus = STATUS_NOT_IMPLEMENTED;
    }
        if (ntStatus != STATUS_SUCCESS) {
                DbgPrint("DispatchControl: status is not SUCCESS and call TdiDispatchComplete\n");
                ntStatus = TdiDispatchComplete(pDeviceObject, pIrp, FILTER_ALLOW, NULL, NULL);
        }
        return ntStatus;
}

//处理INTERNAL_CONTROL
NTSTATUS DispatchInternalControl(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    PIO_STACK_LOCATION currentIrpStack;
    PIO_STACK_LOCATION nextIrpStack;
    ULONG ulCode;
        PDEVICE_EXTENSION  pDeviceExtension;
       
    DbgPrint("Entering DispatchInternalControl Routine...\n");

    currentIrpStack = IoGetCurrentIrpStackLocation(pIrp);
        nextIrpStack = IoGetNextIrpStackLocation(pIrp);

        pDeviceExtension = (PDEVICE_EXTENSION)pDeviceObject->DeviceExtension;
    ulCode = currentIrpStack->Parameters.DeviceIoControl.IoControlCode;
    DbgPrint("InternalControl, MinorFunction=%x, code=%x\n",
        currentIrpStack->MinorFunction,
        ulCode);
    if (currentIrpStack->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL) {
    }

        return TdiDispatchComplete(pDeviceObject, pIrp, FILTER_ALLOW, NULL, NULL);
//        return IoCallDriver(pDeviceExtension->pTcpDevice, pIrp);
}
//处理CREATE
NTSTATUS DispatchCreate(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    PIO_STACK_LOCATION currentIrpStack;
    PIO_STACK_LOCATION nextIrpStack;
   
    NTSTATUS ntStatus;
    Completion completion;
   
    DbgPrint("Entering DispatchCreate routine...\n");

        memset(&completion, 0, sizeof(Completion));
    currentIrpStack = IoGetCurrentIrpStackLocation(pIrp);
    nextIrpStack = IoGetNextIrpStackLocation(pIrp);
   
        ntStatus = TdiCreate(pDeviceObject, currentIrpStack, pIrp, &completion);
        if (ntStatus != STATUS_SUCCESS) {
                DbgPrint("the TdiCreate not successeed and exit\n");
                return ntStatus;
        }
    return TdiDispatchComplete(pDeviceObject,
                pIrp,
                FILTER_ALLOW,
                completion._routine,
                completion._context);
}

[培训]内核驱动高级班，冲击BAT一流互联网大厂工作，每周日13:00-18:00直播授课