/*
InterruptHook
Copyright (C) 2003  Alexander M.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
*/

#include <ntddk.h>
#include "dispatch.h"
#include "init.h"
#include "hook.h"
#include "iocontrol.h"
#include "handler.h"
#include "debug.h"

NTSTATUS
IoReadWrite( 
		IN	PDEVICE_OBJECT		pDeviceObject, 
		IN	PIRP				pIrp )
{
	NTSTATUS				iStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION		pStack;
	ULONG					iTransfered = 0;

	pStack = IoGetCurrentIrpStackLocation( pIrp );
	
	pIrp->IoStatus.Status		= iStatus;
	pIrp->IoStatus.Information	= iTransfered;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return iStatus;
}

NTSTATUS
IoDeviceControl( 
		IN	PDEVICE_OBJECT		pDeviceObject, 
		IN	PIRP				pIrp )
{
	NTSTATUS				iStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION		pStack;
	ULONG					iTransfered = 0;
	IDT						Idt;
	INT_VECTOR				Vec;

	pStack = IoGetCurrentIrpStackLocation( pIrp );
	
	switch( pStack->Parameters.DeviceIoControl.IoControlCode )
	{
		case IOCTL_HOOK_INT:
			if( pStack->Parameters.DeviceIoControl.InputBufferLength != sizeof(UCHAR) )
			{
				iStatus = STATUS_INVALID_PARAMETER;
				iTransfered = sizeof(UCHAR);
				break;
			}

			DPRINT( "IoDeviceControl: Hooking Int 0x%.2X\n", *(UCHAR *)pIrp->AssociatedIrp.SystemBuffer );

			LoadIDT( &Idt );
			LoadINTVector( 
				&Idt, 
				*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
				&Vec );
			DWORD_TO_VEC_OFFSET( Vec, InternalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer] );
			SaveINTVector( 
				&Idt, 
				*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
				&Vec );
			
			DPRINT( "IoDeviceControl: Hooking successful\n" );
			break;

		case IOCTL_UNHOOK_INT:
			if( pStack->Parameters.DeviceIoControl.InputBufferLength != sizeof(UCHAR) )
			{
				iStatus = STATUS_INVALID_PARAMETER;
				iTransfered = sizeof(UCHAR);
				break;
			}

			DPRINT( "IoDeviceControl: Unhooking Int 0x%.2X\n", *(UCHAR *)pIrp->AssociatedIrp.SystemBuffer );
			
			LoadIDT( &Idt );
			LoadINTVector( 
				&Idt, 
				*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
				&Vec );
			Vec.wHighOffset = OriginalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer].wHighOffset;
			Vec.wLowOffset = OriginalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer].wLowOffset;
			SaveINTVector( 
				&Idt, 
				*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
				&Vec );
			
			DPRINT( "IoDeviceControl: Unhooking successful\n" );
			break;

		case IOCTL_HOOK_ALL_INT:
			if( StartTime.QuadPart != 0 )
				break;

			DPRINT( "IoDeviceControl: Hooking all interrupts\n" );
			HookAllInterrupts();
			DPRINT( "IoDeviceControl: Hooking successful\n" );
			break;

		case IOCTL_DUMP_IDT:
			if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(INT_VECTOR) * 256 )
			{
				iStatus = STATUS_BUFFER_TOO_SMALL;
				iTransfered = sizeof(INT_VECTOR) * 256;
				break;
			}

			memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&OriginalHandlers, sizeof(INT_VECTOR) * 256 );
			iTransfered = sizeof(INT_VECTOR) * 256;
			break;

		case IOCTL_GET_INT_COUNTS:
			if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(__int64) * 256 )
			{
				iStatus = STATUS_BUFFER_TOO_SMALL;
				iTransfered = sizeof(__int64) * 256;
				break;
			}

			DPRINT( "IoDeviceControl: Retrieving interrupt counts\n" );

			__asm{ PUSHFD };
			__asm{ CLI };
			memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&InterruptCounts, sizeof(__int64) * 256 );
			__asm{ POPFD };
			iTransfered = sizeof(__int64) * 256;
			break;

		case IOCTL_GET_START_TIME:
			if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(LARGE_INTEGER) )
			{
				iStatus = STATUS_BUFFER_TOO_SMALL;
				iTransfered = sizeof(LARGE_INTEGER);
				break;
			}

			memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&StartTime, sizeof(LARGE_INTEGER) );
			iTransfered = sizeof(LARGE_INTEGER);
			break;

		default:
			iStatus = STATUS_INVALID_PARAMETER;
			break;
	}

	pIrp->IoStatus.Status		= iStatus;
	pIrp->IoStatus.Information	= iTransfered;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return iStatus;
}

NTSTATUS
IoDispatch( 
		IN	PDEVICE_OBJECT		pDeviceObject, 
		IN	PIRP				pIrp )
{
	NTSTATUS				iStatus = STATUS_SUCCESS;

	pIrp->IoStatus.Status		= iStatus;
	pIrp->IoStatus.Information	= 0;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return iStatus;
}

VOID
Unload( 
		IN	PDRIVER_OBJECT		pDriverObject )
{
	UNICODE_STRING		SymlinkName;

	RtlInitUnicodeString( &SymlinkName, L"\\DosDevices\\InterruptHook" );
	IoDeleteSymbolicLink( &SymlinkName );

	if( gpCtrlDevice )
		IoDeleteDevice( gpCtrlDevice );

	RestoreNtVectors();
}