在注册表中设置用户或系统环境变量的值

// Copyright © 2022, 飞麦 <[email protected]>, All rights reserved.

// 程序: 在注册表中设置用户或系统环境变量的值
// 作者: 飞麦, mailto:[email protected], 2010-10-26~2023-12-22

#include <windows.h>
#include <tchar.h>
#include <stdio.h>
#include "Shlwapi.h"  // StrStrI, PathFileExists, PathIsDirectory


#define show() Show(_T(__FILE__), __LINE__)


static const size_t MAX_DATA = MAX_PATH * 256;
static const size_t MAX_DIR = MAX_PATH * 16;

typedef TCHAR TypeData[MAX_DATA];
typedef TCHAR TypeDir[MAX_DIR];

// 可执行程序查找路径
static LPCTSTR PATH = _T("Path");
// Java 类查找路径
static LPCTSTR CLASSPATH = _T("ClassPath");
// Ruby 类查找路径
static LPCTSTR RUBYLIB = _T("RubyLib");
// 环境变量性质
static LPCTSTR szMode = _T("用户");

// 直接使用 stdout 难以正确输出非英文 Unicode 字符
static FILE* fout = NULL;
// 本程序运行结果输出文件
static TypeDir szLogFile;
// Unicode 编码文件头
static const TCHAR BOM = L'\xFEFF';
// 用户环境变量子键
static LPCTSTR USR_SUB_KEY = _T("Environment");
// 系统环境变量子键
static LPCTSTR SYS_SUB_KEY = _T("SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment");
// 帮助信息
static LPCTSTR HELP_MSG =
_T("程序: envy\n")
_T("作者: 飞麦, mailto:[email protected], 2010-10-26~2023-12-22.\n")
_T("功能: 在注册表中设置、增加、减少、清除系统(/s)或用户(/u,默认)环境变量，涉及系统环境变量需以管理员身份运行。\n")
_T("用法: envy [/s|/u] var[+|-]=\"[item]\"\n")
_T("示例: envy RUBYLIB=\"C:\\R\" /s Path-=\"D:\\jdk1.5\\bin\" /u Path+=\"D:\\jdk1.6\\bin\" NOTHING=\"\"\n")
_T("解释: 1设置用户环境变量，2减少系统环境变量的值，3增加用户环境变量的值，4清除用户环境变量。\n")
_T("========================================\n");


// 打开日志文件
static bool log_on() {
    // 运行结果输出目录
    static TypeDir szLogDir;
    // 获取临时目录
    (void)GetTempPath(MAX_DIR, szLogDir);
    (void)_stprintf_s(szLogFile, MAX_DIR, _T("%s#envy.tmp"), szLogDir);
    // b 输出Unicode字符, 不影响type, 影响notepad(换行不加回车)；t输出ANSI/MBCS字符(换行自动加回车)
    _tfopen_s(&fout, szLogFile, _T("wb"));
    if (fout == NULL) {
        // Only English is allowed in _tprintf
        _tprintf(_T("Fail on open %s for write.\n"), szLogFile);
        return false;
    }
    // TYPE 命令根据 BOM 来判断此文件是否为 Unicode 编码
    _ftprintf(fout, _T("%c%s"), BOM, HELP_MSG);
    return true;
}


// 关闭并显示日志文件
static void log_off() {
    (void)fclose(fout);

    static TypeDir szCmdLine;
    _stprintf_s(szCmdLine, MAX_DIR, _T("cmd.exe /c type %s"), szLogFile);
    STARTUPINFO startInfo;
    GetStartupInfo(&startInfo);
    PROCESS_INFORMATION proInfo;
    DWORD dwCreaFlags = CREATE_UNICODE_ENVIRONMENT;
    (void)CreateProcess(NULL, szCmdLine, NULL, NULL, FALSE, dwCreaFlags, NULL, NULL, &startInfo, &proInfo);
    (void)WaitForSingleObject(proInfo.hProcess, INFINITE);
    (void)CloseHandle(proInfo.hThread);
    (void)CloseHandle(proInfo.hProcess);
    (void)DeleteFile(szLogFile);
}


// 显示系统函数的错误信息
static bool Show(LPCTSTR szFile, int iLine) {
    DWORD dwMsgId = GetLastError();
    LPVOID lpMsgBuf = NULL;
    DWORD dwFlags = FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM;
    DWORD dwLangId = MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT);
    LPTSTR lpBuffer = LPTSTR(&lpMsgBuf);
    DWORD dwLen = FormatMessage(dwFlags, NULL, dwMsgId, dwLangId, lpBuffer, 0, NULL);
    LPCTSTR szMsg = dwLen > 0 ? LPTSTR(lpMsgBuf) : _T("\n");
    _ftprintf(fout, _T("文件=%s 行号=%d 编码=%lu 错误=%s"), szFile, iLine, dwMsgId, szMsg);
    (void)LocalFree(lpMsgBuf);
    return false;
}


// 查询注册表值
static bool getRegistry(HKEY hKey, LPCTSTR lpSubKey, LPCTSTR szVar, LPDWORD pdwType, LPTSTR szData, LPDWORD lpcb) {
    HKEY hkResult;
    LONG lResult = RegOpenKeyEx(hKey, lpSubKey, NULL, KEY_QUERY_VALUE, &hkResult);
    if (lResult != ERROR_SUCCESS)
        (void)show();
    else {
        lResult = RegQueryValueEx(hkResult, szVar, NULL, pdwType, LPBYTE(szData), lpcb);
        if (lResult != ERROR_SUCCESS)
            (void)show();
        *lpcb /= sizeof(TCHAR);
        LONG lResult2 = RegCloseKey(hkResult);
        if (lResult2 != ERROR_SUCCESS) {
            (void)show();
            lResult = lResult2;
        }
    }
    return lResult == ERROR_SUCCESS;
}


// 设置注册表值
static bool setRegistry(HKEY hKey, LPCTSTR lpSubKey, LPCTSTR szVar, DWORD dwType, LPTSTR szData, DWORD cb) {
    HKEY hkResult;
    LONG lResult = RegOpenKeyEx(hKey, lpSubKey, NULL, KEY_ALL_ACCESS, &hkResult);
    if (lResult != ERROR_SUCCESS)
        (void)show();
    else {
        cb *= sizeof(TCHAR);
        lResult = RegSetValueEx(hkResult, szVar, NULL, dwType, LPBYTE(szData), cb);
        if (lResult != ERROR_SUCCESS)
            (void)show();
        LONG lResult2 = RegCloseKey(hkResult);
        if (lResult2 != ERROR_SUCCESS) {
            (void)show();
            lResult = lResult2;
        }
    }
    return lResult == ERROR_SUCCESS;
}


// 删除注册表值
static bool delRegistry(HKEY hKey, LPCTSTR lpSubKey, LPCTSTR szVar) {
    HKEY hkResult;
    LONG lResult = RegOpenKeyEx(hKey, lpSubKey, NULL, KEY_ALL_ACCESS, &hkResult);
    if (lResult != ERROR_SUCCESS)
        (void)show();
    else {
        lResult = RegDeleteValue(hkResult, szVar);
        if (lResult != ERROR_SUCCESS)
            (void)show();
        LONG lResult2 = RegCloseKey(hkResult);
        if (lResult2 != ERROR_SUCCESS) {
            (void)show();
            lResult = lResult2;
        }
    }
    return lResult == ERROR_SUCCESS;
}


// 将环境变量 szVar (值为 szData) 中的 szItem 消除掉, 再在后面加上 szItem
static void adjustEnvVariable(LPCTSTR szVar, TypeData szData, LPTSTR szItem) {
    // 判断 szVar 只是为了防止 Lint 告警，szItem 为空串时不需要处理
    if (*szItem == NULL || *szVar == NULL)
        return;
    // 可能为空串, 如 ClassPath 初始没有被设置
    static TypeData szDataPlus;
    // 分号是 Windows 的路径分隔符
    _stprintf_s(szDataPlus, MAX_DATA, _T(";%s%s"), szData, *szData == NULL ? _T("") : _T(";"));
    static TypeData szItemPlus;
    int iLen = _stprintf_s(szItemPlus, MAX_DATA, _T(";%s%s"), szItem, *szItem == NULL ? _T("") : _T(";"));
    // 将 szDataPlus 中的所有 szItemPlus 均替换为分号
    for (;;) {
        // 不区分大小写的查找, Windows 的路径不区分大小写
        LPTSTR szFind = StrStrI(szDataPlus, szItemPlus);
        // 如果找不到了则退出
        if (szFind == NULL)
            break;
        // 将一个 szItemPlus 替换为分号, 通过前后各增加分号, 确保被搜索串的唯一性与完整性
        _tcscpy_s(szFind + 1, MAX_DATA, szFind + iLen);
    }

    // 将 szDataPlus(不要前面的分号) 和 sItemPlus(不要前面的分号) 合并后放到 szData 中
    iLen = _stprintf_s(szData, MAX_DATA, _T("%s%s"), szDataPlus + 1, szItemPlus + 1);
    // 删除最后的分号
    szData[iLen - 1] = NULL;
}


// 清洁 szVar 环境变量, 将其中不存在的目录或文件清除
static void cleanEnvVariable(LPCTSTR szVar, LPTSTR szData) {
    // maximum characters in expanded string
    DWORD dwLenPath = MAX_DATA;
    static TypeData szExpData;
    dwLenPath = ExpandEnvironmentStrings(szData, szExpData, dwLenPath);
    if (dwLenPath == 0) {
        _ftprintf(fout, _T("%s=%s：扩展环境变量失败！"), szVar, szData);
        return;
    }
    LPTSTR szExpLoc = szExpData;
    LPTSTR szLoc = szData;
    static TypeData szTmpData;
    *szTmpData = NULL;
    for (;;) {
        LPTSTR szExpSemi = _tcschr(szExpLoc, static_cast<int>(_T(';')));
        LPTSTR szSemi = _tcschr(szLoc, static_cast<int>(_T(';')));
        static TypeData szExpDir;
        static TypeData szDir;
        if (szExpSemi == NULL && szSemi == NULL) {
            _tcscpy_s(szExpDir, MAX_DATA, szExpLoc);
            _tcscpy_s(szDir, MAX_DATA, szLoc);
        }
        else if (szExpSemi != NULL && szSemi != NULL) {
            *szExpSemi = NULL;
            _tcscpy_s(szExpDir, MAX_DATA, szExpLoc);
            *szExpSemi = _T(';');
            *szSemi = NULL;
            _tcscpy_s(szDir, MAX_DATA, szLoc);
            *szSemi = _T(';');
            szExpLoc = szExpSemi + 1;
            szLoc = szSemi + 1;
        }
        else {
            _ftprintf(fout, _T("%s=%s：扩展后分号数量不同！\n"), szVar, szData);
            return;
        }
        bool bJar = _tcsicmp(szVar, CLASSPATH) == 0 && PathFileExists(szExpDir);
        bool bDot = _tcscmp(szExpDir, _T(".")) == 0;
        bool bDir = PathIsDirectory(szExpDir);
        if (bJar || bDot || bDir) {
            if (*szTmpData != NULL)
                _tcscat_s(szTmpData, MAX_DATA, _T(";"));
            _tcscat_s(szTmpData, MAX_DATA, szDir);
        }
        else if (*szDir != NULL)
            _ftprintf(fout, _T("%s：目录或文件因不存在而被抛弃！\n"), szDir);
        if (szExpSemi == NULL)
            break;
    }
    _tcscpy_s(szData, MAX_DATA, szTmpData);
}


// 设置注册表值
// hKey: HKEY_LOCAL_MACHINE, HKEY_CURRENT_USER, ...
// lpSubKey: (HKEY_CURRENT_USER\\Environment, SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment)
// szVar: Path, Classpath, RUBYLIB, ...
// dwType: REG_EXPAND_SZ (%Path_HKCU%), REG_SZ, ...
// szData: D:\\jdk1.5\\bin, ...
// cb: length of szData
static bool setRegistry2(HKEY hKey, LPCTSTR lpSubKey, LPCTSTR szVar, DWORD dwType, LPTSTR szData, int chgType) {
    static TypeData szNewData;
    if (chgType == 0)
        _tcscpy_s(szNewData, MAX_DATA, szData);
    else {
        static TypeData szOldData;
        DWORD cbData = MAX_DATA * sizeof(TCHAR);
        DWORD dwType0 = dwType;
        if (!getRegistry(hKey, lpSubKey, szVar, &dwType0, szOldData, &cbData)) {
            if (chgType == 1)
                _tcscpy_s(szNewData, MAX_DATA, szData);
            else
                *szNewData = NULL;
        }
        else {
            if (chgType == 1) {
                _tcscpy_s(szNewData, MAX_DATA, szOldData);
                adjustEnvVariable(szVar, szNewData, szData);
            }
            else {
                _stprintf_s(szNewData, MAX_DATA, _T(";%s%s"), szOldData, *szOldData == NULL ? _T("") : _T(";"));
                static TypeData szItemPlus;
                // 分号是 Windows 的路径分隔符
                int iLen = _stprintf_s(szItemPlus, MAX_DATA, _T(";%s;"), szData);
                // 将 szNewData 中的所有 szItemPlus 均替换为分号
                for (;;) {
                    // 不区分大小写的查找, Windows 的路径不区分大小写
                    LPTSTR szFind = StrStrI(szNewData, szItemPlus);
                    if (szFind == NULL)
                        // 如果找不到了则退出
                        break;
                    // 将一个 szItemPlus 替换为分号，通过前后各增加分号, 确保被搜索串的唯一性与完整性
                    _tcscpy_s(szFind + 1, MAX_DATA, szFind + iLen);
                }
            }
        }
    }
    if (*szNewData == NULL) {
        _ftprintf(fout, _T("删除环境变量：%s\n"), szVar);
        return delRegistry(hKey, lpSubKey, szVar);
    }
    if (_tcsicmp(szVar, PATH) == 0 || _tcsicmp(szVar, CLASSPATH) == 0 || _tcsicmp(szVar, RUBYLIB) == 0)
        cleanEnvVariable(szVar, szNewData);
    _ftprintf(fout, _T("设置%s环境变量 %s 为 %s\n"), szMode, szVar, szNewData);
    DWORD cb = DWORD(_tcslen(szNewData) + 1);
    return setRegistry(hKey, lpSubKey, szVar, dwType, szNewData, cb);
}


// 设置 szVar 环境变量的值为 szData
static bool setEnvVariable(LPCTSTR szVar, LPTSTR szData, int chgType) {
    if (_tcscmp(szMode, _T("系统")) == 0)
        return setRegistry2(HKEY_LOCAL_MACHINE, SYS_SUB_KEY, szVar, REG_EXPAND_SZ, szData, chgType);
    else
        return setRegistry2(HKEY_CURRENT_USER, USR_SUB_KEY, szVar, REG_EXPAND_SZ, szData, chgType);
}


// 处理单个环境变量
static int deal(LPTSTR szVarExpr) {
    LPTSTR szEqu = _tcschr(szVarExpr, int(_T('=')));
    LPCTSTR szAction = _T("设置");
    LPCTSTR szBy = _T("为");
    LPTSTR szVarTail = NULL;
    int chgType = 0;
    if (szEqu == NULL) {
        _ftprintf(fout, _T("参数 %s 中应有等号！\n"), szVarExpr);
        return 3;
    }
    else if (szEqu == szVarExpr) {
        _ftprintf(fout, _T("参数 %s 中等号前应有环境变量名！\n"), szVarExpr);
        return 4;
    }
    else if (*(szEqu - 1) == _T('+')) {
        *(szEqu - 1) = NULL;
        szVarTail = szEqu - 2;
        chgType = 1;
        szAction = _T("增加");
        szBy = _T("用");
    }
    else if (*(szEqu - 1) == _T('-')) {
        *(szEqu - 1) = NULL;
        szVarTail = szEqu - 2;
        chgType = -1;
        szAction = _T("减少");
        szBy = _T("以");
    }
    else {
        *szEqu = NULL;
        szVarTail = szEqu - 1;
    }
    LPTSTR szVar = szVarExpr;
    while (*szVar == _T(' '))
        szVar++;
    while (*szVarTail == _T(' '))
        *szVarTail-- = NULL;
    LPTSTR szData = szEqu + 1;
    while (*szData == _T(' '))
        szData++;
    LPTSTR szNil = _tcschr(szData, 0);
    if (szNil != NULL) {
        LPTSTR szDataTail = szNil - 1;
        while (*szDataTail == _T(' '))
            *szDataTail-- = NULL;
    }
    int iRe = setEnvVariable(szVar, szData, chgType) ? 0 : 3;
    LPCTSTR szRes = iRe == 0 ? _T("成功。") : _T("失败！");
    _ftprintf(fout, _T("%s%s环境变量 %s %s %s：%s\n"), szAction, szMode, szVar, szBy, szData, szRes);
    return iRe;
}


// 广播环境变量到所有窗口
static int broadcast() {
    DWORD_PTR pdwRetValue = NULL;
    LPARAM param = (LPARAM)_T("Environment");
    LRESULT iR = SendMessageTimeout(HWND_BROADCAST, WM_SETTINGCHANGE, 0, param, SMTO_ABORTIFHUNG, 5000, &pdwRetValue);
    // 若出错
    if (iR == 0) {
        (void)show();
        _ftprintf(fout, _T("广播环境变量到所有窗口时出错！\n"));
        return 1;
    }
    return 0;
}


// 命令行参数是否含有 /s 选项
static bool has_sys_opt(int argc, LPTSTR argv[]) {
    for (int i = 1; i < argc; i++)
        if (lstrcmpi(argv[i], _T("/s")) == 0)
            return true;
    return false;
}


// 是否有管理员权限
static bool has_admin_right() {
    HKEY hkResult;
    LONG lResult = RegOpenKeyEx(HKEY_LOCAL_MACHINE, SYS_SUB_KEY, NULL, KEY_ALL_ACCESS, &hkResult);
    bool ok = lResult == ERROR_SUCCESS;
    if (ok)
        (void)RegCloseKey(hkResult);
    return ok;
}


// 处理命令行参数
static int inner(int argc, LPTSTR argv[]) {
    if (argc <= 1)
        return 1;
    int iRet = 0;
    if (has_sys_opt(argc, argv) && !has_admin_right()) {
        _ftprintf(fout, _T("当使用 /s 选项时，应以管理员身份运行。\n"));
        return 9;
    }
    for (int i = 1; i < argc; i++) {
        if (*argv[i] == _T('/')) {
            if (lstrcmpi(argv[i], _T("/s")) == 0)
                szMode = _T("系统");
            else if (lstrcmpi(argv[i], _T("/u")) == 0)
                szMode = _T("用户");
            else
                _ftprintf(fout, _T("%s：错误选项！\n"), argv[i]);
        }
        else
            iRet += deal(argv[i]);
    }
    iRet += broadcast();
    return iRet;
}


// 主程序
int _tmain(int argc, LPTSTR argv[]) {
    if (!log_on())
        return 2;
    int iRet = inner(argc, argv);
    log_off();
    return iRet;
}