Add a read-only list_registry MCP tool for querying a remote Windows host's registry. An empty path returns the five root keys; a path such as HKEY_LOCAL_MACHINE plus subkeys returns that key's immediate subkeys and values (name, type, and formatted data). Client side: enumerate with KEY_READ instead of KEY_ALL_ACCESS and map the full value-type set (REG_SZ, REG_DWORD, REG_BINARY, REG_EXPAND_SZ, REG_MULTI_SZ, REG_QWORD, REG_NONE) so unknown types are no longer mis-reported as REG_SZ. The registry manager now always sends both the TOKEN_REG_PATH and TOKEN_REG_KEY packets (an empty packet when a part is empty), making the two-packet reply deterministic for the server. Server side: add the three-phase flow (COMMAND_REGEDIT, COMMAND_REG_FIND, then PATH plus KEY) with per-host pending state, parse the fixed-width wire format with bounds checks, format value data for JSON, and guard registry paths against exceeding MAX_PATH to protect the client stack. Co-Authored-By: deepseek-v4-pro
197 lines
6.4 KiB
C++
197 lines
6.4 KiB
C++
// RegisterOperation.cpp: implementation of the RegisterOperation class.
|
||
//
|
||
//////////////////////////////////////////////////////////////////////
|
||
|
||
#include "stdafx.h"
|
||
#include "RegisterOperation.h"
|
||
#include "Common.h"
|
||
#include <IOSTREAM>
|
||
|
||
//////////////////////////////////////////////////////////////////////
|
||
// Construction/Destruction
|
||
//////////////////////////////////////////////////////////////////////
|
||
|
||
enum MYKEY {
|
||
MHKEY_CLASSES_ROOT,
|
||
MHKEY_CURRENT_USER,
|
||
MHKEY_LOCAL_MACHINE,
|
||
MHKEY_USERS,
|
||
MHKEY_CURRENT_CONFIG
|
||
};
|
||
|
||
|
||
enum KEYVALUE {
|
||
MREG_SZ, // REG_SZ
|
||
MREG_DWORD, // REG_DWORD
|
||
MREG_BINARY, // REG_BINARY
|
||
MREG_EXPAND_SZ, // REG_EXPAND_SZ
|
||
MREG_MULTI_SZ, // REG_MULTI_SZ
|
||
MREG_QWORD, // REG_QWORD
|
||
MREG_NONE // REG_NONE / 未知类型
|
||
};
|
||
|
||
struct REGMSG {
|
||
int count; //名字个数
|
||
DWORD size; //名字大小
|
||
DWORD valsize; //值大小
|
||
|
||
};
|
||
RegisterOperation::RegisterOperation(char bToken)
|
||
{
|
||
switch(bToken) {
|
||
case MHKEY_CLASSES_ROOT:
|
||
MKEY=HKEY_CLASSES_ROOT;
|
||
break;
|
||
case MHKEY_CURRENT_USER:
|
||
MKEY=HKEY_CURRENT_USER;
|
||
break;
|
||
case MHKEY_LOCAL_MACHINE:
|
||
MKEY=HKEY_LOCAL_MACHINE;
|
||
break;
|
||
case MHKEY_USERS:
|
||
MKEY=HKEY_USERS;
|
||
break;
|
||
case MHKEY_CURRENT_CONFIG:
|
||
MKEY=HKEY_CURRENT_CONFIG;
|
||
break;
|
||
default:
|
||
MKEY=HKEY_LOCAL_MACHINE;
|
||
break;
|
||
}
|
||
|
||
ZeroMemory(KeyPath,MAX_PATH);
|
||
}
|
||
|
||
RegisterOperation::~RegisterOperation()
|
||
{
|
||
|
||
}
|
||
|
||
|
||
char* RegisterOperation::FindPath()
|
||
{
|
||
char *szBuffer=NULL;
|
||
HKEY hKey; //注册表返回句柄
|
||
/*打开注册表 User kdjfjkf\kdjfkdjf\ */
|
||
if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_READ,&hKey)==ERROR_SUCCESS) { //打开
|
||
DWORD dwIndex=0,NameCount,NameMaxLen;
|
||
DWORD KeySize,KeyCount,KeyMaxLen,MaxDataLen;
|
||
//这就是枚举了
|
||
if(RegQueryInfoKey(hKey,NULL,NULL,NULL,&KeyCount, //14
|
||
&KeyMaxLen,NULL,&NameCount,&NameMaxLen,&MaxDataLen,NULL,NULL)!=ERROR_SUCCESS) {
|
||
return NULL;
|
||
}
|
||
//一点保护措施
|
||
KeySize=KeyMaxLen+1;
|
||
if(KeyCount>0&&KeySize>1) {
|
||
int Size=sizeof(REGMSG)+1;
|
||
|
||
DWORD DataSize=KeyCount*KeySize+Size+1; //[TOKEN_REG_PATH][2 11 ccccc\0][11][11]
|
||
szBuffer=(char*)LocalAlloc(LPTR, DataSize);
|
||
if (szBuffer == NULL) {
|
||
return NULL;
|
||
}
|
||
ZeroMemory(szBuffer,DataSize);
|
||
szBuffer[0]=TOKEN_REG_PATH; //命令头
|
||
REGMSG msg; //数据头
|
||
msg.size=KeySize;
|
||
msg.count=KeyCount;
|
||
memcpy(szBuffer+1,(void*)&msg,sizeof(REGMSG));
|
||
|
||
char * szTemp=new char[KeySize];
|
||
for(dwIndex=0; dwIndex<KeyCount; dwIndex++) { //枚举项
|
||
ZeroMemory(szTemp,KeySize);
|
||
DWORD i=KeySize; //21
|
||
RegEnumKeyEx(hKey,dwIndex,szTemp,&i,NULL,NULL,NULL,NULL);
|
||
strcpy(szBuffer+dwIndex*KeySize+Size,szTemp);
|
||
}
|
||
delete[] szTemp;
|
||
RegCloseKey(hKey);
|
||
}
|
||
}
|
||
|
||
return szBuffer;
|
||
}
|
||
|
||
|
||
void RegisterOperation::SetPath(char *szPath)
|
||
{
|
||
ZeroMemory(KeyPath,MAX_PATH);
|
||
strcpy(KeyPath,szPath);
|
||
}
|
||
|
||
char* RegisterOperation::FindKey()
|
||
{
|
||
char *szValueName; //键值名称
|
||
LPBYTE szValueData; //键值数据
|
||
|
||
char *szBuffer=NULL;
|
||
HKEY hKey; //注册表返回句柄
|
||
if(RegOpenKeyEx(MKEY,KeyPath,0,KEY_READ,&hKey)==ERROR_SUCCESS) { //打开
|
||
DWORD dwIndex=0,NameSize,NameCount,NameMaxLen,Type;
|
||
DWORD KeyCount,KeyMaxLen,DataSize,MaxDataLen;
|
||
//这就是枚举了
|
||
if(RegQueryInfoKey(hKey,NULL,NULL,NULL,
|
||
&KeyCount,&KeyMaxLen,NULL,&NameCount,&NameMaxLen,&MaxDataLen,NULL,NULL)!=ERROR_SUCCESS) {
|
||
return NULL;
|
||
}
|
||
if(NameCount>0&&MaxDataLen>0) {
|
||
DataSize=MaxDataLen+1;
|
||
NameSize=NameMaxLen+100;
|
||
REGMSG msg;
|
||
msg.count=NameCount; //总个数
|
||
msg.size=NameSize; //名字大小
|
||
msg.valsize=DataSize; //数据大小
|
||
const int msgsize=sizeof(REGMSG);
|
||
// 头 标记 名字 数据
|
||
DWORD size=sizeof(REGMSG)+
|
||
sizeof(BYTE)*NameCount+ NameSize*NameCount+DataSize*NameCount+10;
|
||
szBuffer = (char*)LocalAlloc(LPTR, size);
|
||
if (szBuffer==NULL) {
|
||
return NULL;
|
||
}
|
||
ZeroMemory(szBuffer,size);
|
||
szBuffer[0]=TOKEN_REG_KEY; //命令头
|
||
memcpy(szBuffer+1,(void*)&msg,msgsize); //数据头
|
||
|
||
szValueName=(char *)malloc(NameSize);
|
||
szValueData=(LPBYTE)malloc(DataSize);
|
||
if (szValueName==NULL||szValueData == NULL) {
|
||
return NULL;
|
||
}
|
||
char *szTemp=szBuffer+msgsize+1;
|
||
for(dwIndex=0; dwIndex<NameCount; dwIndex++) { //枚举键值
|
||
ZeroMemory(szValueName,NameSize);
|
||
ZeroMemory(szValueData,DataSize);
|
||
|
||
DataSize=MaxDataLen+1;
|
||
NameSize=NameMaxLen+100;
|
||
|
||
RegEnumValue(hKey,dwIndex,szValueName,&NameSize,
|
||
NULL,&Type,szValueData,&DataSize);//读取键值
|
||
|
||
// 值类型映射:REG_SZ/DWORD/BINARY/EXPAND_SZ 之外补 MULTI_SZ/QWORD/NONE,
|
||
// 未识别类型一律 MREG_NONE(避免误报为 REG_SZ)。
|
||
switch (Type) {
|
||
case REG_SZ: szTemp[0] = MREG_SZ; break;
|
||
case REG_EXPAND_SZ: szTemp[0] = MREG_EXPAND_SZ; break;
|
||
case REG_MULTI_SZ: szTemp[0] = MREG_MULTI_SZ; break;
|
||
case REG_DWORD: szTemp[0] = MREG_DWORD; break;
|
||
case REG_QWORD: szTemp[0] = MREG_QWORD; break;
|
||
case REG_BINARY: szTemp[0] = MREG_BINARY; break;
|
||
case REG_NONE: szTemp[0] = MREG_NONE; break;
|
||
default: szTemp[0] = MREG_NONE; break;
|
||
}
|
||
szTemp+=sizeof(BYTE);
|
||
strcpy(szTemp,szValueName);
|
||
szTemp+=msg.size;
|
||
memcpy(szTemp,szValueData,msg.valsize);
|
||
szTemp+=msg.valsize;
|
||
}
|
||
free(szValueName);
|
||
free(szValueData);
|
||
}
|
||
}
|
||
return szBuffer;
|
||
}
|