psubst.cpp

来自「以C++示范使用微软操作系统内建指令”SUBST”的范例工具程序.在档案总管中加」· C++ 代码 · 共 233 行

#include <windows.h>
#include <stdio.h>
#include <string.h>

static const char help_text[] = 
    "PSUBST v1.0a (c) by Alter, Home site http://www.alter.org.ua\n"
    "Usage:\n"
    "\tpsubst\t[<command>] <drive letter>: <target path>\n"
    "Commands:\n"
    "\t/D\tdelete existing subst drive letter\n"
    "\t/DF\tforce delete existing subst drive letter\n"
    "\t/P\tcreate persistent subst (across system reboots)\n"
    "\t/?,/H\tdisplay this help message\n"
    "\t\twithout commands operates like standard SUBST utility\n"
    "Examples:\n"
    "\tpsubst\t/P Y: C:\\temp\n"
    "\tpsubst\t/DF Y:\n"
    ;

void
usage(void)
{
    HANDLE stdOuth = GetStdHandle(STD_OUTPUT_HANDLE);
    ULONG WrittenBytes;

    if(stdOuth && stdOuth != ((HANDLE)(-1))) {
        WriteFile(stdOuth, help_text, strlen(help_text), &WrittenBytes, NULL);
    } else {
        MessageBox(NULL, help_text, "Usage", MB_OK);
    }
}


/// Start console formatter with desired parameters
DWORD
WINAPI LauncherRoutine(
    PWCHAR ExecStr
    )
{
    STARTUPINFOW proc_startup_info;
    PROCESS_INFORMATION proc_info;
    INT  index;
    ULONG RetCode;

    proc_startup_info.cb = sizeof(proc_startup_info);
    proc_startup_info.lpReserved = 0;
    proc_startup_info.lpReserved2 = 0;
    proc_startup_info.cbReserved2 = 0;
    proc_startup_info.lpDesktop = 0;
    proc_startup_info.lpTitle = 0;
    proc_startup_info.dwFlags = STARTF_USESTDHANDLES;

    proc_startup_info.hStdInput = GetStdHandle(STD_OUTPUT_HANDLE);;
    proc_startup_info.hStdOutput = GetStdHandle(STD_INPUT_HANDLE);;
    proc_startup_info.hStdError = GetStdHandle(STD_ERROR_HANDLE);;

    if(CreateProcessW(NULL, ExecStr, 0,0, TRUE, NORMAL_PRIORITY_CLASS,
                  0,0, &proc_startup_info, &proc_info)) {

        WaitForSingleObject(proc_info.hProcess, -1);
        GetExitCodeProcess(proc_info.hProcess, &RetCode);
        index=0;
        
        CloseHandle(proc_info.hThread);
        CloseHandle(proc_info.hProcess);
    } else {
        return -1;
    }
    return RetCode;
}

#define PSUBST_REG_PATH L"SYSTEM\\CurrentControlSet\\Control\\Session Manager\\DOS Devices"

PWCHAR
FixSlashes(
    PWCHAR Path,
    BOOLEAN Drv
    )
{
    PWCHAR ret;
    WCHAR a;
    ULONG i;

    ret = (PWCHAR)GlobalAlloc(GMEM_DISCARDABLE, (wcslen(Path)*2)*sizeof(WCHAR));
    if(!ret) {
        printf("Insufficient resources\n");
        ExitProcess(-2);
    }
    i = 0;
    while(a = *Path) {
        if(a == '/') {
            ret[i] = '\\';
        } else {
            ret[i] = a;
        }
        Path++;
        i++;
    }
    ret[i] = 0;
    if(i) {
        if( ret[i-1] == '\\' ) {
            if(i>1 && (Drv || (!Drv && ret[i-2] != ':'))) {
                ret[i-1] = 0;
            }
        } else
        if(!Drv && ret[i-1] == ':' ) {
            ret[i] = '\\';
            ret[i+1] = 0;
        }
    }
    return ret;
}

DWORD
CreatePSubst(
    PWCHAR DrvStr,
    PWCHAR TargetStr
    )
{
    HKEY hKey;
    ULONG len;
    WCHAR subst_name[MAX_PATH*2];

    if(RegOpenKeyExW(HKEY_LOCAL_MACHINE, PSUBST_REG_PATH, 0, KEY_SET_VALUE, &hKey) != ERROR_SUCCESS) {
        return -1;
    }
    _wcsupr(DrvStr);
    len = _snwprintf(subst_name, MAX_PATH*2, L"\\??\\%s\0", TargetStr);
    if(RegSetValueExW(hKey, DrvStr, 0, REG_SZ, (PUCHAR)subst_name, (len+1)*sizeof(WCHAR)) != ERROR_SUCCESS) {
        RegCloseKey(hKey);
        return -1;
    }
    RegCloseKey(hKey);

    return 0;
}

DWORD
DeletePSubst(
    PWCHAR DrvStr
    )
{
    HKEY hKey;

    if(RegOpenKeyExW(HKEY_LOCAL_MACHINE, PSUBST_REG_PATH, 0, KEY_SET_VALUE, &hKey) != ERROR_SUCCESS) {
        return -1;
    }
    _wcsupr(DrvStr);
    if(RegDeleteValueW(hKey, DrvStr) != ERROR_SUCCESS) {
        RegCloseKey(hKey);
        return -1;
    }
    RegCloseKey(hKey);

    return 0;
}

int
main (void)
{
    int argc;
    WCHAR** argv;
    WCHAR*  CmdLine;
    BOOLEAN force = FALSE;

    PWCHAR DrvStr;
    PWCHAR TargetStr;

    ULONG RetCode = 0;
    WCHAR cmd_str[MAX_PATH*2];

    CmdLine = GetCommandLineW();
    argv = CommandLineToArgvW(CmdLine, &argc);

    if ( argc < 2 ) {
        usage();
        ExitProcess(1);
    }
    if(!wcscmp(argv[1], L"/?") ||
       !wcsicmp(argv[1], L"/h")) {
        usage();
        ExitProcess(1);
    }
    if(!wcsicmp(argv[1], L"/D") ||
       (force = !wcsicmp(argv[1], L"/DF"))) {
        // delete subst
        if ( argc != 3 ) {
            usage();
            ExitProcess(1);
        }
        DrvStr = FixSlashes(argv[2], TRUE);
        _snwprintf(cmd_str, MAX_PATH*2, L"subst /D %s\0", DrvStr);
        RetCode = LauncherRoutine(cmd_str);
        if(!RetCode || force) {
            if(!DeletePSubst(DrvStr) && RetCode) {
                printf("Changes will take effect after reboot\n");
            }
        }
        GlobalFree(DrvStr);
    } else
    if(!wcsicmp(argv[1], L"/P")) {
        // create persistent subst
        if ( argc != 4 ) {
            usage();
            ExitProcess(1);
        }
        DrvStr    = FixSlashes(argv[2], TRUE);
        TargetStr = FixSlashes(argv[3], FALSE);
        _snwprintf(cmd_str, MAX_PATH*2, L"subst %s \"%s\"\0", DrvStr, TargetStr);
        //printf("Exec: %ws\n", cmd_str);
        RetCode = LauncherRoutine(cmd_str);
        if(!RetCode) {
            CreatePSubst(DrvStr, TargetStr);
        }
        GlobalFree(DrvStr);
        GlobalFree(TargetStr);
    } else {
        // create subst
        if ( argc != 3 ) {
            usage();
            ExitProcess(1);
        }
        DrvStr    = FixSlashes(argv[1], TRUE);
        TargetStr = FixSlashes(argv[2], FALSE);
        _snwprintf(cmd_str, MAX_PATH*2, L"subst %s \"%s\"\0", DrvStr, TargetStr);
        RetCode = LauncherRoutine(cmd_str);
        GlobalFree(DrvStr);
        GlobalFree(TargetStr);
    }
    ExitProcess(RetCode);
    // make compile happy
    return 0;
}