/*
    tdi防火墙 先关注TCP协议   采用过滤驱动
    Codz by Zero 2010-04-08
*/
/*
	typedef struct _NETWORK
	{
		PID                 224
		NAME                iexplore.exe
		PATH                C:\Program Files\Internet Explorer\iexplorer.exe
		LOCAL_IP			192.168.1.22
		LOCAL_PORT			1747
		REMOTE_IP			192.168.1.102
		REMOTE_PORT			8000
		TIME				2010-04-08-09-49-33
	}NETWORK,*PNETWORK;
*/
#ifndef _TDI
#define _TDI 1

#include <ntddk.h>

#endif
#include "mytdi.h"  
#include "filter.h" 
#include <tdikrnl.h>

PDEVICE_OBJECT tcpfltobj = NULL;
PDEVICE_OBJECT tcpoldobj = NULL;

typedef struct _COMPLETION 
{
	PIO_COMPLETION_ROUTINE	routine;
	PVOID					context;
}COMPLETION;

typedef struct
{
	TDI_ADDRESS_INFO	*tai;            //地址信息保存在这里面		
	PFILE_OBJECT		fileobj;	     //irps保存的文件对象
} TDI_CREATE_ADDROBJ2_CTX;

NTKERNELAPI
UCHAR *
PsGetProcessImageFileName(
						  PEPROCESS Process);
//===========================================
NTSTATUS DriverEntry(PDRIVER_OBJECT pDrvObject, PUNICODE_STRING pRegString);
NTSTATUS DispatchCreate(PDEVICE_OBJECT pDevObject, PIRP pIrp);
NTSTATUS DispatchRoutine(PDEVICE_OBJECT pDevObject, PIRP pIrp);
VOID DriverUnload(PDRIVER_OBJECT pDrvObject);
NTSTATUS DispatchIoctl(PDEVICE_OBJECT pDevObject, PIRP pIrp);
VOID GetCurrentProcess(PULONG pid, PCHAR name, PCHAR path);
VOID Get_Process(PIRP irp, COMPLETION *completion);
NTSTATUS Pass_Irp(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context);
NTSTATUS Get_Address(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context);


//==========================================
NTSTATUS DriverEntry(PDRIVER_OBJECT pDrvObject, PUNICODE_STRING pRegString)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING ustrLinkName;
	UNICODE_STRING ustrDevName;    
	PDEVICE_OBJECT pDevObject;
	int i;
	KdPrint(("[tdi] DriverEntry: %S\n",pRegString->Buffer));
	//设置28个分发例程
	for(i=0;i<28;i++)
	{
		pDrvObject->MajorFunction[i]=DispatchRoutine;
	}
	pDrvObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
	pDrvObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchIoctl;
	pDrvObject->DriverUnload = DriverUnload;

	//创建设备绑定Tcp
	status = Filter_AttachDevice(pDrvObject,&tcpfltobj,&tcpoldobj,L"\\Device\\Tcp");
	if(!NT_SUCCESS(status))
	{
		return status;
	}
	//创建符号链接
	RtlInitUnicodeString(&ustrDevName, DEVICE_NAME);
	
	status = IoCreateDevice(pDrvObject, 
		0,
		&ustrDevName, 
		FILE_DEVICE_UNKNOWN,
		0,
		FALSE,
		&pDevObject);
	
	KdPrint(("[sae] Device Name %S",ustrDevName.Buffer));
	
	if(!NT_SUCCESS(status))
	{
		KdPrint(("[sae] IoCreateDevice = 0x%x\n", status));
		return status;
	}
	RtlInitUnicodeString(&ustrLinkName, LINK_NAME);
	
	status = IoCreateSymbolicLink(&ustrLinkName, &ustrDevName);  
	if(!NT_SUCCESS(status))
	{
		KdPrint(("[tdi] IoCreateSymbolicLink = 0x%x\n", status));
		IoDeleteDevice(pDevObject);  
		return status;
	}
	
	KdPrint(("[tdi] SymbolicLink:%S",ustrLinkName.Buffer));
	return STATUS_SUCCESS;
}


VOID DriverUnload(PDRIVER_OBJECT pDrvObject)
{	
	UNICODE_STRING strLink;
	RtlInitUnicodeString(&strLink, LINK_NAME);
	IoDeleteSymbolicLink(&strLink);
	IoDeleteDevice(pDrvObject->DeviceObject);
	//对已生成和绑定的设备进行解绑和删除
	if(tcpfltobj!=NULL)
	{
		IoDeleteDevice(tcpfltobj);
	}
	if(tcpoldobj!=NULL)
	{
		IoDetachDevice(tcpoldobj);
	}
	KdPrint(("[tdi] Unloaded\n"));
}

NTSTATUS DispatchCreate(PDEVICE_OBJECT pDevObject, PIRP pIrp)
{
	
	if(pDevObject == tcpfltobj)
	{
		PFILE_FULL_EA_INFORMATION ea;
		COMPLETION completion;
		memset(&completion, 0, sizeof(COMPLETION));
		ea = (PFILE_FULL_EA_INFORMATION)pIrp->AssociatedIrp.SystemBuffer;
		if(ea != NULL)
		{
			if(	ea->EaNameLength == TDI_TRANSPORT_ADDRESS_LENGTH&&
				memcmp(ea->EaName,TdiTransportAddress,TDI_TRANSPORT_ADDRESS_LENGTH)==0)
			{
				//这里解析得到本地IP和端口
				KdPrint(("Get Local\n"));
				//得到进程 本地IP 端口
				IoCopyCurrentIrpStackLocationToNext(pIrp);
				Get_Process(pIrp,&completion);
				IoCallDriver( tcpoldobj, pIrp ); 
				return STATUS_PENDING;
			}
			else if(ea->EaNameLength == TDI_CONNECTION_CONTEXT_LENGTH&&
				    memcmp(ea->EaName,TdiConnectionContext,TDI_CONNECTION_CONTEXT_LENGTH) == 0)
			{
				//这里解析远程IP 和端口
				KdPrint(("Get Remote\n"));
				IoSkipCurrentIrpStackLocation (pIrp);
				return IoCallDriver( tcpoldobj, pIrp );
			}
		}
		IoSkipCurrentIrpStackLocation (pIrp);
		return IoCallDriver( tcpoldobj, pIrp ); 
	}
	return STATUS_SUCCESS;
}

NTSTATUS DispatchRoutine(PDEVICE_OBJECT pDevObject, PIRP pIrp)
{
	//下发IRP
	if(pDevObject == tcpfltobj)
	{
		IoSkipCurrentIrpStackLocation (pIrp);
		return IoCallDriver( tcpoldobj, pIrp ); 
	}
	return STATUS_SUCCESS;	
}

NTSTATUS DispatchIoctl(PDEVICE_OBJECT pDevObject, PIRP pIrp)
{
	PIO_STACK_LOCATION pIrpStack;
	ULONG uIoControlCode;
	PVOID pIoBuffer;
	ULONG uInSize;
	ULONG uOutSize;
	NTSTATUS status = STATUS_INVALID_DEVICE_REQUEST;

	if(pDevObject == tcpfltobj)
	{
		IoSkipCurrentIrpStackLocation (pIrp);
		return IoCallDriver( tcpoldobj, pIrp );
	}
	
	

	pIrpStack = IoGetCurrentIrpStackLocation(pIrp);
	uIoControlCode = pIrpStack->Parameters.DeviceIoControl.IoControlCode;
	pIoBuffer = pIrp->AssociatedIrp.SystemBuffer;
	uInSize = pIrpStack->Parameters.DeviceIoControl.InputBufferLength;
	uOutSize = pIrpStack->Parameters.DeviceIoControl.OutputBufferLength;

	switch(uIoControlCode)
	{
	case IOCTL_HELLO:
		{
			KdPrint(("[tdi] Hello,I have benn called.."));
			status = STATUS_SUCCESS;
		}
		break;
	}

	if(status == STATUS_SUCCESS)
		pIrp->IoStatus.Information = uOutSize;
	else
		pIrp->IoStatus.Information = 0;

	pIrp->IoStatus.Status = status;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);

	return status;
}


//依据EPROCESS得到进程全路径
VOID GetFullPathByEprocess( ULONG eprocess,PCHAR ProcessImageName )
{
	//原理Eprocess->sectionobject(0x138)->Segment(0x014)->ControlAera(0x000)->FilePointer(0x024)->(FileObject->FileName,FileObject->DeviceObject)
	ULONG object;
	PFILE_OBJECT FileObject;
	UNICODE_STRING FilePath; 
	UNICODE_STRING DosName; 
	STRING AnsiString; 
	
	FileObject = NULL; 
	FilePath.Buffer = NULL; 
	FilePath.Length = 0; 
	*ProcessImageName = 0;  
	
	if(MmIsAddressValid((PULONG)(eprocess+0x138)))//Eprocess->sectionobject(0x138)
	{
		object=(*(PULONG)(eprocess+0x138));
        //KdPrint(("[GetProcessFileName] sectionobject :0x%x\n",object));
		if(MmIsAddressValid((PULONG)((ULONG)object+0x014)))
		{
			object=*(PULONG)((ULONG)object+0x014);
			//KdPrint(("[GetProcessFileName] Segment :0x%x\n",object));
			if(MmIsAddressValid((PULONG)((ULONG)object+0x0)))
			{
				object=*(PULONG)((ULONG_PTR)object+0x0);
				//KdPrint(("[GetProcessFileName] ControlAera :0x%x\n",object));
				if(MmIsAddressValid((PULONG)((ULONG)object+0x024)))
				{
					object=*(PULONG)((ULONG)object+0x024);
					//KdPrint(("[GetProcessFileName] FilePointer :0x%x\n",object));
				}
				else
					return ;
			}
			else
				return ;
		}
		else
			return ;
	}
	else
		return ;
    FileObject=(PFILE_OBJECT)object;
	
	FilePath.Buffer = ExAllocatePool(PagedPool,0x200);
	FilePath.MaximumLength = 0x200; 
    //KdPrint(("[GetProcessFileName] FilePointer :%wZ\n",&FilePointer->FileName));
	ObReferenceObjectByPointer((PVOID)FileObject,0,NULL,KernelMode);//引用计数+1，操作对象
	
	RtlVolumeDeviceToDosName(FileObject-> DeviceObject, &DosName); 
	RtlCopyUnicodeString(&FilePath, &DosName); 
	RtlAppendUnicodeStringToString(&FilePath, &FileObject->FileName); 
	ObDereferenceObject(FileObject); 
	
	RtlUnicodeStringToAnsiString(&AnsiString, &FilePath, TRUE); 
	if ( AnsiString.Length >= 216 ) 
	{ 
		memcpy(ProcessImageName, AnsiString.Buffer, 0x100u); 
		*(ProcessImageName + 215) = 0; 
	} 
	else 
	{ 
		memcpy(ProcessImageName, AnsiString.Buffer, AnsiString.Length); 
		ProcessImageName[AnsiString.Length] = 0; 
	} 
	RtlFreeAnsiString(&AnsiString); 
	ExFreePool(DosName.Buffer); 
	ExFreePool(FilePath.Buffer); 
}
//
VOID GetCurrentProcess(PULONG pid, PCHAR name, PCHAR path)
{
	PEPROCESS Cprocess;
	Cprocess = PsGetCurrentProcess();
	*pid = *(PULONG)((ULONG)Cprocess+0x84);
	strcpy(name ,PsGetProcessImageFileName(Cprocess));
	GetFullPathByEprocess((ULONG)Cprocess,path);
}
//
VOID Get_Process(PIRP irp,COMPLETION* completion)
{
	/*进程信息*/
	ULONG pid;
	CHAR name[126]={0};
	CHAR path[256]={0};
	/*获取地址*/
	/*创建查询IRP*/
	PIRP query_irp;
	GetCurrentProcess(&pid, name, path);
	query_irp = TdiBuildInternalDeviceControlIrp(	TDI_QUERY_INFORMATION,
													tcpoldobj, 
													irps->FileObject, 
													NULL, 
													NULL);
	if (query_irp == NULL) 
	{
		KdPrint(("[tdi_fw] tdi_create: TdiBuildInternalDeviceControlIrp\n"));
		return ;
	}
	/*设置completion*/
	completion->routine = Pass_Irp;      //生成请求的完成函数
	completion->context = query_irp;
	/*设置回调例程*/
	IoSetCompletionRoutine(irp,completion->routine,completion->context, TRUE, TRUE, TRUE);
	/*打印信息*/
	KdPrint(("pid:%d\n",pid));
	KdPrint(("name:%s\n",name));
	KdPrint(("path:%s\n",path));

}
//回调例程Query_Irp
NTSTATUS Pass_Irp(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	NTSTATUS status;
	/*发送query_irp*/
	PIO_STACK_LOCATION irps = IoGetCurrentIrpStackLocation(Irp);
	/*Context参数是我们申请的query_irp*/
	PIRP query_irp = (PIRP)Context;
	/*填充query_irp*/
	TDI_ADDRESS_INFO *ctx = NULL;
	PMDL mdl = NULL;

	if(Irp->PendingReturned)
	{
		IoMarkIrpPending(Irp);
	}

	KdPrint(("Enter Pass_Irp\n"));

	ctx = (TDI_ADDRESS_INFO *)ExAllocatePool(NonPagedPool,sizeof(TDI_ADDRESS_INFO));
	if(ctx == NULL)
	{
		KdPrint(("malloc error\n"));
		return STATUS_SUCCESS;
	}
	mdl = IoAllocateMdl(ctx,sizeof(TDI_ADDRESS_INFO),FALSE,FALSE,NULL);
	if(mdl == NULL)
	{
		KdPrint(("malloc error\n"));
		return STATUS_SUCCESS;
	}
	MmBuildMdlForNonPagedPool(mdl);
	//设置query_irp
	TdiBuildQueryInformation(query_irp,
							tcpoldobj,
							irps->FileObject,
							Get_Address,
							Context,
							TDI_QUERY_ADDRESS_INFO,
							mdl);
	status = IoCallDriver(tcpfltobj,query_irp);
	status = STATUS_SUCCESS;
	return status;
}                                                                                                                                                                                                   
NTSTATUS Get_Address(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp, IN PVOID Context)
{
	if(Irp->MdlAddress)
	{
		TDI_ADDRESS_INFO *ctx = (TDI_ADDRESS_INFO *)MmGetSystemAddressForMdl(Irp->MdlAddress);
		TA_ADDRESS *addr = ctx->Address.Address;
		KdPrint(("Enter Get_Address\n"));
		KdPrint(("[tdi_fw] tdi_create_addrobj_complete2: address: %u\n", 
		 ((TDI_ADDRESS_IP *)(addr->Address))->sin_port));

	}
	return STATUS_SUCCESS;
}

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