// com.c
// Copyright : 2021-11-17 Yutaka Sawada
// License : GPL

#ifndef _UNICODE
#define _UNICODE
#endif
#ifndef UNICODE
#define UNICODE
#endif
#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0600	// Windows Vista or later
#endif

#include <windows.h>
#include <Wbemidl.h>
#include <comdef.h>

#pragma comment(lib, "wbemuuid.lib")

extern "C" {	// C 言語の関数呼び出しに対応する

// Detecting SSD in Windows
// https://stackoverflow.com/questions/23363115/detecting-ssd-in-windows

int check_msft(void)
{
	HRESULT hres;
	IWbemLocator *wbemLocator = NULL;
	IWbemServices *wbemServices = NULL;

	hres = CoInitializeEx(0, COINIT_MULTITHREADED);
	if (FAILED(hres)){
		return 1;
	}

	// Step 2: --------------------------------------------------
	// Set general COM security levels --------------------------
	hres = CoInitializeSecurity(
		NULL,
		-1,                          // COM authentication
		NULL,                        // Authentication services
		NULL,                        // Reserved
		RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication 
		RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation  
		NULL,                        // Authentication info
		EOAC_NONE,                   // Additional capabilities 
		NULL                         // Reserved
	);
	if (FAILED(hres)){
		CoUninitialize();
		return 2;
	}

	// Step 3: ---------------------------------------------------
	// Obtain the initial locator to WMI -------------------------
	hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID *)&wbemLocator);
	if (FAILED(hres)){
		CoUninitialize();
		return 3;
	}

	// Step 4: -----------------------------------------------------
	// Connect to WMI through the IWbemLocator::ConnectServer method

	// Connect to the ROOT\\\microsoft\\windows\\storage namespace with
	// the current user and obtain pointer pSvc
	// to make IWbemServices calls.
	hres = wbemLocator->ConnectServer(
		_bstr_t(L"ROOT\\microsoft\\windows\\storage"), // Object path of WMI namespace
		NULL,                    // User name. NULL = current user
		NULL,                    // User password. NULL = current
		0,                       // Locale. NULL indicates current
		NULL,                    // Security flags.
		0,                       // Authority (for example, Kerberos)
		0,                       // Context object 
		&wbemServices            // pointer to IWbemServices proxy
	);
	if (FAILED(hres)){
		wbemLocator->Release();
		CoUninitialize();
		return 4;
	}

	// Step 5: --------------------------------------------------
	// Set security levels on the proxy -------------------------
	hres = CoSetProxyBlanket(
		wbemServices,                // Indicates the proxy to set
		RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
		RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
		NULL,                        // Server principal name 
		RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx 
		RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
		NULL,                        // client identity
		EOAC_NONE                    // proxy capabilities 
	);
	if (FAILED(hres)){
		wbemServices->Release();
		wbemLocator->Release();
		CoUninitialize();
		return 5;
	}

	IEnumWbemClassObject* storageEnumerator = NULL;
	hres = wbemServices->ExecQuery(
		bstr_t("WQL"),
		bstr_t("SELECT * FROM MSFT_PhysicalDisk"),
		WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
		NULL,
		&storageEnumerator
	);
	if (FAILED(hres)){
		wbemServices->Release();
		wbemLocator->Release();
		CoUninitialize();
		return 6;
	}

	IWbemClassObject *storageWbemObject = NULL;
	ULONG uReturn = 0;

	while (storageEnumerator){
		HRESULT hr = storageEnumerator->Next(WBEM_INFINITE, 1, &storageWbemObject, &uReturn);
		if (0 == uReturn || hr != S_OK){
			break;
		}


		VARIANT deviceId;
		VARIANT friendlyName;
		VARIANT busType;		// 7 = USB, 8 = RAID, 11 = SATA, 17 = NVMe
		VARIANT healthStatus;	// 0 = Healthy, 1 = Warning, 2 = Unhealthy
		VARIANT spindleSpeed;	// 0 = SSD
		VARIANT mediaType;		// 3 = HDD, 4 = SSD, 5 = SCM

		storageWbemObject->Get(L"DeviceId", 0, &deviceId, 0, 0);
		storageWbemObject->Get(L"FriendlyName", 0, &friendlyName, 0, 0);
		storageWbemObject->Get(L"BusType", 0, &busType, 0, 0);
		storageWbemObject->Get(L"HealthStatus", 0, &healthStatus, 0, 0);
		storageWbemObject->Get(L"SpindleSpeed", 0, &spindleSpeed, 0, 0);
		storageWbemObject->Get(L"MediaType", 0, &mediaType, 0, 0);

/*
        storageDevice.DeviceId = deviceId.bstrVal == NULL ? "" : _bstr_t(deviceId.bstrVal);
        storageDevice.BusType = busType.uintVal;
        storageDevice.HealthStatus = healthStatus.uintVal;
        storageDevice.SpindleSpeed = spindleSpeed.uintVal;
        storageDevice.MediaType = mediaType.uintVal;

*/
		if (deviceId.bstrVal != NULL)
			printf("DeviceId = %S\n", (LPCTSTR)_bstr_t(deviceId.bstrVal));
		if (friendlyName.bstrVal != NULL)
			printf("FriendlyName = %S\n", (LPCTSTR)_bstr_t(friendlyName.bstrVal));
		printf("BusType = %d\n", busType.uintVal);
		printf("HealthStatus = %d\n", healthStatus.uintVal);
		printf("SpindleSpeed = %d\n", spindleSpeed.uintVal);
		printf("MediaType = %d\n", mediaType.uintVal);

		storageWbemObject->Release();
	}



	wbemServices->Release();
	wbemLocator->Release();
	CoUninitialize();
	return 0;
}

}	// end of extern "C"

