// 只贴 ring3 部分,这也是最主要的代码,驱动只是负责把得到的原始值写到 ssdt,其他啥也不做
// 简单说一下,找 KiServiceTable 表过程, 是直接读取原始内核文件,寻找输出符号,然后结合重定位信息,谁向 KeServiceDescriptorTable 赋值,从而确定 KiServiceTable 的位置。
// 这是第一写驱动相关的代码,写的很挫,高手就不用看了。  

#include <tchar.h>
#include <Windows.h>
#include <winioctl.h>
#include "RSSDT.h"
/*
// RSSDT.h

typedef struct _tagSST_ENTRY
{
  ULONG cbSize;
  ULONG Index;
  ULONG OrgValue;    // 类型 RVA, 传给驱动的时做 Set Value
  ULONG NowValue; // 类型 RVA, 传给驱动做 get value
} SST_ENTRY, *PSST_ENTRY;
#define IOCTL_RSSDT_ENTRY_SET  CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_ANY_ACCESS)
*/

ULONG Rav2Raw(PVOID pFileData, ULONG rav)
{
  PIMAGE_DOS_HEADER pDosH;
  PIMAGE_NT_HEADERS pNtH;
  PIMAGE_SECTION_HEADER pSectH;
  ULONG Index;

  pDosH = (PIMAGE_DOS_HEADER)pFileData;
  pNtH = (PIMAGE_NT_HEADERS)((ULONG)pFileData + pDosH->e_lfanew);
  pSectH = IMAGE_FIRST_SECTION(pNtH);

  if ( rav >= 0 && rav < pSectH->VirtualAddress )
    return rav;

  for ( Index = 0; Index < pNtH->FileHeader.NumberOfSections; Index++ )
  {
    if ( rav >= pSectH->VirtualAddress && rav < pSectH->VirtualAddress + pSectH->Misc.VirtualSize )
      return rav - pSectH->VirtualAddress + pSectH->PointerToRawData;
    pSectH++;
  }

  return (ULONG)-1;
}
PVOID NTAPI 
GetImageDirEntry(
  IN ULONG ImageBase, 
  IN ULONG DirIndex,
  OUT OPTIONAL PIMAGE_NT_HEADERS* ppNtH,
  OUT OPTIONAL PIMAGE_DATA_DIRECTORY* ppDataDir
  )
{
  PIMAGE_DOS_HEADER pDosH;
  PIMAGE_NT_HEADERS pNtH;
  PIMAGE_DATA_DIRECTORY pDataDir;

  pDosH = (PIMAGE_DOS_HEADER)ImageBase;
  if ( pDosH->e_magic != IMAGE_DOS_SIGNATURE )
  {
    return NULL;
  }
  
  pNtH = (PIMAGE_NT_HEADERS)(ImageBase + pDosH->e_lfanew);
  if ( pNtH->Signature != IMAGE_NT_SIGNATURE )
  {
    return NULL;
  }

  if ( ppNtH != NULL )
  {
    *ppNtH = pNtH;
  }
  
  pDataDir = &pNtH->OptionalHeader.DataDirectory[DirIndex];

  if ( pDataDir->VirtualAddress == 0 )
  {
    return NULL;
  }

  if ( ppDataDir != NULL )
  {
    *ppDataDir = pDataDir;
  }
  
  return (PVOID)(ImageBase + Rav2Raw((PVOID)ImageBase, pDataDir->VirtualAddress));
}

typedef BOOLEAN (NTAPI* LOOPUPXREF_CALLBACK)(PULONG RefAddr, PVOID Param);

ULONG NTAPI // 返回找到的引用个数
pLookupImageXRef(
  IN ULONG ImageBase, // 参考地址
  IN PIMAGE_DATA_DIRECTORY pBrDir,
  IN PIMAGE_BASE_RELOCATION pBr,
  IN ULONG SymAddr,
  IN LOOPUPXREF_CALLBACK LookupXRefCallback, // 回调
  IN PVOID pParam
  )
{
  ULONG nRefCount = 0;

  __try
  {
    ULONG Size = 0;
    while ( pBr->SizeOfBlock != 0 && Size < pBrDir->Size )
    {
      PWORD pRelItem = (PWORD)(pBr + 1);
      PIMAGE_BASE_RELOCATION pNextBr = (PIMAGE_BASE_RELOCATION)((ULONG)pBr + pBr->SizeOfBlock);

      while ( (ULONG)pRelItem < (ULONG)pNextBr )
      {
        ULONG RefBlock = Rav2Raw((PVOID)ImageBase, pBr->VirtualAddress);

        if ( (pRelItem[0] >> 12) == 3 )
        {
          PULONG RefAddr = (PULONG)(ImageBase + RefBlock + (pRelItem[0] & 0x0FFF));

          // 验证内存是否有效
          if ( !IsBadReadPtr(RefAddr, 1) && RefAddr[0] == SymAddr )
          {
            nRefCount++;
            if ( LookupXRefCallback(RefAddr, pParam)  )
              break;
          }
        }
        
        pRelItem++;
      }
      
      Size += pBr->SizeOfBlock;
      pBr = pNextBr;
    }
  }
  __except ( EXCEPTION_EXECUTE_HANDLER )
  {
  //  DbgPrint("pLookupImageXRef: Found a Exception!\n");
  }
  
  return nRefCount;
}

ULONG NTAPI
LookupImageXRef(
  IN ULONG ImageBase, // 有效地址
  IN ULONG SymAddr,
  IN LOOPUPXREF_CALLBACK LookupXRefCallback, // 回调
  IN PVOID Param
  )
{
  if ( LookupXRefCallback != NULL )
  {
    PIMAGE_DATA_DIRECTORY pBrDir;
    PIMAGE_BASE_RELOCATION pBr;
    
    pBr = GetImageDirEntry(ImageBase, IMAGE_DIRECTORY_ENTRY_BASERELOC, NULL, &pBrDir);
    if ( pBr != NULL )
    {
      return pLookupImageXRef(ImageBase, pBrDir, pBr, SymAddr, LookupXRefCallback, Param);
    }
  }

  return 0;
}

typedef struct tag_SYSTEM_SERVICE_TABLE {
  PULONG  ServiceTable;  // array of entry points
  PULONG  CounterTable;  // array of usage counters
  ULONG  ServiceLimit;    // number of table entries
  PCHAR  ArgumentTable;  // array of argument counts
} SYSTEM_SERVICE_TABLE, *PSYSTEM_SERVICE_TABLE, **PPSYSTEM_SERVICE_TABLE;

BOOLEAN WINAPI LookXRefCallback(PULONG pRefAddr, PVOID pParam)
{
  USHORT OpCode = *((PUSHORT)pRefAddr - 1);
  if ( OpCode == 0x05C7 ) // mov     ds:_KeServiceDescriptorTable, offset _KiServiceTable
  {
    if ( pParam != NULL )
    {
      *(PULONG)pParam = pRefAddr[1]; // KiServiceTable
    }

    return FALSE;
  }

  return TRUE;
}

BOOL GetSSTEntry(PCTSTR pszServiceName, PSST_ENTRY pSSTEntry)
{
  HANDLE hFile, hFileMap;
  LPVOID pFileData;
  HMODULE hNtdll;
  FARPROC pfnService;
  BOOL bResult;
  TCHAR szNtoskrnlPath[MAX_PATH];
  
  hNtdll = GetModuleHandle(_T("ntdll.dll"));
  pfnService = GetProcAddress(hNtdll, pszServiceName);
  if ( pfnService == NULL || *(PBYTE)pfnService != 0xB8 )
    return FALSE;

  // mov eax, ??
  pSSTEntry->Index = *(PULONG)((PBYTE)pfnService + 1);

  GetSystemDirectory(szNtoskrnlPath, MAX_PATH);
  _tcscat(szNtoskrnlPath, _T("\\ntoskrnl.exe")); // or ntkrnlpa.exe
  
  hFile = CreateFile(szNtoskrnlPath, GENERIC_READ, FILE_SHARE_WRITE | FILE_SHARE_READ,
    NULL, OPEN_EXISTING, 0, NULL);
  if ( hFile == INVALID_HANDLE_VALUE )
    return FALSE;

  hFileMap = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
  CloseHandle(hFile);

  if ( hFileMap == NULL )
    return FALSE;
  
  pFileData = MapViewOfFile(hFileMap, FILE_MAP_READ, 0, 0, 0);
  CloseHandle(hFileMap);
  
  if ( NULL == pFileData )
    return FALSE;
  
  bResult = FALSE;
  __try
  {
    PIMAGE_DOS_HEADER pDosH;
    PIMAGE_NT_HEADERS pNtH;
    PIMAGE_EXPORT_DIRECTORY pExpDir;
    PULONG pFunName, pFunAddr;
    PUSHORT pFunNameOrd;
    ULONG Index;
    
    pDosH = (PIMAGE_DOS_HEADER)pFileData;
    pNtH = (PIMAGE_NT_HEADERS)((ULONG)pFileData + pDosH->e_lfanew);
    pExpDir = (PIMAGE_EXPORT_DIRECTORY)((ULONG)pFileData + pNtH->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);

    pFunName = (PULONG)((ULONG)pFileData + Rav2Raw(pFileData, pExpDir->AddressOfNames));
    pFunNameOrd = (PUSHORT)((ULONG)pFileData + Rav2Raw(pFileData, pExpDir->AddressOfNameOrdinals));
    pFunAddr = (PULONG)((ULONG)pFileData + Rav2Raw(pFileData, pExpDir->AddressOfFunctions));

    for ( Index = 0; Index < pExpDir->NumberOfNames; Index++ )    
    {
      PCSTR pName = (PCSTR)((ULONG)pFileData + Rav2Raw(pFileData, pFunName[Index]));
      if ( strcmp(pName, "KeServiceDescriptorTable") == 0 )
      {
        ULONG _KiServiceTable, _KeServiceDescriptorTable; // 类型都是 va
        
        _KiServiceTable = 0;
        _KeServiceDescriptorTable = pNtH->OptionalHeader.ImageBase + pFunAddr[pFunNameOrd[Index]];
  
        LookupImageXRef((ULONG)pFileData, _KeServiceDescriptorTable, LookXRefCallback, &_KiServiceTable);
        if ( _KiServiceTable != 0 )
        {
          PULONG ServiceTable;
          _KiServiceTable -= pNtH->OptionalHeader.ImageBase;

          ServiceTable = (PULONG)((ULONG)pFileData + Rav2Raw(pFileData, _KiServiceTable));
          pSSTEntry->OrgValue = ServiceTable[pSSTEntry->Index] - pNtH->OptionalHeader.ImageBase;
        }

        bResult = TRUE;
        break;
      }
    }
  }
  __except ( EXCEPTION_EXECUTE_HANDLER )
  {
    bResult = FALSE;
  }

  UnmapViewOfFile(pFileData);
  return bResult;

}

HANDLE GetServiceByName(IN PCTSTR pszServiceName, OUT OPTIONAL SC_HANDLE* pScManager)
{
  SC_HANDLE hScManager, hScService;
  
  hScService = NULL;
  hScManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
  
  if ( pScManager != NULL )
    *pScManager = hScManager;

  if ( hScManager != NULL )
  {
    hScService = OpenService(hScManager, pszServiceName, SERVICE_ALL_ACCESS);
    if ( pScManager == NULL )
      CloseServiceHandle(hScManager);
  }

  return hScService;
}

SC_HANDLE InstallDriver(IN PCTSTR pszServiceName, IN PCTSTR pszDriverFile)
{
  SC_HANDLE hScManager, hScService;
  
  hScManager = NULL;
  hScService = GetServiceByName(pszServiceName, &hScManager);

  if ( hScService == NULL && hScManager != NULL )
  {
    hScService = CreateService(hScManager, pszServiceName, pszServiceName, SERVICE_ALL_ACCESS,
      SERVICE_KERNEL_DRIVER, SERVICE_DEMAND_START, SERVICE_ERROR_NORMAL, pszDriverFile, NULL,
      NULL, NULL, NULL, NULL
      );
//    GetLastError();
  }

  if ( hScManager != NULL )
    CloseServiceHandle(hScManager);
  
  return hScService;
}

BOOL UninstallDriver(SC_HANDLE hScService)
{
  SERVICE_STATUS Status;
  
  if ( ControlService(hScService, SERVICE_CONTROL_STOP, &Status) )
  {
    DeleteService(hScService);
    CloseServiceHandle(hScService);

    return TRUE;
  }
//  GetLastError();

  return FALSE;
}

BOOL ResetSSTEntry(PCSTR lpServiceName)
{
  SC_HANDLE hScService;
  HANDLE hDevice;
  SST_ENTRY SstEntry;
  ULONG cbReturn;
  TCHAR szDriverPath[MAX_PATH];
  PTSTR pszDriverName;
  BOOL bResult;

  SstEntry.cbSize = sizeof (SST_ENTRY);
  if ( !GetSSTEntry(lpServiceName, &SstEntry) )
    return FALSE;
  
  GetModuleFileName(NULL, szDriverPath, MAX_PATH);
  pszDriverName = _tcsrchr(szDriverPath, _T('\\'));
  _tcscpy(pszDriverName, _T("\\RSSDT.sys"));
  
  hScService = GetServiceByName(_T("RSSDT"), NULL);
  if ( hScService == NULL )
  {
    hScService = InstallDriver(_T("RSSDT"), szDriverPath);
    if ( hScService == NULL )
      return FALSE;
  }

  bResult = FALSE;
  if ( StartService(hScService, 0, NULL) || GetLastError() == ERROR_ALREADY_EXISTS )
  {
    hDevice = CreateFile(_T("\\\\.\\RSSDT"), 
      GENERIC_READ, 
      FILE_SHARE_READ | FILE_SHARE_WRITE,
      NULL, 
      OPEN_EXISTING, 
      0, 
      NULL
      );
    GetLastError();

    if ( hDevice != INVALID_HANDLE_VALUE )
    {
      bResult = DeviceIoControl(hDevice, IOCTL_RSSDT_ENTRY_SET, &SstEntry, sizeof (SstEntry), 
        &SstEntry, sizeof (SstEntry), &cbReturn, NULL);

      CloseHandle(hDevice);
    }
  }

  UninstallDriver(hScService);
  return bResult;
}

#if 1
int __cdecl main()
{
  ResetSSTEntry("NtCreateKey");
  ResetSSTEntry("NtCreatePagingFile");
  ResetSSTEntry("NtEnumerateValueKey");
  ResetSSTEntry("NtOpenKey");
  ResetSSTEntry("NtQueryKey");
  ResetSSTEntry("NtQueryValueKey");
  ResetSSTEntry("NtSetSystemPowerState");

  return 0;
}

#endif