/*
 * Copyright (c) 2004 Security Architects Corporation. All rights reserved.
 *
 * Module Name:
 *
 *		network.c
 *
 * Abstract:
 *
 *		This module defines various routines used for hooking the Transport Driver Interface (TDI) network routines.
 *
 * Author:
 *
 *		Eugene Tsyrklevich 12-Mar-2004
 *
 * Revision History:
 *
 *		None.
 */


#include <NTDDK.h>
#include <tdikrnl.h>
#include <ctype.h>
#include "network.h"
#include "hookproc.h"
#include "userland.h"
#include "learn.h"
#include "policy.h"
#include "log.h"


#ifdef ALLOC_PRAGMA
#pragma alloc_text (INIT, InstallNetworkHooks)
#pragma alloc_text (PAGE, RemoveNetworkHooks)
#endif


//XXX fast io is not handled. TdiDispatchFastDeviceControl


PDEVICE_OBJECT	pTcpDevice = NULL, pTcpDeviceOriginal = NULL;
PDEVICE_OBJECT	pUdpDevice = NULL, pUdpDeviceOriginal = NULL;
PDEVICE_OBJECT	pIpDevice = NULL, pIpDeviceOriginal = NULL;

#if DBG
int	HookedTDIRunning = 0;
#endif


/*
 * TdiStub() XXX remove
 *
 * Description:
 *		.
 *
 * Parameters:
 *		pIrp - IRP (I/O Request Packet) request.
 *		pIrpStack - .
 *		pCompletion - .
 *
 * Returns:
 *		STATUS_SUCCESS.
 */

NTSTATUS
TdiStub(IN PIRP pIrp, IN PIO_STACK_LOCATION pIrpStack, OUT PTDI_CALLBACK pCompletion, IN ULONG DeviceType)
{
//	LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("TdiStub(%x %x %x)\n", pIrp, pIrpStack, pCompletion));
	return STATUS_SUCCESS;
}



NTSTATUS
TdiSetEventHandler(IN PIRP pIrp, IN PIO_STACK_LOCATION pIrpStack, OUT PTDI_CALLBACK pCompletion, IN ULONG DeviceType)
{
	PTDI_REQUEST_KERNEL_SET_EVENT	r = (PTDI_REQUEST_KERNEL_SET_EVENT) &pIrpStack->Parameters;


	if (r->EventType != TDI_EVENT_CONNECT)
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiSetEventHandler: %x %x %x\n", CURRENT_PROCESS_PID, r->EventType, r->EventHandler, r->EventContext));
		return STATUS_SUCCESS;
	}


	if (r->EventHandler == NULL)
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiSetEventHandler: TDI_EVENT_CONNECT deregistration %x %x %x\n", CURRENT_PROCESS_PID, r->EventHandler, r->EventContext, pIrpStack->FileObject));
		return STATUS_SUCCESS;
	}


	LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiSetEventHandler: TDI_EVENT_CONNECT %x %x %x\n", CURRENT_PROCESS_PID, r->EventHandler, r->EventContext, pIrpStack->FileObject));


	return STATUS_SUCCESS;
}



NTSTATUS
TdiConnect(IN PIRP pIrp, IN PIO_STACK_LOCATION pIrpStack, OUT PTDI_CALLBACK pCompletion, IN ULONG DeviceType)
{
	/*
	 * IrpSp->Parameters 
	 *
	 * Pointer to a TDI_REQUEST_KERNEL_CONNECT structure, equivalent to the TDI_REQUEST_KERNEL structure.
	 */

	PTDI_REQUEST_KERNEL_CONNECT		ConnectInfo = (PTDI_REQUEST_KERNEL_CONNECT) &pIrpStack->Parameters;
	PTRANSPORT_ADDRESS				pTransportAddress;
	PTA_ADDRESS						pAddress;
	PTDI_ADDRESS_IP					ip;
	CHAR							NETWORKNAME[MAX_PATH];
	PCHAR							FunctionName = "TdiConnect";


	HOOK_ROUTINE_ENTER();


	if (! MmIsAddressValid(ConnectInfo) ||
		! MmIsAddressValid(ConnectInfo->RequestConnectionInformation) ||
		! MmIsAddressValid(ConnectInfo->RequestConnectionInformation->RemoteAddress))
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("TdiConnect: MmIsAddressValid failed\n"));
		HOOK_ROUTINE_EXIT(STATUS_SUCCESS);
	}


	pTransportAddress = (PTRANSPORT_ADDRESS) ConnectInfo->RequestConnectionInformation->RemoteAddress;

	pAddress = (PTA_ADDRESS) pTransportAddress->Address;

	/* verify that the specified address is a single IP address */
	if (pTransportAddress->TAAddressCount != 1 ||
		pAddress->AddressType != TDI_ADDRESS_TYPE_IP ||
		pAddress->AddressLength != TDI_ADDRESS_LENGTH_IP)        
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiConnect: Invalid address detected\n", CURRENT_PROCESS_PID));
		HOOK_ROUTINE_EXIT(STATUS_SUCCESS);
	}

	ip = (PTDI_ADDRESS_IP) &pAddress->Address;


	LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiConnect(%x %x %x). %d %x:%u (%s)\n", (ULONG) PsGetCurrentProcessId(), pIrp, pIrpStack, pCompletion, pTransportAddress->TAAddressCount, ntohl(ip->in_addr), ntohs(ip->sin_port), inet_ntoa2(ip->in_addr)));


	inet_ntoa(ip->in_addr, NETWORKNAME);

	if (LearningMode == FALSE)
	{
		POLICY_CHECK_OPTYPE_NAME(NETWORK, DeviceType == NET_DEVICE_TYPE_TCP ? OP_TCPCONNECT : OP_UDPCONNECT);
	}
	else
	{
		// learning mode
		AddRule(RULE_NETWORK, NETWORKNAME, DeviceType == NET_DEVICE_TYPE_TCP ? OP_TCPCONNECT : OP_UDPCONNECT);
	}


	HOOK_ROUTINE_EXIT(STATUS_SUCCESS);
}



NTSTATUS
TdiListen(IN PIRP pIrp, IN PIO_STACK_LOCATION pIrpStack, OUT PTDI_CALLBACK pCompletion, IN ULONG DeviceType)
{
	/*
	 * IrpSp->Parameters 
	 *
	 * Pointer to a TDI_REQUEST_KERNEL_LISTEN structure, equivalent to the TDI_REQUEST_KERNEL structure.
	 */

	PTDI_REQUEST_KERNEL_LISTEN		ListenInfo = (PTDI_REQUEST_KERNEL_LISTEN) &pIrpStack->Parameters;
	PTRANSPORT_ADDRESS				pTransportAddress;
	PTA_ADDRESS						pAddress;
	PTDI_ADDRESS_IP					ip;


	if (! MmIsAddressValid(ListenInfo) ||
		! MmIsAddressValid(ListenInfo->RequestConnectionInformation) ||
		! MmIsAddressValid(ListenInfo->RequestConnectionInformation->RemoteAddress))
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("TdiListen: MmIsAddressValid failed\n"));
		HOOK_ROUTINE_EXIT(STATUS_SUCCESS);
	}

	pTransportAddress = (PTRANSPORT_ADDRESS) ListenInfo->RequestConnectionInformation->RemoteAddress;

	pAddress = (PTA_ADDRESS) pTransportAddress->Address;

	/* verify that the specified address is a single IP address */
	if (pTransportAddress->TAAddressCount != 1 ||
		pAddress->AddressType != TDI_ADDRESS_TYPE_IP ||
		pAddress->AddressLength != TDI_ADDRESS_LENGTH_IP)        
	{
		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d TdiListen: Invalid address detected\n", CURRENT_PROCESS_PID));
		HOOK_ROUTINE_EXIT(STATUS_SUCCESS);
	}

	ip = (PTDI_ADDRESS_IP) &pAddress->Address;


	LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("TdiListen(%x %x %x). %d %x:%u (%s)\n", pIrp, pIrpStack, pCompletion, pTransportAddress->TAAddressCount, ntohl(ip->in_addr), ntohs(ip->sin_port), inet_ntoa2(ip->in_addr)));


	return STATUS_SUCCESS;
}



NTSTATUS
TdiAccept(IN PIRP pIrp, IN PIO_STACK_LOCATION pIrpStack, OUT PTDI_CALLBACK pCompletion, IN ULONG DeviceType)
{
	/*
	 * IrpSp->Parameters 
	 *
	 * Specifies a TDI_REQUEST_KERNEL_ACCEPT structure.
	 */

	PTDI_REQUEST_KERNEL_ACCEPT		AcceptInfo = (PTDI_REQUEST_KERNEL_ACCEPT) &pIrpStack->Parameters;
	PTRANSPORT_ADDRESS				pTransportAddress = (PTRANSPORT_ADDRESS) AcceptInfo->RequestConnectionInformation->RemoteAddress;
	PTA_ADDRESS						pAddress = (PTA_ADDRESS) pTransportAddress->Address;
	PTDI_ADDRESS_IP					ip = (PTDI_ADDRESS_IP) &pAddress->Address;


	LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("TdiAccept(%x %x %x). %d %x:%u (%s)\n", pIrp, pIrpStack, pCompletion, pTransportAddress->TAAddressCount, ntohl(ip->in_addr), ntohs(ip->sin_port), inet_ntoa2(ip->in_addr)));


	return STATUS_SUCCESS;
}


/*
NTSTATUS
GenericCompletion(IN PDEVICE_OBJECT pDevObj, IN PIRP pIrp, IN PVOID pContext)
{             
	if (pIrp->PendingReturned )
		IoMarkIrpPending(pIrp);
    
	return STATUS_SUCCESS;
}
*/

			   
TDI_IOCTL	TdiIoctl[] =
{
	{ TDI_ASSOCIATE_ADDRESS, "TDI_ASSOCIATE_ADDRESS", TdiStub },
	{ TDI_DISASSOCIATE_ADDRESS, "TDI_DISASSOCIATE_ADDRESS", TdiStub },
	{ TDI_CONNECT, "TDI_CONNECT", TdiConnect },
	{ TDI_LISTEN, "TDI_LISTEN", TdiListen },
	{ TDI_ACCEPT, "TDI_ACCEPT", TdiAccept },
	{ TDI_DISCONNECT, "TDI_DISCONNECT", TdiStub },
	{ TDI_SEND, "TDI_SEND", TdiStub },
	{ TDI_RECEIVE, "TDI_RECEIVE", TdiStub },
	{ TDI_SEND_DATAGRAM, "TDI_SEND_DATAGRAM", TdiStub },
	{ TDI_RECEIVE_DATAGRAM, "TDI_RECEIVE_DATAGRAM", TdiStub },
	{ TDI_SET_EVENT_HANDLER, "TDI_SET_EVENT_HANDLER", TdiSetEventHandler },
	{ TDI_QUERY_INFORMATION, "TDI_QUERY_INFORMATION", TdiStub },
	{ TDI_SET_INFORMATION, "TDI_SET_INFORMATION", TdiStub },
	{ TDI_ACTION, "TDI_ACTION", TdiStub },
	{ TDI_DIRECT_SEND, "TDI_DIRECT_SEND", TdiStub },
	{ TDI_DIRECT_SEND_DATAGRAM, "TDI_DIRECT_SEND_DATAGRAM", TdiStub },
};



//XXX this function can be called from HookedNtCreateFile (-> NtCreateFile -> IoCreateFile -> ObOpenObjectbyName -> ... -> TDI)
BOOLEAN
TDIDispatch(PDEVICE_OBJECT pDeviceObject, PIRP pIrp, NTSTATUS *status)
{
	PIO_STACK_LOCATION	pIrpStack;
	TDI_CALLBACK		Callback;
	ULONG				DeviceType = 0;	


	if (pDeviceObject == pTcpDevice)
	{
		DeviceType = NET_DEVICE_TYPE_TCP;
	}
	else if (pDeviceObject == pUdpDevice)
	{
		DeviceType = NET_DEVICE_TYPE_UDP;
	}
	else if (pDeviceObject == pIpDevice)
	{
		DeviceType = NET_DEVICE_TYPE_IP;
	}
	else
	{
		return FALSE;
	}


	HOOK_TDI_ENTER_NORC();


	pIrpStack = IoGetCurrentIrpStackLocation(pIrp);


	memset(&Callback, 0, sizeof(Callback));

//	if (pIrpStack->Parameters.DeviceIoControl.IoControlCode == IOCTL_TDI_QUERY_DIRECT_SEND_HANDLER)
//	{
//		LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("IOCTL_TDI_QUERY_DIRECT_SEND_HANDLER\n"));
//	}

	switch (pIrpStack->MajorFunction)
	{
		case IRP_MJ_CREATE:

			*status = TDICreate(pDeviceObject, pIrp, pIrpStack, &Callback);

			break;


		case IRP_MJ_DEVICE_CONTROL:

//			if (DeviceType == NET_DEVICE_TYPE_IP && pIrpStack->Parameters.DeviceIoControl.IoControlCode == 0x120000)
			if (DeviceType == NET_DEVICE_TYPE_IP)
			{
				LOG(LOG_SS_NETWORK, LOG_PRIORITY_VERBOSE, ("%d pIpDevice in use (%x %x %x)\n", (ULONG) PsGetCurrentProcessId(), pIrpStack->Parameters.DeviceIoControl.IoControlCode, pIrpStack->MajorFunction, pIrpStack->MinorFunction));
//				*status = STATUS_ACCESS_DENIED;
				break;
			}

			if (KeGetCurrentIrql() != PASSIVE_LEVEL || ! NT_SUCCESS(TdiMapUserRequest(pDeviceObject, pIrp, pIrpStack)))
			{
				LOG(LOG_SS_NETWORK, LOG_PRIORITY_VERBOSE, ("TdiMapUserRequest failed: %x (irql %d)\n", pIrpStack->Parameters.DeviceIoControl.IoControlCode, KeGetCurrentIrql()));
				break;
			}

			LOG(LOG_SS_NETWORK, LOG_PRIORITY_VERBOSE, ("IRP_MJ_DEVICE_CONTROL2 %x\n", pIrpStack->Parameters.DeviceIoControl.IoControlCode));

			/* FALLTHROUGH */


		case IRP_MJ_INTERNAL_DEVICE_CONTROL:
		{
			int		i;

			if (DeviceType == NET_DEVICE_TYPE_IP)
				LOG(LOG_SS_NETWORK, LOG_PRIORITY_DEBUG, ("%d pIpDevice in use2\n", (ULONG) PsGetCurrentProcessId()));

			for (i = 0; i < sizeof(TdiIoctl) / sizeof(TdiIoctl[0]); i++)
			{
				if (TdiIoctl[i].MinorFunction == pIrpStack->MinorFunction)
				{
					if (TdiIoctl[i].pfRoutine == TdiStub)
						LOG(LOG_SS_NETWORK, LOG_PRIORITY_VERBOSE, ("%d IRP_MJ_INTERNAL_DEVICE_CONTROL %s\n", (ULONG) PsGetCurrentProcessId(), TdiIoctl[i].Description));

					*status = TdiIoctl[i].pfRoutine(pIrp, pIrpStack, &Callback, DeviceType);

					break;