Windows下映射文件夹为虚拟磁盘

Windows下映射文件夹为虚拟磁盘

/* INCLUDES *****************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <windows.h>
#include <tchar.h>

/* FUNCTIONS ****************************************************************/

void PrintError(DWORD ErrCode)
{
    TCHAR *buffer = (TCHAR*) calloc(2048,
                                    sizeof(TCHAR));
    TCHAR *msg = NULL;

    if (buffer)
    {
        FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
                      NULL,
                      ErrCode,
                      MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      (TCHAR*)&msg,
                      0,
                      NULL);
        _sntprintf(buffer,
                   2048,
                   _T("操作失败！错误码 0x%x：%s\n"),
                   ErrCode,
                   msg);
        _tprintf(_T("%s"),
                 buffer);
        if (msg)
            LocalFree(msg);
        free(buffer);
    }
}

void DisplaySubstUsage(void)
{
    _tprintf(_T("%s"), _T("将一个路径用驱动器号关联。\n\n\
SUBST [驱动器1: [驱动器2:]路径]\n\
SUBST 驱动器1: /D\n\n\
  驱动器1:          指定一个您想要关联到一个路径的虚拟驱动器。\n\
  [驱动器2:]路径    指定一个您想要关联到一个虚拟驱动器的物理驱动器和路径。\n\
  /D                删除一个虚拟驱动器。\n\n\
执行不带参数的 SUBST 将会显示当前虚拟驱动器的列表。\n"));
}

BOOLEAN IsSubstedDrive(TCHAR *Drive)
{
    BOOLEAN Result = FALSE;
    LPTSTR lpTargetPath = NULL;
    DWORD CharCount, dwSize;

    if (_tcslen(Drive) > 2)
        return FALSE;

    dwSize = sizeof(TCHAR) * MAX_PATH;
    lpTargetPath = (LPTSTR) malloc(sizeof(TCHAR) * MAX_PATH);
    if ( lpTargetPath)
    {
        CharCount = QueryDosDevice(Drive,
                                   lpTargetPath,
                                   dwSize / sizeof(TCHAR));
        while (! CharCount &&
               GetLastError() == ERROR_INSUFFICIENT_BUFFER)
        {
            free(lpTargetPath);
            dwSize *= 2;
            lpTargetPath = (LPTSTR) malloc(dwSize);
            if (lpTargetPath)
            {
                CharCount = QueryDosDevice(Drive,
                                           lpTargetPath,
                                           dwSize / sizeof(TCHAR));
            }
        }

        if (CharCount)
        {
            if ( _tcsncmp(lpTargetPath, _T("\\??\\"), 4) == 0 &&
                ( (lpTargetPath[4] >= _T('A') &&
                lpTargetPath[4] <= _T('Z')) ||
                 (lpTargetPath[4] >= _T('a') &&
                lpTargetPath[4] <= _T('z')) ) )
				Result = TRUE;
        }
        free(lpTargetPath);
    }
    return Result;
}

void DumpSubstedDrives(void)
{
    TCHAR Drive[3] = _T("A:");
    LPTSTR lpTargetPath = NULL;
    DWORD CharCount, dwSize;
    INT i = 0;

    dwSize = sizeof(TCHAR) * MAX_PATH;
    lpTargetPath = (LPTSTR) malloc(sizeof(TCHAR) * MAX_PATH);
    if (! lpTargetPath)
        return;

    while (i < 26)
    {
        Drive[0] = _T('A') + i;
        CharCount = QueryDosDevice(Drive,
                                   lpTargetPath,
                                   dwSize / sizeof(TCHAR));
        while (! CharCount &&
               GetLastError() == ERROR_INSUFFICIENT_BUFFER)
        {
            free(lpTargetPath);
            dwSize *= 2;
            lpTargetPath = (LPTSTR) malloc(dwSize);
            if (lpTargetPath)
                CharCount = QueryDosDevice(Drive,
                                           lpTargetPath,
                                           dwSize / sizeof(TCHAR));
        }

        if (! CharCount)
        {
            i++;
            continue;
        }
        else
        {
            if ( _tcsncmp(lpTargetPath, _T("\\??\\"), 4) == 0 &&
                ( (lpTargetPath[4] >= _T('A') &&
                lpTargetPath[4] <= _T('Z')) ||
                 (lpTargetPath[4] >= _T('a') &&
                lpTargetPath[4] <= _T('z')) ) )
            {
                _tprintf(_T("%s\\: => %s\n"),
                         Drive,
                         lpTargetPath + 4);
            }
        }
        i++;
    }
    free(lpTargetPath);
}

int DeleteSubst(TCHAR* Drive)
{
    BOOL Result;

    if (_tcslen(Drive) > 2)
    {
        _tprintf(_T("无效参数 - %s\n"),
                 Drive);
        return 1;
    }

    if (! IsSubstedDrive(Drive))
    {
        _tprintf(_T("无效参数 - %s\n"),
                Drive);
        return 1;
    }

    Result = DefineDosDevice(DDD_REMOVE_DEFINITION,
                             Drive,
                             NULL);
    if (! Result)
    {
        PrintError(GetLastError());
        return 1;
    }
    return 0;
}

int AddSubst(TCHAR* Drive, TCHAR *Path)
{
    BOOL Result;

    if (_tcslen(Drive) != 2)
    {
        _tprintf(_T("无效参数 - %s\n"),
                 Drive);
        return 1;
    }

    if (Drive[1] != _T(':'))
    {
        _tprintf(_T("无效参数 - %s\n"),
                 Drive);
        return 1;
    }

    if (IsSubstedDrive(Drive))
    {
        _tprintf(_T("驱动器已经运行过 SUBST\n"));
        return 1;
    }

    Result = DefineDosDevice(0,
                             Drive,
                             Path);
    if (! Result)
    {
        PrintError(GetLastError());
        return 1;
    }
    return 0;
}

int _tmain(int argc, TCHAR* argv[])
{

    for (INT i = 0; i < argc; i++)
    {
        if (!_tcsicmp(argv[i], _T("/?")))
        {
            DisplaySubstUsage();
            return 0;
        }
    }

    if (argc < 3)
    {
        if (argc >= 2)
        {
            _tprintf(_T("参数个数不正确 - %s\n"),
                     argv[1]);
            return 1;
        }
        DumpSubstedDrives();
        return 0;
    }

    if (argc > 3)
    {
        _tprintf(_T("参数个数不正确 - %s\n"),
                 argv[3]);
        return 1;
    }

    if (! _tcsicmp(argv[1], _T("/D")))
        return DeleteSubst(argv[2]);
    if (! _tcsicmp(argv[2], _T("/D")))
        return DeleteSubst(argv[1]);
    return AddSubst(argv[1], argv[2]);
}