
/**
 * @file: general_handler.c
 *
 * Tento soubor implementuje primitiva umoznujici monitorovat ci menit
 * prubeh systemovych volani. Vyuziva pritom funkci modulu ssdt.c. 
 */

#include <ntddk.h>
#include "driver.h"
#include "Cr0Trick.h"
#include "ssdt.h"
#include "general_handler.h"



static VOID GeneralHandlerCodeDestroy(PVOID CodeAddress);
static NTSTATUS GeneralHandlerCodeCreate(OUT PVOID *CodeAddress, IN PRE_HANDLER PreHnalder, IN PVOID Original, IN POST_HANDLER PostHandler, ULONG Index, ULONG ArgCnt);

// Pocet sluzeb v tabulce systemovych volani
static ULONG NumberOfServices = 0;

static KSPIN_LOCK Lock;

// Delka obecne obsluhy systemoveho volani
static ULONG Length = 0;



/** Obecna obsluha systemoveho volani.
 *
 * Tato metarutina predpoklada, ze adresa nektere jeji instance bude
 * vlozena do tabulky systemovych volani. Tato rutina nejprve zavola
 * tzv. Prehandler - funkci, ktere preda informace o cisle systemoveho
 * volani a jeho parametrech. Pokud Prehandler vrati STATUS_SUCCESS,
 * rutina dale zavola kod originalni systemove sluzby (parametry mohou
 * byt upravene v rutine Prehandler). Dale vola tzv. Postahandler - rutinu,
 * ktera obdrzi informace o cisle systemoveho volani, jeho parametrech a
 * navratove hodnote. Posthandler muze navratovou hodnotu zmenit. Pokud
 * rutina Prehandler nevrati STATUS_SUCCESS, obecna obsluha systemoveho
 * volani nevola originalni kod systemove sluzby, ale metodu Posthandler.
 */
__declspec(naked) VOID GeneralServiceRoutine(VOID)
{
  __asm {
    push esi
    mov esi, esp
    sub esp, 4
    pushad

    lea edi, [esi + 8]
    push 0xAAAAAAAA            // Pocet argumentu
    push edi                   // Adresa argumentu
    push 0xBBBBBBBB            // Index sluzby
    __emit 0xe8                // Volani prehandleru
    __emit 0xCC
    __emit 0xCC
    __emit 0xCC
    __emit 0xCC
    test eax, eax              // Pokud nevrati STATUS_SUCCESS, volame hned posthandler
    jnz post
    popad

    pushad
    mov [esi-0x4], ecx         // ulozeni hodnoty ECX
    mov ecx, 0xAAAAAAAA        // pocet argumentu
    cmp ecx, 0                  
    jz orig_call
  push_par:                    // Ukladani argumentu na zasobnik
    push [4*ecx+esi+4]
    loop push_par
  orig_call:
    mov ecx, [esi-0x4]         // Nacteni puvodni hodnoty ECX
    __emit 0xe8                // Volani originalni funkce
    __emit 0xDD
    __emit 0xDD
    __emit 0xDD
    __emit 0xDD
    add esp, 0x20
    pushad

  post:
    lea edi, [esi + 0x8]
    push eax                   // Vysledek volani orig. funkce (nebo prehandleru)   
    push 0xAAAAAAAA            // Pocet argumentu
    push edi                   // Adresa argumentu
    push 0xBBBBBBBB            // Index sluzby
    __emit 0xe8                // Volani posthandleru
    __emit 0xEE
    __emit 0xEE
    __emit 0xEE
    __emit 0xEE
    mov [esi-4], eax

    popad
    add esp, 4
    mov eax, [esi-4]
    pop esi
    ret 0xFFFF
    __emit 0xFF
    __emit 0xFF
  }
}


/** Vyplni do objektu obecne obsluhy systemove sluzby udaje o jeho cilove
 *  systemove sluzbe.
 *
 * @param GH Adresa struktury GENERAL_HANDLER, ktera reprezentuje objekt
 * obecne obsluhy systemove sluzby.
 * @param ServiceIndex Cislo systemove sluzby, kterou bude objekt obecne
 * obsluhy v budoucnu obsluhovat.
 */
static VOID _GeneralHandlerFillServiceArguments(IN PGENERAL_HANDLER GH, IN ULONG ServiceIndex)
{
  KdPrint(("general_handler.c: _GeneralHandlerFillServiceArguments(GH=0x%p; ServiceIndex=%u)\n", GH, ServiceIndex));

  GH->ServiceIndex = ServiceIndex;
  GH->ArgumentCount = GetNumberOfArguments(ServiceIndex);
  GH->Original = GetServiceAddress(ServiceIndex);

  KdPrint(("general_handler.c: _GeneralHandlerFillServiceArguments(-)\n"));
  return;
}


/** Vytvori objekt obecne obsluhy systemoveho volani a aplikuje jej na
 *  zadanou systemovou sluzbu.
 *
 *  Rutina alokuje strukturu GENERAl_HANDLER, vyplni ji potrebnymi udajemi,
 *  zajisti vytvoreni instance kodu obecne obsluhy systemoveho volani a provede
 *  zmenu adresy prislusne systemove sluzby na sebe. Tim dojde k tomu, ze
 *  veskere pokusy o systemove volani s danym cislem povedou na vykonani
 *  rutin spjatych s timto objektem obecne obsluhy systemoveho volani.
 *
 *  @param OH Adresa promenne, ktera v pripade uspechu obdrzi adresu struktury
 *  GENERAL_HANDLER noveho objektu obecne obsluhy systemoveho volani.
 *  @param ServiceIndex Cislo systemove sluzby, kterou si chce volajici
 *  podmanit pomoci objektu obecne obsluhy systemoveho volani.
 *  @param PreHandler Adresa rutiny Prehandler (viz komentar u obecne obsluzne
 *  metarutiny systemoveho volani).
 *  @param PostHandler Adresa rutiny PostHandler (viz komentar u obecne obsluzne
 *  metarutiny systemoveho volani).
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech operace.
 */
NTSTATUS GeneralHandlerCreate(OUT PGENERAL_HANDLER *GH, IN ULONG ServiceIndex, IN PRE_HANDLER PreHandler, IN POST_HANDLER PostHandler)
{
  KIRQL OldIrql;
  PGENERAL_HANDLER o = NULL;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("general_handler.c: GeneralHandlerCreate(GH=0x%p; ServiceIndex=%u; PreHandler=0x%p; PostHandler=0x%p)\n", GH, ServiceIndex, PreHandler, PostHandler));

  if (GH != NULL) {
    if (ServiceIndex < NumberOfServices) {
      o = (PGENERAL_HANDLER)ExAllocatePoolWithTag(NonPagedPool, sizeof(GENERAL_HANDLER), POOL_TAG);
      if (o != NULL) {
        RtlZeroMemory(o, sizeof(GENERAL_HANDLER));
        o->PreHandler = PreHandler;
        o->PostHandler = PostHandler;
        KeAcquireSpinLock(&Lock, &OldIrql);
        _GeneralHandlerFillServiceArguments(o, ServiceIndex);
        Status = GeneralHandlerCodeCreate(&o->CodeAddress, o->PreHandler, o->Original, o->PostHandler, o->ServiceIndex, o->ArgumentCount);
        if (NT_SUCCESS(Status)) {
          SetServiceAddress(ServiceIndex, o->CodeAddress);
          *GH = o;
        } else ExFreePoolWithTag(o, POOL_TAG);

        KeReleaseSpinLock(&Lock, OldIrql);
      } else Status = STATUS_INSUFFICIENT_RESOURCES;
    } else Status = STATUS_INVALID_PARAMETER;
  } else Status = STATUS_INVALID_PARAMETER;

  KdPrint(("general_handler.c: GeneralHandlerCreate(-):0x%x\n", Status));
  return Status;
}


/** Zrusi objekt obecne obsluhy systemoveho volani.
 *
 *  Rutina modifikuje adresu systemove sluzby spjate s objektem na jeji
 *  puvodni hodnotu, odstrani instanci metarutiny a dealokuje strukturu
 *  GENERAL_HANDLER reprezentujici dany objekt obecne obsluhy systemoveho
 *  volani.
 *
 *  @param GH Adresa struktury GENERAL_HANDLER reprezentujici objekt obecne
 *  obsluhy systemoveho volani, ktery chce volajici zrusit.
 */
VOID GeneralHandlerDestroy(PGENERAL_HANDLER GH)
{
  KIRQL OldIrql;
  KdPrint(("general_handler.c: GeneralHandlerDestroy(GH=0x%p)\n", GH));

  KeAcquireSpinLock(&Lock, &OldIrql);
  SetServiceAddress(GH->ServiceIndex, GH->Original);
  KeReleaseSpinLock(&Lock, OldIrql);
  GeneralHandlerCodeDestroy(GH->CodeAddress);
  ExFreePoolWithTag(GH, POOL_TAG);

  KdPrint(("general_handler.c: GeneralHandlerDestroy(-)\n"));
  return;
}


/** Vytvori instanci kodu metarutiny pro obsluhu systemoveho volani se
 *  zadanymi parametry.
 *
 *  Rutina prekopiruje kod metarutiny do nove alokovane oblasti pameti a
 *  zmeni casti kodu nasledujicim zpusobem:
 *  - 0xAAAAAAAA za pocet argumentu systemove sluzby,
 *  - 0xBBBBBBBB za cislo systemove sluzby,
 *  - 0xCCCCCCCC za relativni adresu rutiny Prehandler,
 *  - 0xDDDDDDDD za relativni adresu originalniho kodu systemove sluzby,
 *  - 0xEEEEEEEE za relativni adresu rutiny Posthandler,
 *  - 0xFFFFFFFF za pocet bajtu, ktery zaberou argumenty systemove sluzby na zasobniku.
 *
 *  @param CodeAddress Adresa promenne, ktera v pripade uspechu obdrzi
 *  adresu nove vytvorene instance metarutiny pro obecnou obsluhu systemoveho
 *  volani.
 *  @param PreHandler Adresa rutiny Prehandler.
 *  @param Original Adresa originalniho kodu systemove sluzby.
 *  @param PostHandler Adresa rutiny Posthandler.
 *  @param Index Cislo systemove sluzby.
 *  @param ArgCnt Pocet bajtu, ktere zaberou argumenty systemove sluzby na
 *  zasobniku.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech
 *  operace.
 */
NTSTATUS GeneralHandlerCodeCreate(OUT PVOID *CodeAddress, IN PRE_HANDLER PreHandler, IN PVOID Original, IN POST_HANDLER PostHandler, ULONG Index, ULONG ArgCnt)
{
  PULONG Address = NULL;
  PVOID StopMark = NULL;
  ULONG NumberReplaces = 0;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("general_handler.c: GeneralHandlerCodeCreate(CodeAddress=0x%p; PreHandler=0x%p; Original=0x%p; PostHandler=0x%p; Index=%u; ArgCnt=%u)\n", CodeAddress, PreHandler, Original, PostHandler, Index, ArgCnt));

  if (CodeAddress != NULL) {
    *CodeAddress = NULL;
    Address = (PULONG)ExAllocatePoolWithTag(NonPagedPool, Length, POOL_TAG);
    if (Address != NULL) {
      *CodeAddress = Address;
      RtlCopyMemory(Address, &GeneralServiceRoutine, Length);
      StopMark = (PCHAR)Address + Length;
      Status = STATUS_UNSUCCESSFUL;
      CR0TrickBegin();
      while (Address != StopMark) {
        switch (*Address) {
          case 0xAAAAAAAA:
            *Address = (ArgCnt >> 2);
            NumberReplaces++;
            break;
          case 0xBBBBBBBB:
            *Address = Index;
            NumberReplaces++;
            break;
          case 0xCCCCCCCC:
            *Address = (ULONG)PreHandler - ((ULONG)Address - 1 + 5);
            NumberReplaces++;
            break;
          case 0xDDDDDDDD:
            *Address = (ULONG)Original - ((ULONG)Address - 1 + 5);
            NumberReplaces++;
            break;
          case 0xEEEEEEEE:
            *Address = (ULONG)PostHandler - (ULONG)((ULONG)Address - 1 + 5);
            NumberReplaces++;
            break;
          case 0xFFFFFFFF:
            *Address = ArgCnt;
            if (NumberReplaces == 8)
              Status = STATUS_SUCCESS;

            break;
        }

        Address = (PULONG)((PCHAR)Address + 1);
      }

      CR0TrickEnd();
      if (!NT_SUCCESS(Status)) {
        ExFreePoolWithTag((PCHAR)StopMark - Length, POOL_TAG);
        *CodeAddress = NULL;
      }
    } else Status = STATUS_INSUFFICIENT_RESOURCES;
  } else Status = STATUS_INVALID_PARAMETER;

  KdPrint(("general_handler.c: GeneralHandlerCodeCreate(-):0x%x\n", Status));
  return Status;
}


/** Zrusi instanci metarutiny pro obecnou obsluhu systemoveho volani.
 *
 *  Rutina pouze uvolni pamet, kterou kod teto instance zabira.
 *
 *  @param CodeAddress Adresa kodu prislusne instance metarutiny.
 */
VOID GeneralHandlerCodeDestroy(PVOID CodeAddress)
{
  KdPrint(("general_handler.c: GeneralHandlerCodeDestroy(CodeAddress=0x%p)\n", CodeAddress));

  ExFreePoolWithTag(CodeAddress, POOL_TAG);

  KdPrint(("general_handler.c: GeneralHandlerCodeDestroy(-)\n"));
  return;
}


/** Inicializuje modul pro snadne monitorovani a zmenu chovani systemovych
  * sluzeb.
  *
  *  Rutina zjisti delku kodu metarutiny pro obecnou obsluhu systemoveho
  *  volani.
  *
  *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech
  *  operace.
 */
NTSTATUS GeneralHandlerInit(VOID)
{
  ULONG i = 0;
  PULONG Address = (PULONG)&GeneralServiceRoutine;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("general_handler.c: GeneralHandlerInit()\n"));

  NumberOfServices = GetNumberOfServices();
  KeInitializeSpinLock(&Lock);
  Length = 0;
  while ((ULONG)Address - (ULONG)&GeneralServiceRoutine < 4096) {
    if (*Address == 0xFFFFFFFF) {
      Length = (ULONG)Address - (ULONG)&GeneralServiceRoutine + 4;
      Status = STATUS_SUCCESS;
      break;
    }

    Address = (PULONG)((PCHAR)Address + 1);
  }

  KdPrint(("general_handler.c: GeneralHandlerInit(-):0x%x\n", Status));
  return Status;
}


/** Provede uklid po modulu pro monitorovani a zmenu chovani systemovych
 *  sluzeb.
 *
 *  Protoze tento modul pri sve inicializaci nijak nezasahuje do internich
 *  datovych struktur jadra, rutina nedela vlastne nic.
 */
VOID GeneralHandlerFinit(VOID)
{
  KdPrint(("general_handler.c: GeneralHandlerFinit()\n"));

  KdPrint(("general_handler.c: GeneralHandlerFinit(-)\n"));
  return;
}
