
/**
 * @file: common.c
 *
 * Tento soubor implementuje ruzne pomocne rutiny, aby ostatni soucasti
 * projektu nemusely sahat primo na rozhrani Windows API. Usnadneni
 * se tyka povolovani opravneni ci alokaci a dealokaci pameti na halde.
 */

#include <stdio.h>
#include <windows.h>
#include <Strsafe.h>
#include "common.h"

// Uchovava adresu funkce RtlInitUnicodeString
static RTLINITUNICODESTRING _RtlInitUnicodeString = NULL;



/** Prevede retezce DriverName a Prefix do tvaru Prefix\DriverName a vysledny
 *  retezec vrati reprezentovany strukturou UNICODE_STRING.
 *
 * @param DriverName Retezec, ktery bude umisten za zpetnym lomitkem.
 * @param Prefix Retezec, ktery bude umisten pred zpetnym lomitkem.
 * @param uFullName Adresa struktury UNICODE_STRING, ktera bude v pripade
 * uspechu popisovat vysledny retezec "Prefix\DriverName".
 *
 * @return Funkce vraci TRUE v pripade uspechu. Jinak vraci FALSE.
 */
BOOL _PrepareFullName(IN PWCHAR DriverName, IN PWCHAR Prefix, OUT PUNICODE_STRING uFullName)
{
  BOOL ret = FALSE;
  ULONG FullNameLength = 0;
  PWCHAR FullName = NULL;
  DebugMsg(("common.c: _PrepareFullName(DriverName=%S; Prefix=%S; uFullName=%p)\n", DriverName, Prefix, uFullName));

  FullNameLength = (ULONG)(wcslen(Prefix) + wcslen(DriverName) + 3);
  FullName = (PWCHAR)AllocMem(FullNameLength * sizeof(WCHAR));
  ret = FullName != NULL;
  if (ret) {
    StringCbPrintfW(FullName, FullNameLength * sizeof(WCHAR), L"%s\\%s", Prefix, DriverName);
	_RtlInitUnicodeString(uFullName, FullName);
  } else DebugMsg(("common.c: _PrepareFullName(): HeapAlloc failed. Error %u\n", GetLastError()));

  DebugMsg(("common.c: _PrepraeFullName(-):%u\n", ret));
  return ret;
}


/** Uvolni pamet, kterou zabira retezec vytvoreny pomoci rutiny
 *  _PrepareFullName.
 *
 * Rutina uvolni pamet, na kterou odkazuje prvek Buffer struktury
 * UNICODE_STRING. Tuto strukturu v zapeti vynuluje.
 *
 * @param uFullName. Adresa struktury UNICODE_STRING, jejiz prostredky
 * maji byt uvolneny.
 */
VOID _FreeFullName(IN PUNICODE_STRING uFullName)
{
  DebugMsg(("common.c: _FreeFullName(uFullName=%p)\n", uFullName->Buffer));

  FreeMem(uFullName->Buffer);
  ZeroMemory(uFullName, sizeof(UNICODE_STRING));

  DebugMsg(("common.c: _FreeFullName(-)\n"));
  return;
}


/** Povoli aktualnimu procesu privilegium zadaneho jmena.
 *
 * @param PrivilegeName Nazev ciloveho privilegia.
 *
 * @return Funkce vrati TRUE v pripade uspechu. V opacnem pripade vraci
 * FALSE.
 */
BOOL EnablePrivilege(PWCHAR PrivilegeName)
{
  BOOL ret = FALSE;
  HANDLE htoken;
  LUID luid;
  TOKEN_PRIVILEGES privileges;
  DebugMsg(("common.c: EnablePrivilege(PrivilegeName=%S)\n", PrivilegeName));

  ret = OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &htoken);
  if (ret) {
    ret = LookupPrivilegeValueW(NULL, PrivilegeName, &luid);
    if (ret) {
      DWORD ret_len = 0;

	  privileges.PrivilegeCount = 1;
      privileges.Privileges[0].Luid = luid;
      privileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
      ret = AdjustTokenPrivileges(htoken, FALSE, &privileges, sizeof(TOKEN_PRIVILEGES), NULL, &ret_len);
      if (!ret)
        DebugMsg(("common.c: EnablePrivilege(): AdjustTokenPrivileges() failed. Error %u\n", ret));
    } else DebugMsg(("common.c: EnablePrivilege(): LookupPrivilegeValue() failed. Error %u\n", ret));

    CloseHandle(htoken);
  } else DebugMsg(("common.c: EnablePrivilege(): OpenProcessToken() failed. Error %u\n", GetLastError()));

  DebugMsg(("common.c: EnabledPrivilege(-):%u\n", ret));
  return ret;
}


/** Alokuje blok pameti z haldy.
 *
 * @param Size Pozadovana velikost bloku.
 *
 * @return Funkce vrati adresu nove alokovaneho bloku, nebo NULL, pokud
 * se alokace nezdarila.
 */
PVOID AllocMem(SIZE_T Size)
{
  PVOID ret = NULL;

  ret = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, Size);

  return ret;
}


/** Zkonvertuje retezec znaku ANSI na retezec Unicode.
 *
 * @param AnsiString Retezec znaku ANSi, ktery si volajici preje zkonvertovat.
 *
 * @return Funkce vraci adresu pocatku nove vytvoreneho retezce Unicode, nebo
 * NULL, pokud se konverze nepodari. Misto pro novy retezec je alokovano
 * z haldy a volajici je povinen je uvolnit, jakmile prestane s retezcem
 * pracovat.
 */
PWCHAR ConvertAnsiStringToUnicode(PCHAR AnsiString)
{
  HRESULT res = 0;
  PWCHAR ret = NULL;
  SIZE_T retlength = (strlen(AnsiString) + 1) * sizeof(WCHAR);
  DebugMsg(("common.c: ConvertAnsiStringToUnicode(AnsiString=%s)\n", AnsiString));

  ret = AllocMem(retlength);
  if (ret != NULL) {
    res = StringCbPrintfW(ret, retlength, L"%S", AnsiString);
    if (res != S_OK) {
      DebugMsg(("common.c: ConvertAnsiStringtoUnicode(): StringCbPrintfW failed with error %u\n", res));
	  FreeMem(ret);
	  ret = NULL;
    }
  } else DebugMsg(("common.c: ConvertAnsiStringToUnicode(): Failed to allocate %u bytes to store unicode string. Error %u\n", retlength, GetLastError()));

  if (ret != NULL)
    DebugMsg(("common.c: ConvertAnsiStringToUnicode():0x%p;%S\n", ret, ret));
  else DebugMsg(("common.c: ConvertAnsiStringToUnicode():0x%p\n", ret));
  return ret;
}


/** Uvolni retezec vytvoreny funkce ConvertAnsiStringToUnicodeString.
 *
 * @param UnicodeString Retezec Unicode vytvoreny funkci 
 * ConvertAnsiStringToUnicodeString.
 */
VOID FreeUnicodeString(PWCHAR UnicodeString)
{
  FreeMem(UnicodeString);

  return;
}


/** Uvolni blok pameti alokovany rutinou AllocMem.
 *
 * @param BUffer Adresa bloku pameti, ktery si volajici preje uvolnit.
 */
VOID FreeMem(PVOID Buffer)
{
  HeapFree(GetProcessHeap(), 0, Buffer);

  return;
}


/** Odstrani klic registru specifikovany pomoci handle klice a jeho podklice.
 *
 *  Funkce si umi poradit i s klicem, ktery ma dalsi podklice. Rekurzivne se
 *  vola na jednotlive podklice a maze je. Podobna rutina existuje i ve Windows API,
 *  ale je pritomna az od Windows Vista.
 *
 *  @param hKey Handle ke klici.
 *  @param aSubkey Nazev podklice, ktery si volajici preje smazat.
 *
 *  @return Funkce vraci TRUE v pripade uspechu. V pripade neuspechu
 *  vrati FALSE.
 */
BOOL DeleteRegistryKey(HKEY hKey, PWCHAR aSubkey)
{
  BOOL ret = FALSE;
  HKEY hSubkey = NULL;
  LONG res = ERROR_SUCCESS;
  PWCHAR SubkeyName = NULL;
  DWORD Index = 0;
  DWORD SubkeyNameLen = MAX_PATH * sizeof(WCHAR);
  DebugMsg(("common.c: DeleteRegistryKey(hKey=0x%p; aSubkey=%S)\n", hKey, aSubkey));

  SubkeyName = (PWCHAR)AllocMem(SubkeyNameLen);
  ret = SubkeyName != NULL;
  if (ret) {
    res = RegOpenKeyExW(hKey, aSubkey, 0, DELETE | KEY_ENUMERATE_SUB_KEYS, &hSubkey);
    ret = res == ERROR_SUCCESS;
    if (ret) {
      Index = 0;
      do {
        SubkeyNameLen = MAX_PATH * sizeof(WCHAR);
	    ZeroMemory(SubkeyName, SubkeyNameLen);
        res = RegEnumKeyExW(hSubkey, Index, SubkeyName, &SubkeyNameLen, NULL, NULL, NULL, NULL);
	    ret = res == ERROR_SUCCESS;
        if (ret) {
          Index++;
          ret = DeleteRegistryKey(hSubkey, SubkeyName);
        }

	  } while (res == ERROR_SUCCESS);

	  RegCloseKey(hSubkey);
	  ret = res == ERROR_NO_MORE_ITEMS;
	  if (ret) {
        res = RegDeleteKeyW(hKey, aSubkey);
        ret = res == ERROR_SUCCESS;
		if (!ret)
          DebugMsg(("common.c: DeleteRegistryKey(): Failed to delete subkey. Error %u\n", res));
      } else DebugMsg(("common.c: DeleteRegistryKey(): Failed to enumerate subkeys. Error %u\n", res));
    } else DebugMsg(("common.c: DeleteRegistryKey(): Failed to open subkey. Error %u\n", res));

    FreeMem(SubkeyName);
  } else DebugMsg(("common.c: DeleteRegistryKey(): Failed to allocate memory. Error %u\n", GetLastError()));
  
  DebugMsg(("common.c: DeleteRegistryKey(-):%u\n", ret));
  return ret;
}


/** Inicializuje modul podpurnych funkci.
 *
 *  Rutina pouze nalezne adresu funkce RtlInitUnicodeString.
 *
 *  @return Funkce vraci TRUE v pripade uspechu a FALSE v pripade
 *  neuspechu.
 */
BOOL CommonInit(VOID)
{
  HMODULE hntdll = NULL;
  BOOL ret = FALSE;
  DebugMsg(("common.c: CommonInit()\n"));

  hntdll = GetModuleHandleW(L"ntdll.dll");
  ret = hntdll != NULL;
  if (ret) {
    _RtlInitUnicodeString = (RTLINITUNICODESTRING) GetProcAddress(hntdll, "RtlInitUnicodeString");
	ret = _RtlInitUnicodeString != NULL;
	if (!ret)
      DebugMsg(("common.c: CommonInit(): Failed to find RtlInitUnicodeString export. Error %u\n", GetLastError()));
  } else DebugMsg(("common.c: CommonInit(): Failed to find ntdll.dll. Error %u\n", GetLastError()));

  DebugMsg(("common.c: CommonInit(-):%u\n", ret));
  return ret;
}


/** Privadi uklid po modulu podpurnych funkci.
 *
 * Jelikoz tento modul pri inicializaci nijak nezasahuje do struktur
 * operacniho systemu, tato funkce nema co uklizet, a tak nedela nic.
 */
VOID CommonFinit(VOID)
{
  DebugMsg(("common.c: CommonFInit()\n"));

  DebugMsg(("common.c: CommonFinit(-)\n"));
  return;
}
