DeviceMonitor.cpp

#include "DeviceMonitor.h"
#include "DeviceUtils.h"
#include "DiskMonitor.h"
#include "Shlwapi.h"
#include <algorithm>
#include <stdexcept>
#include <winioctl.h>

using namespace std;
using namespace DeviceUtils;
#pragma comment(lib, "Shlwapi.lib")

DeviceMonitor::DeviceMonitor()
	: m_stopWork(FALSE),
	m_windowManager(make_unique<WindowManager>(this)) {
}

DeviceMonitor::~DeviceMonitor() {
	StopMonitor();
}

bool DeviceMonitor::StartMonitor(std::wstring driveLetters, SpecifiedMode specifiedMode) {
	m_stopWork = FALSE;
	m_specifiedMode = specifiedMode;
	try {
		m_windowManager->InitDeviceNotification();
		wprintf(L"设备通知初始化成功\n");
	}
	catch (const exception& e) {
		wprintf(L"初始化设备通知失败: %S\n", e.what());
		return false;
	}
	WORD oldErrorMode = SetErrorMode(SEM_FAILCRITICALERRORS);

	if (specifiedMode == SM_ALL) {
		m_specifiedMode = SM_EXCLUDE;
		m_specifiedDisk = L""; // 排除空列表=监控所有
		wprintf(L"模式：监控所有可用磁盘\n");
	}
	else {
		// 清理输入盘符
		m_specifiedDisk.clear();
		for (wchar_t c : driveLetters) {
			wchar_t upper = towupper(c);
			if (upper >= L'A' && upper <= L'Z') {
				m_specifiedDisk.push_back(upper);
			}
		}
		wprintf(L"模式：%s，目标盘符：%s\n",
			specifiedMode == SM_INCLUDE ? L"包含" : L"排除",
			m_specifiedDisk.c_str());
	}

	// 获取当前所有盘符
	wstring currentDrives = GetCurrentDrives();
	wprintf(L"当前可用盘符：%s\n", currentDrives.empty() ? L"无" : currentDrives.c_str());

	// 根据模式计算需要监控的目标盘符
	wstring targetDrives;
	switch (m_specifiedMode) {
	case SM_INCLUDE:
		// 仅监控指定盘符中存在的部分
		for (wchar_t c : m_specifiedDisk) {
			if (currentDrives.find(c) != wstring::npos) {
				targetDrives.push_back(c);
			}
		}
		break;
	case SM_EXCLUDE:
		// 监控当前盘符中不在排除列表的部分
		for (wchar_t c : currentDrives) {
			if (m_specifiedDisk.find(c) == wstring::npos) {
				targetDrives.push_back(c);
			}
		}
		break;
	}

	// 初始化目标盘符的监控
	if (targetDrives.empty()) {
		wprintf(L"无符合条件的磁盘可监控\n");
		SetErrorMode(oldErrorMode);
		return false;
	}
	for (wchar_t c : targetDrives) {
		if (!InitSingleDiskMonitoring(c)) {
			wprintf(L"警告：无法监控磁盘 %c\n", c);
		}
	}

	SetErrorMode(oldErrorMode);

	// 检查是否有成功监控的设备
	bool hasMonitored = false;
	{
		lock_guard<mutex> lock(m_devicesMutex);
		for (const auto& device : m_devices) {
			if (device && device->IsMonitoring()) {
				hasMonitored = true;
				break;
			}
		}
	}
	if (!hasMonitored) {
		wprintf(L"没有可监控的磁盘（请检查磁盘是否正常挂载）\n");
		return false;
	}

	wprintf(L"\n等待设备事件...（输入q停止监控）\n");
	return true;
}

void DeviceMonitor::StopMonitor() {
	if (m_stopWork) return;
	m_stopWork = TRUE;

	UninitAllDisksMonitoring();
	wprintf(L"所有监控已停止\n");
}

void DeviceMonitor::InitAllDisksMonitoring() {
	auto currentDrives = GetCurrentDrives();
	WORD oldErrorMode = SetErrorMode(SEM_FAILCRITICALERRORS);

	for (wchar_t c : currentDrives) {
		InitSingleDiskMonitoring(c);
	}

	SetErrorMode(oldErrorMode);
}

bool DeviceMonitor::InitSingleDiskMonitoring(wchar_t driveLetter) {
	wstring diskPath = RootPath(driveLetter);
	wprintf(L"\n");
	// 检查磁盘是否已监控
	if (IsDiskMonitored(driveLetter)) {
		wprintf(L"磁盘 %c 已在监控中\n", driveLetter);
		return true;
	}

	// 检查路径有效性
	if (!PathIsDirectoryW(diskPath.c_str())) {
		wprintf(L"路径 %s 不是有效的目录\n", diskPath.c_str());
		return false;
	}

	bool isRemovable = DeviceUtils::IsRemovableDrive(driveLetter);
	UINT sysDriveType = GetDriveType(diskPath.c_str());  // 系统定义的驱动器类型
	wprintf(L"DeviceMonitor: 磁盘 %s 类型信息\n", diskPath.c_str());
	wprintf(L"  - 系统定义类型: %d USB设备: %s\n", sysDriveType, isRemovable ? L"是" : L"否");

	// 创建设备监控实例
	auto device = make_shared<DiskMonitor>();
	if (!device->FromDeviceLetter(driveLetter)) {
		wprintf(L"初始化设备监控实例失败: %s\n", diskPath.c_str());
		return false;
	}

	// 同时设置两种类型到DiskMonitor
	device->SetDriveType(sysDriveType);         // 系统定义类型
	device->SetIsRemovable(isRemovable);

	// 获取并设置设备实例ID
	device->SetInstanceId(DeviceUtils::GetInstanceId(diskPath));
	wprintf(L"  - 设备实例ID: %s\n", device->GetInstanceId().c_str());

	// 记录插入时间
	FILETIME insertionTime;
	GetSystemTimeAsFileTime(&insertionTime);
	device->SetInsertionTime(insertionTime);

	// 为可移动设备注册通知
	HDEVNOTIFY hNotify = NULL;
	if (isRemovable) {
		hNotify = RegisterRemovableDeviceNotification(driveLetter, device);
		device->SetNotifyHandle(hNotify);
	}

	// 添加到设备列表
	{
		lock_guard<mutex> lock(m_devicesMutex);
		m_devices.push_back(device);
	}

	// 启动监控
	return device->StartMonitoring();
}

void DeviceMonitor::UninitAllDisksMonitoring() {
	vector<shared_ptr<DiskMonitor>> devicesCopy;
	{
		lock_guard<mutex> lock(m_devicesMutex);
		devicesCopy = m_devices;
	}

	for (const auto& device : devicesCopy) {
		UninitSingleDevice(device, true);
	}

	lock_guard<mutex> lock(m_devicesMutex);
	m_devices.clear();
}

void DeviceMonitor::UninitSingleDevice(const shared_ptr<DiskMonitor>& device, bool removeFromList) {
	if (!device) return;

	device->StopMonitoring();
	HDEVNOTIFY hNotify = device->GetNotifyHandle();
	if (hNotify) {
		::UnregisterDeviceNotification(hNotify);
		device->SetNotifyHandle(NULL);
	}

	HANDLE hDevice = device->GetDeviceHandle();
	if (hDevice != INVALID_HANDLE_VALUE) {
		CloseHandle(hDevice);
		device->SetDeviceHandle(INVALID_HANDLE_VALUE);
	}

	// 从列表移除
	if (removeFromList) {
		auto it = find_if(m_devices.begin(), m_devices.end(),
			[&](const shared_ptr<DiskMonitor>& d) {
				return d->GetDriveLetter() == device->GetDriveLetter();
			});
		if (it != m_devices.end()) {
			m_devices.erase(it);
		}
	}
}

// 判断指定盘符是否需要监控
bool DeviceMonitor::NeedMonitor(wchar_t driveLetter) {
	// 统一转换为大写字母，确保判断一致性
	wchar_t upperLetter = towupper(driveLetter);
	if (upperLetter < L'A' || upperLetter > L'Z') {
		wprintf(L"无效盘符: %c，无需监控\n", driveLetter);
		return false;
	}

	// 根据当前模式判断是否需要监控
	switch (m_specifiedMode) {
	case SM_ALL:
		return true;

	case SM_INCLUDE:
		return (m_specifiedDisk.find(upperLetter) != wstring::npos);

	case SM_EXCLUDE:
		return (m_specifiedDisk.find(upperLetter) == wstring::npos);

	default:
		wprintf(L"未知模式，默认不监控盘符: %c\n", upperLetter);
		return false;
	}
}

wstring DeviceMonitor::GetCurrentDrives() {
	wstring driveLetters;
	DWORD logicalDrives = GetLogicalDrives();
	if (logicalDrives == 0) return driveLetters;

	// 遍历A-Z盘符，拼接存在的盘符字母
	for (wchar_t letter = L'A'; letter <= L'Z'; ++letter) {
		if (logicalDrives & (1 << (letter - L'A'))) {
			wstring drivePath = DeviceUtils::RootPath(letter);
			if (PathFileExistsW(drivePath.c_str())) {
				driveLetters.push_back(letter);
			}
		}
	}
	return driveLetters;
}

bool DeviceMonitor::IsDiskMonitored(wchar_t driveLetter) {
	lock_guard<mutex> lock(m_devicesMutex);
	for (const auto& device : m_devices) {
		if (device->GetDriveLetter() == driveLetter && device->IsMonitoring()) {
			return true;
		}
	}
	return false;
}

HDEVNOTIFY DeviceMonitor::RegisterRemovableDeviceNotification(wchar_t driveLetter, shared_ptr<DiskMonitor> device) {
	// 仅为可移动设备注册额外通知
	HANDLE hDevice = CreateFile(
		DeviceUtils::RootPath(driveLetter).c_str(),
		GENERIC_READ,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		FILE_FLAG_BACKUP_SEMANTICS | FILE_ATTRIBUTE_NORMAL,
		NULL
	);

	if (hDevice == INVALID_HANDLE_VALUE) {
		wprintf(L"打开设备失败，错误码: %d\n", GetLastError());
		return NULL;
	}

	// 保存设备句柄到DiskMonitor实例
	device->SetDeviceHandle(hDevice);

	DEV_BROADCAST_HANDLE notifyFilter = { 0 };
	notifyFilter.dbch_size = sizeof(DEV_BROADCAST_HANDLE);
	notifyFilter.dbch_devicetype = DBT_DEVTYP_HANDLE;
	notifyFilter.dbch_handle = hDevice;

	HDEVNOTIFY hNotify = RegisterDeviceNotification(
		m_windowManager->GetHwnd(), // 使用窗口管理器的窗口句柄
		&notifyFilter,
		DEVICE_NOTIFY_WINDOW_HANDLE
	);

	if (!hNotify) {
		wprintf(L"注册设备通知失败，错误码: %d\n", GetLastError());
		// 如果注册失败，需要关闭已经打开的句柄
		CloseHandle(hDevice);
		device->SetDeviceHandle(INVALID_HANDLE_VALUE);
	}

	return hNotify;
}

BOOL DeviceMonitor::HandleDeviceChange(UINT eventType, LPARAM lParam) {
	if (eventType == DBT_DEVICEARRIVAL) return OnDeviceArrival(lParam);
	else if (eventType == DBT_DEVICEQUERYREMOVE) return OnDeviceQueryRemove(lParam);
	else if (eventType == DBT_DEVICEREMOVECOMPLETE) return OnDeviceRemoved(lParam);
	return TRUE;
}

BOOL DeviceMonitor::OnDeviceArrival(LPARAM lParam) {
	PDEV_BROADCAST_DEVICEINTERFACE pInf = (PDEV_BROADCAST_DEVICEINTERFACE)lParam;
	std::wstring drive = DeviceUtils::FromInstanceId(pInf->dbcc_name);

	wprintf(L"\n");
	wprintf(L"===== 检测到设备插入 =====\n");
	wprintf(L"dbcc_name: %s\n", pInf->dbcc_name);

	if (!drive.empty()) {
		for (wchar_t driveLetter : drive) {
			if (NeedMonitor(driveLetter)) {
				InitSingleDiskMonitoring(driveLetter);
			}
		}
	}

	// 获取当前所有可用盘符
	auto currentDrives = GetCurrentDrives();
	if (currentDrives.empty()) {
		wprintf(L"未检测到任何可用盘符\n");
		return TRUE;
	}

	// 获取当前已监控的盘符
	wstring monitoredDrives;
	{
		lock_guard<mutex> lock(m_devicesMutex);
		for (const auto& device : m_devices) {
			monitoredDrives += device->GetDriveLetter();
		}
	}

	// 对比找出新插入的盘符
	wstring newDrives;
	for (const auto& curr : currentDrives) {
		bool isNew = true;
		for (const auto& monitored : monitoredDrives) {
			if (curr == monitored) {
				isNew = false;
				break;
			}
		}
		if (isNew) {
			newDrives.push_back(curr);
		}
	}

	// 处理新发现的盘符
	if (newDrives.empty()) {
		return TRUE;
	}

	for (wchar_t newDrive : newDrives) {
		if (NeedMonitor(newDrive)) {
			InitSingleDiskMonitoring(newDrive);
		}
	}

	return TRUE;
}

BOOL DeviceMonitor::OnDeviceQueryRemove(LPARAM lParam) {
	wprintf(L"\n===== 检测到设备弹出请求 =====\n");
	// 停止对应设备的监控
	auto pDevHandle = (PDEV_BROADCAST_HANDLE)lParam;

	lock_guard<mutex> lock(m_devicesMutex);
	for (const auto& device : m_devices) {
		if (device->GetDeviceHandle() == pDevHandle->dbch_handle) {
			UninitSingleDevice(device, false);
		}
	}
	return TRUE;
}

BOOL DeviceMonitor::OnDeviceRemoved(LPARAM lParam) {
	PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR)lParam;
	if (!pHdr) {
		return TRUE;
	}

	bool handled = false;
	vector<shared_ptr<DiskMonitor>> devicesToRemove;

	auto currentDrives = GetCurrentDrives();

	if (pHdr->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE) {
		wprintf(L"\n===== 检测到设备移除 =====");
		// 查找已知的已移除设备
		auto pInf = (PDEV_BROADCAST_DEVICEINTERFACE)lParam;
		auto instanceId = (wstring)L"\\\\?\\" + DeviceUtils::ParseUsbStoragePath(pInf->dbcc_name);

		shared_ptr<DiskMonitor> removedDevice;

		{
			lock_guard<mutex> lock(m_devicesMutex);
			for (const auto& device : m_devices) {
				if (device->GetInstanceId() == instanceId) {
					removedDevice = device;
					break;
				}
			}
			if (removedDevice) {
				wprintf(L"\n已移除设备: %c（可移动设备: %s）",
					removedDevice->GetDriveLetter(),
					removedDevice->IsRemovable() ? L"是" : L"否");
				UninitSingleDevice(removedDevice, true);
				handled = true;
			}
			for (const auto& device : m_devices) {
				bool exists = false;
				for (wchar_t c : currentDrives) {
					if (device->GetDriveLetter() == c) {
						exists = true;
						break;
					}
				}
				if (!exists) {
					devicesToRemove.push_back(device);
				}
			}
			if (!devicesToRemove.empty()) {
				for (const auto& device : devicesToRemove) {
					wprintf(L"\n已移除设备: %c（可移动设备: %s）",
						device->GetDriveLetter(),
						device->IsRemovable() ? L"是" : L"否");

					UninitSingleDevice(device, true);
					handled = true;
				}
			}
		}

	}

	if (handled) {
		wprintf(L"\n");
	}
	return TRUE;
}

shared_ptr<DiskMonitor> DeviceMonitor::FindDeviceByHandleInDev(PDEV_BROADCAST_HANDLE pDevHandle) {
	lock_guard<mutex> lock(m_devicesMutex);
	for (const auto& device : m_devices) {
		if (device->GetDevHandle() && device->GetDevHandle()->dbch_handle == pDevHandle->dbch_handle) {
			return device;
		}
	}
	return nullptr;
}

wstring DeviceMonitor::GetFixedDrives() {
	wstring result;
	lock_guard<mutex> lock(m_devicesMutex);
	for (const auto& device : m_devices) {
		if (!device->IsRemovable()) {
			result += device->GetDriveLetter();
		}
	}
	return result;
}

wstring DeviceMonitor::GetRemovableDrives() {
	wstring result;
	lock_guard<mutex> lock(m_devicesMutex);
	for (const auto& device : m_devices) {
		if (device->IsRemovable()) {
			result += device->GetDriveLetter();
		}
	}
	return result;
}

DWORD WINAPI DeviceMonitor::AsyncPostQueryRemoveHandler(LPVOID lpParam) {
	DeviceMonitor* pThis = static_cast<DeviceMonitor*>(lpParam);
	if (!pThis) return 1;

	wprintf(L"等待1000ms后重新读取盘符并初始化监控...\n");
	Sleep(1000);

	// 获取当前可用盘符并重新初始化
	auto currentDrives = pThis->GetCurrentDrives();
	wprintf(L"重新读取盘符完成，当前有效盘符: %d 个\n", (int)currentDrives.size());

	if (!currentDrives.empty()) {
		wprintf(L"开始重新初始化磁盘监控...\n");
		pThis->InitAllDisksMonitoring();

		bool hasMonitored = false;
		{
			lock_guard<mutex> lock(pThis->m_devicesMutex);
			for (const auto& device : pThis->m_devices) {
				if (device->IsMonitoring()) {
					hasMonitored = true;
					break;
				}
			}
		}

		if (hasMonitored) {
			wprintf(L"磁盘监控重新初始化完成\n");
		}
		else {
			wprintf(L"重新初始化后没有可监控的磁盘\n");
		}
	}
	else {
		wprintf(L"没有检测到可用磁盘，无需重新初始化\n");
	}

	return 0;
}