/*
     Programmed By ChenChengQin(tsu00@263.net)!
     Welcome to http://webcrazy.yeah.net for more information!
     History:
             10/05/2000 Initial edition!
                        Portion from internet!
             10/07/2000 Bug fixes about deadlock when new process created!
             10/22/2000 Add thread filter!
*/

extern "C" {
#include "ntddk.h"
#include "windef.h"
}
 
#define NT_DEVICE_NAME      L"\\Device\\Strace"
#define DOS_DEVICE_NAME     L"\\DosDevices\\Strace"
#define SYSNAME             "System"

#include "Stracedef.h"

STRACE_BUF_FIELD * StraceBuf1;
STRACE_BUF_FIELD * StraceBuf2;
LONG StraceBufCount1;
LONG StraceBufCount2;
LONG StraceBufMaxCount1 = 1024;
LONG StraceBufMaxCount2 = 1024;

char *keSDTParameter;
char *w32SDTParameter;

DWORD StracePID=0;
BOOLEAN SpyOn=FALSE;
ULONG ProcessNameOffset;

LONG StraceCurrentBufId; 

WORD ProcessFilterEntriesCount;
WORD *ProcessFilterIdArray;

DWORD filter[6];
STRACE_FILTER* strace_filter;

KEVENT StraceBufOverflowEvent;
FAST_MUTEX StraceBuf1Mutex;
FAST_MUTEX StraceBuf2Mutex;
FAST_MUTEX StraceCountMutex;
FAST_MUTEX StraceFilterMutex;

PVOID OldSystemServiceISR;

NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
    )
{

    PDEVICE_OBJECT deviceObject = NULL;
    NTSTATUS status;
    UNICODE_STRING lNtNameString;
    UNICODE_STRING lWin32NameString;

    RtlInitUnicodeString( &lNtNameString, NT_DEVICE_NAME );

    status = IoCreateDevice(
                 DriverObject,
                 0,                   
                 &lNtNameString,
                 FILE_DEVICE_UNKNOWN,
                 0,                     
                 FALSE,                
                 &deviceObject
                 );

    if ( NT_SUCCESS(status) )
    {

        DriverObject->MajorFunction[IRP_MJ_CREATE] = StraceOpen;
        DriverObject->MajorFunction[IRP_MJ_CLOSE] = StraceClose;
        DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = StraceDeviceControl;

        DriverObject->DriverUnload = StraceUnload;

        RtlInitUnicodeString( &lWin32NameString, DOS_DEVICE_NAME );
    
        
        status = IoCreateSymbolicLink( &lWin32NameString, &lNtNameString );


	KeInitializeEvent(
		&StraceBufOverflowEvent,
		NotificationEvent,
		FALSE
	);


	ExInitializeFastMutex(&StraceBuf1Mutex);
	ExInitializeFastMutex(&StraceBuf2Mutex);
	ExInitializeFastMutex(&StraceCountMutex);
	ExInitializeFastMutex(&StraceFilterMutex);


	StraceBuf1 = (STRACE_BUF_FIELD *) ExAllocatePool( NonPagedPool, (StraceBufMaxCount1+1)* sizeof(STRACE_BUF_FIELD) );
	StraceBuf2 = (STRACE_BUF_FIELD *) ExAllocatePool( NonPagedPool, (StraceBufMaxCount2+1)* sizeof(STRACE_BUF_FIELD) );

	StraceBufCount1 =0;
	StraceBufCount2 =0;

	StraceCurrentBufId = STRACE_FIRST_BUF ;

	ProcessFilterEntriesCount =0;
	ProcessFilterIdArray =0;
        ProcessNameOffset=GetProcessNameOffset();


	StraceSetSystemServiceHook(TRUE);

    }
    else
    {
        DbgPrint ("Strace: Create the device failure!\n") ;
    }
    return status;
}

NTSTATUS
StraceOpen(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;

    IoCompleteRequest( Irp, IO_NO_INCREMENT );

    return STATUS_SUCCESS;
}

NTSTATUS
StraceClose(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;

    IoCompleteRequest( Irp, IO_NO_INCREMENT );

    return STATUS_SUCCESS;
}

NTSTATUS
StraceDeviceControl(
    PDEVICE_OBJECT DeviceObject,
    PIRP Irp
    )

{

    PIO_STACK_LOCATION  irpStack;
    PVOID               inputBuffer;
    PVOID               outputBuffer;
    ULONG               inputBufferLength;
    ULONG               outputBufferLength;
    ULONG               ioControlCode;

    Irp->IoStatus.Status      = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;

    irpStack = IoGetCurrentIrpStackLocation (Irp);

    inputBuffer     = Irp->AssociatedIrp.SystemBuffer;
    inputBufferLength   = irpStack->Parameters.DeviceIoControl.InputBufferLength;
    outputBuffer    = Irp->AssociatedIrp.SystemBuffer;
    outputBufferLength  = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
    ioControlCode   = irpStack->Parameters.DeviceIoControl.IoControlCode;

    switch (irpStack->MajorFunction) {

    case IRP_MJ_DEVICE_CONTROL:

        if( Irp->MdlAddress ) {

            outputBuffer = MmGetSystemAddressForMdl( Irp->MdlAddress );
        }

        Irp->IoStatus.Status = StraceProcessDeviceControl(inputBuffer, inputBufferLength, 
                    outputBuffer, outputBufferLength,
                    ioControlCode,&Irp->IoStatus.Information);
        break;
    }

    IoCompleteRequest( Irp, IO_NO_INCREMENT );
    return STATUS_SUCCESS;   

}

VOID
StraceUnload(
    IN PDRIVER_OBJECT DriverObject
    )
{
       StraceSetSystemServiceHook(FALSE);

       ExFreePool(StraceBuf1);
       ExFreePool(StraceBuf2);
       if(keSDTParameter) ExFreePool(keSDTParameter);
       if(w32SDTParameter) ExFreePool(w32SDTParameter);
       if(ProcessFilterIdArray) ExFreePool(ProcessFilterIdArray);

       UNICODE_STRING lWin32NameString;

       RtlInitUnicodeString( &lWin32NameString, DOS_DEVICE_NAME );
       
       IoDeleteSymbolicLink( &lWin32NameString );

       IoDeleteDevice( DriverObject->DeviceObject );
}


void __fastcall StraceInterceptSystemCall()
{

	unsigned int CallNumber;
	int status;
        int ParameterTable;

	__asm{
		mov CallNumber,eax;
                mov ParameterTable,edx
	}

        if(!SpyOn) return;

	KIRQL OldIrql;

	if(CallNumber<strace_filter->ServiceIDFrom||CallNumber>strace_filter->ServiceIDTo)
		return;

        if(ParameterTable<0x80000000&&strace_filter->CallFrom==CALLFROMKERNEL)
             return;
        if(ParameterTable>0x80000000&&strace_filter->CallFrom==CALLFROMUSER)
             return;
	DWORD ProcessId,ThreadId,IsInFilter = 0;
        if((ProcessId=(unsigned int)PsGetCurrentProcessId())==(DWORD)StracePID) return;
        ThreadId=(unsigned int)PsGetCurrentThreadId();
        if(strace_filter->Thrd1!=0||strace_filter->Thrd2!=0||strace_filter->Thrd1!=0){
           if(strace_filter->Thrd1!=0&&ThreadId==strace_filter->Thrd1) IsInFilter=TRUE;
           if(strace_filter->Thrd2!=0&&ThreadId==strace_filter->Thrd2) IsInFilter=TRUE;
           if(strace_filter->Thrd3!=0&&ThreadId==strace_filter->Thrd3) IsInFilter=TRUE;
           if(!IsInFilter) return;
        }
        
	IsInFilter = 0;

        KeRaiseIrql(DISPATCH_LEVEL, &OldIrql);
	ExAcquireFastMutex(&StraceFilterMutex);
        if(ProcessFilterEntriesCount==0) IsInFilter=TRUE;
        else 
            for(register i=0;i < ProcessFilterEntriesCount;i++)
		{
			if( ProcessFilterIdArray[i] == ProcessId )
			{
				IsInFilter = TRUE;
				break;
			}
		}
	ExReleaseFastMutex(&StraceFilterMutex);
	if(!IsInFilter){
	    KeLowerIrql(OldIrql);
            return ;
	}

	ExAcquireFastMutex(&StraceCountMutex);
	if( StraceBufCount1 > StraceBufMaxCount1 )
	{
		StraceBufCount1 = StraceBufMaxCount1;
		StraceBufCount2 =0;
		StraceCurrentBufId = STRACE_SECOND_BUF;
		KeSetEvent(
			&StraceBufOverflowEvent,
			0,
			FALSE
			);
	}
	else if( StraceBufCount2 > StraceBufMaxCount2 )
	{
		StraceBufCount2 = StraceBufMaxCount2;
		StraceBufCount1 =0;
		StraceCurrentBufId = STRACE_FIRST_BUF;
		KeSetEvent(
			&StraceBufOverflowEvent,
			0,
			FALSE
			);
			
	}
	ExReleaseFastMutex(&StraceCountMutex);
	if(StraceCurrentBufId == STRACE_FIRST_BUF){
		ExAcquireFastMutex(&StraceBuf1Mutex);
		StraceBuf1[StraceBufCount1].CallNumber = (WORD)CallNumber;
		KeQuerySystemTime(&(StraceBuf1[StraceBufCount1].Time));
		StraceBuf1[StraceBufCount1].ProcessId = (WORD)ProcessId;
		StraceBuf1[StraceBufCount1].ThreadId = (WORD)ThreadId;
                StraceBuf1[StraceBufCount1].ParameterTable = ParameterTable;
                GetProcessName(StraceBuf1[StraceBufCount1].ProcessName);

                if(CallNumber<SERVICEIDUSERFROM)
	  	   StraceBuf1[StraceBufCount1].ParameterNumbers=(*(char *)((int)keSDTParameter+CallNumber))/4;
                else 
		   StraceBuf1[StraceBufCount1].ParameterNumbers=(*(char *)((int)w32SDTParameter+CallNumber-SERVICEIDUSERFROM))/4;

                for(CallNumber=0;CallNumber<StraceBuf1[StraceBufCount1].ParameterNumbers;CallNumber++)
                   StraceBuf1[StraceBufCount1].Parameter[CallNumber]=*(DWORD *)(ParameterTable+CallNumber*sizeof(DWORD));