#include "libEncrypt_Hardware.h"
#include "libEncrypt_MD5.h"

#include <Windows.h>
#include <strsafe.h>

#include <intrin.h>
#include <string>
#include <vector>

//get cpu id
namespace 
{

std::string getCPUID() 
{
    int cpuInfo[4] = { -1 };
    char cpuId[32] = { 0 };

    __cpuid(cpuInfo, 1);
    memset(cpuId, 0, sizeof(cpuId));
    StringCbPrintfA(cpuId, 32, "%08X%08X", cpuInfo[3], cpuInfo[0]);
    return std::string(cpuId);
}

std::string GetDiskSerial(void) 
{
    std::string result;
    HANDLE hDevice = CreateFileW(L"\\\\.\\PhysicalDrive0", 0, FILE_SHARE_READ | FILE_SHARE_WRITE,
        NULL, OPEN_EXISTING, 0, NULL);

    if (hDevice == INVALID_HANDLE_VALUE) 
    {
        return result;
    }

    STORAGE_PROPERTY_QUERY query;
	SecureZeroMemory(&query, sizeof(query));
    query.PropertyId = StorageDeviceProperty;
    query.QueryType = PropertyStandardQuery;

    STORAGE_DESCRIPTOR_HEADER header;
    SecureZeroMemory(&header, sizeof(header));

    DWORD bytesReturned = 0;

    if (DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(query),
        &header, sizeof(header), &bytesReturned, NULL)) 
    {
        std::vector<char> buffer(header.Size);

        if (DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY,
            &query, sizeof(query), buffer.data(), header.Size, &bytesReturned, NULL)) 
        {
            STORAGE_DEVICE_DESCRIPTOR* deviceDescriptor = reinterpret_cast<STORAGE_DEVICE_DESCRIPTOR*>(buffer.data());
            if (deviceDescriptor->SerialNumberOffset) 
            {
                result = buffer.data() + deviceDescriptor->SerialNumberOffset;
            }
        }
    }

    CloseHandle(hDevice);
    return result;
}

}

std::array<uint8_t, 16> getHardwareID()
{
    std::string hardwareInfo = getCPUID() + GetDiskSerial();// +GetBaseboardSerial() + GetMACAddress();

    MD5Context md5Context;
	md5Init(&md5Context);
	md5Update(&md5Context, (uint8_t*)hardwareInfo.c_str(), (uint32_t)hardwareInfo.size());
	MD5Digest digest = md5Finalize(&md5Context);

	return digest;
}