DiskMonitor/main.cpp at master · howdy213/DiskMonitor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#include "ConsoleManager.h"
#include "DeviceMonitor.h"
#include <algorithm>
#include <Windows.h>
#include <iostream>

using namespace std;

// 打印帮助信息
void PrintUsage() {
    wprintf(L"  监控所有磁盘:          DiskMonitor\n");
    wprintf(L"  包含指定磁盘:          DiskMonitor -include <盘符>\n");
    wprintf(L"  排除指定磁盘:          DiskMonitor -exclude <盘符>\n");
}

int main(int argc, char* argv[]) {
    try {
        DeviceMonitor monitor;
        ConsoleManager consoleManager(&monitor);
        consoleManager.SetMainThreadId(GetCurrentThreadId());
        consoleManager.InitConsole();

        // 解析宽字符命令行参数
        int wArgc = 0;
        LPWSTR* wArgv = CommandLineToArgvW(GetCommandLineW(), &wArgc);
        if (!wArgv) {
            wprintf(L"转换命令行参数失败，错误代码: %d\n", GetLastError());
            return 1;
        }

        // 默认为监控所有磁盘模式
        DeviceMonitor::SpecifiedMode specifiedMode = DeviceMonitor::SM_ALL;
        wstring driveLetters;

        // 解析参数
        if (wArgc > 1) {
            wstring modeArg = wArgv[1];
            transform(modeArg.begin(), modeArg.end(), modeArg.begin(), towlower);

            if (modeArg == L"-include") {
                specifiedMode = DeviceMonitor::SM_INCLUDE;
            }
            else if (modeArg == L"-exclude") {
                specifiedMode = DeviceMonitor::SM_EXCLUDE;
            }
            else {
                wprintf(L"错误: 无效选项 %s\n", modeArg.c_str());
                PrintUsage();
                LocalFree(wArgv);
                return 1;
            }

            if (wArgc > 2) driveLetters = wArgv[2];
            else {
                wprintf(L"错误: 选项需要指定盘符\n");
                PrintUsage();
                LocalFree(wArgv);
                return 1;
            }
        }

        // 启动控制台输入线程
        consoleManager.StartConsoleThread();

        // 启动监控
        if (!monitor.StartMonitor(driveLetters, specifiedMode)) {
            LocalFree(wArgv);
            return 1;
        }

        // 消息循环
        MSG msg;
        while (GetMessage(&msg, NULL, 0, 0) > 0) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }

        LocalFree(wArgv);
    }
    catch (const exception& e) {
        wprintf(L"程序发生错误: %S\n", e.what());
        return 1;
    }
    return 0;
}