
/**
 * @file: ntregDrivers.c
 *
 * Tento modul obsahuje funkce, ktere umoznuji pracovat s ovladaci jadra.
 * Dovoluji ovladac nainstalovat, nacist do pameti jadra, uvolnit a
 * odinstalovat. Instalace a odinstalace se provadi primym zapisem do
 * registru, nacitani a uvolnovani z jadra obstaravaji nativni funkce
 * NtLoadDriver a NtUnloadDriver.
 */

#include <stdio.h>
#include <windows.h>
#include <winsvc.h>
#include "..\shared\common.h"
#include "ntregDrivers.h"


// Uchovava adresu rutiny NtLoadDriver
static NTLOADDRIVER _NtLoadDriver = NULL;

// Uchovava adresu rutiny NtUnloadDriver
static NTUNLOADDRIVER _NtUnloadDriver = NULL;

// Uchovava otevrene handle na klic, kde jsou ulozeny ve forme podklicu
// informace o kazde sluzbe (tedy i ovladaci) nainstalovane v systemu.
static   HKEY serviceskey = NULL;



/** Nainstaluje ovladac.
 *
 * Rutina zapise potrebne informace o ovladaci do registru. V podstate
 * tak provede instalaci nove sluzby, ktera bude ovladac reprezentovat.
 *
 * @param DriverName Interni jmeno noveho ovladace.
 * @param FileName Plny nazev souboru noveho ovladace.
 *
 * @return Rutina vrati TRUE v pripade uspechu a FALSE v pripade neuspechu
 * operace.
 */
BOOL ntldInstallDriver (PWCHAR DriverName, PWCHAR FileName)
{
  DWORD start = SERVICE_DEMAND_START;
  DWORD type = SERVICE_KERNEL_DRIVER;
  BOOL ret = FALSE;
  HKEY driverkey;
  LONG res = 0;
  UNICODE_STRING uFullFileName;
  DWORD FileNameSize = (DWORD) (wcslen(FileName) + 1) * sizeof(WCHAR);
  DebugMsg(("ntregDrivers.c: ntldInstallDriver(DriverName=%S; FileName=%S)\n", DriverName, FileName));

  res = RegCreateKeyExW(serviceskey, DriverName, 0, NULL, 0, KEY_ALL_ACCESS, NULL, &driverkey, NULL);
  ret = res == ERROR_SUCCESS;
  if (ret) {
    _PrepareFullName(FileName, L"\\??\\", &uFullFileName); 
	if (ret) {
		res = RegSetValueExW(driverkey, L"ImagePath", 0, REG_SZ, (PVOID)uFullFileName.Buffer, uFullFileName.Length + sizeof(WCHAR));
      ret = res == ERROR_SUCCESS;
      if (ret) {
        res = RegSetValueExW(driverkey, L"Start", 0, REG_DWORD, (PVOID)&start, sizeof(start));
        ret = res == ERROR_SUCCESS;
        if (ret) {
          res = RegSetValueExW(driverkey, L"Type", 0, REG_DWORD, (PVOID)&type, sizeof(type));
          ret = res == ERROR_SUCCESS;
          if (!ret) {
            DebugMsg(("ntregDrivers.c: ntldInstallDriver(): RegSetValueExW(Type) failed, Error %u\n", res));
            RegDeleteKeyW(serviceskey, DriverName);
          }
        } else {
         DebugMsg(("ntregDrivers.c: ntldInstallDriver(): RegSetValueExW(Start) failed, Error %u\n", res));
         RegDeleteKeyW(serviceskey, DriverName);
        }
      } else {
        DebugMsg(("ntregDrivers.c: ntldInstallDriver(): RegSetValueExW(ImagePath) failed, Error %u\n", res));
        RegDeleteKeyW(serviceskey, DriverName);
      }

      _FreeFullName(&uFullFileName);
	} else DebugMsg(("ntregDrivers.c: ntldInstallDriver(): Failed to format driver file name\n"));

    RegCloseKey(driverkey);
  } else DebugMsg(("ntregDrivers.c: ntldInstallDriver(): RegCreateKeyExW() failed. Error %u\n", res));

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


/** Nacte ovladac do jadra.
 *
 * Rutina zavola nativni funkci NtLoadDriver s odkazem do registru na
 * informace o sluzbe ovladace.
 *
 * @param DriverName Interni jmeno ovladace.
 *
 * @return Funkce vraci TRUE v pripade uspechu a FALSE v pripade neuspechu.
 */
BOOL ntldLoadDriver (PWCHAR DriverName)
{
  BOOL ret = FALSE;
  UNICODE_STRING uFullName;
  NTSTATUS status = 0xC0000001;
  DebugMsg(("ntRegDrivers.c: ntldLoadDriver(DriverName=%S)\n", DriverName));

  ret = _PrepareFullName(DriverName, NTLD_NAME_PREFIX, &uFullName);
  if (ret) {
    DebugMsg(("ntRegDriver.c: ntldLoadDriver(): Driver name is %S\n", uFullName.Buffer));
    status = _NtLoadDriver(&uFullName);
    ret = status == 0;
    if (!ret)
      DebugMsg(("ntregDrivers.c: NtLoadDriver() failed. Status 0x%x\n", status));

    _FreeFullName(&uFullName);
  } else DebugMsg(("ntregDrivers.c: ntldLoadDriver(): HeapAlloc() failed. Error %u\n", GetLastError()));

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


/** Uvolni ovladac z jadra.
 *
 * Rutina zavola nativni funkci NtUnloadDriver s odkazem do registru na
 * informace o sluzbe ovladace.
 *
 * @param DriverName Interni jmeno ovladace.
 *
 * @return Funkce vraci TRUE v pripade uspechu a FALSE v pripade neuspechu.
 */
BOOL ntldUnloadDriver (PWCHAR DriverName)
{
  BOOL ret = FALSE;
  UNICODE_STRING uFullName;
  NTSTATUS status = 0xC0000001;
  DebugMsg(("ntRegDrivers.c: ntldUnloadDriver(DriverName=%S)\n", DriverName));

  ret = _PrepareFullName(DriverName, NTLD_NAME_PREFIX, &uFullName);
  if (ret) {
	status = _NtUnloadDriver(&uFullName);
	ret = status == 0;
    if (!ret)
      DebugMsg(("ntregDrivers.c: NtUnloadDriver() failed. Status %x\n", status));
	
    _FreeFullName(&uFullName);
  }

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


/** Odinstaluje ovladac.
 *
 * Rutina smaze podklic odpovidajici sluzbe ovladace z registru.
 *
 * @param DriverName Interni jmeno ovladace.
 *
 * @return Funkce vrati TRUE v pripade uspechu a FALSE v pripade neuspechu.
 */
BOOL ntldUninstallDriver (PWCHAR DriverName)
{
  BOOL ret = FALSE;
  LONG res = 0;
  DebugMsg(("ntregDrivers.c: ntldUninstallDriver(DriverName=%S)\n", DriverName));

  ret = DeleteRegistryKey(serviceskey, DriverName);
  if (!ret)
    DebugMsg(("ntregDrivers.c: ntldUninstallDriver(): failed to delete driver key. Error %u\n", res));

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



/** Inicializuje modul pro praci s ovladaci pomoci primeho pristupu do
 *  registru a nativnich funkci NtLoadDriver a NtUnloadDriver.
 *
 *  Rutina povoli opravneni nacitat ovladace do jadra, najde adresu funkci
 *  NtLoadDriver a NtUnloadDriver v knihovne ntdll.dll a ziska handle ke
 *  klici, kde jsou ulozney udaje o vsech nainstalovanych sluzbach daneho
 *  systemu.
 *
 *  @return Funkce vraci TRUE v pripade uspechu a FALSE v pripade 
 *  neuspechu.
 */
BOOL ntldInit(VOID)
{
  LONG res = 0;
  HMODULE hntdll = NULL;
  BOOL ret = FALSE;
  DebugMsg(("ntregDrivers.c: ntldInit()\n"));

  ret = EnablePrivilege(LOAD_DRIVER_PRIVILEGE);
  if (ret) {
    hntdll = GetModuleHandleW(L"ntdll.dll");
    ret = hntdll != NULL;
    if (ret) {
      _NtLoadDriver = (NTLOADDRIVER)GetProcAddress(hntdll, "NtLoadDriver");
      _NtUnloadDriver = (NTUNLOADDRIVER)GetProcAddress(hntdll, "NtUnloadDriver");
      ret = (_NtLoadDriver != NULL && _NtUnloadDriver != NULL);
	  if (ret) {
        res = RegOpenKeyExW(HKEY_LOCAL_MACHINE, SERVICES_KEY, 0, KEY_CREATE_SUB_KEY, &serviceskey);
        ret = res == ERROR_SUCCESS;
	    if (!ret)
          DebugMsg(("ntredDrivers(): ntldInit(): Failed to open services key. Error %u\n", res));
	  } else {
        if (_NtLoadDriver == NULL)
          DebugMsg(("ntregDrivers.c: ntldInit(): Failed to find NtLoadDriver export. Error %u.\n", GetLastError()));

        if (_NtUnloadDriver == NULL)
          DebugMsg(("ntregDrivers.c: ntldInit(): Failed to find NtUnloadDriver export. Error %u.\n", GetLastError()));
 	  }
    } else DebugMsg(("ntregDrivers.c: ntldInit(): Failed to find ntdll.dl. Error %u.\n", GetLastError()));
  } else DebugMsg(("ntregDrivers.c: ntldInit(): Failed to enable SeLoadDriver privilege\n"));

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


/** Provede uklid po modulu pro praci s ovladaci pres primy pristup do
 *  registru a nativni funkce NtLoadDriver a NtUnloadDriver.
 *
 *  Rutina zrusi handle ke klici sluzeb.
 */
VOID ntldFinit(VOID)
{
  DebugMsg(("ntregDrivers.c: ntldFinit()\n"));

  RegCloseKey(serviceskey);

  DebugMsg(("ntregDrivers.c: ntldFinit(-)\n"));
  return;
}
