hwid 0.1.0

use std::arch::x86_64::CpuidResult;

use hex::ToHex;
use hmac::{Hmac, Mac};
use sha2::{Digest, Sha256};
use smbios::{SMBiosBaseboardInformation, SMBiosStruct, SMBiosSystemInformation, SystemUuidData};
use windows::Win32::{
    Foundation::{ERROR_BUFFER_OVERFLOW, ERROR_SUCCESS, WIN32_ERROR},
    NetworkManagement::IpHelper::{
        GAA_FLAG_SKIP_ANYCAST, GAA_FLAG_SKIP_DNS_SERVER, GAA_FLAG_SKIP_FRIENDLY_NAME,
        GAA_FLAG_SKIP_MULTICAST, GetAdaptersAddresses, IP_ADAPTER_ADDRESSES_LH,
    },
    Networking::WinSock::AF_UNSPEC,
};

use crate::error::{HwidError, HwidResult};

pub mod error;
mod smbios;

const CPUID_LEAF_MANUFACTURER: u32 = 0;
const CPUID_LEAF_FEATURES: u32 = 1;

#[inline]
pub unsafe fn cpuid(leaf: u32, sub_leaf: u32) -> CpuidResult {
    let eax;
    let ebx;
    let ecx;
    let edx;

    unsafe {
        core::arch::asm!(
            "mov {0:r}, rbx",
            "cpuid",
            "xchg {0:r}, rbx",
            out(reg) ebx,
            inout("eax") leaf => eax,
            inout("ecx") sub_leaf => ecx,
            out("edx") edx,
            options(nostack, preserves_flags),
        );
    }

    CpuidResult { eax, ebx, ecx, edx }
}

pub trait HardwareIdentifierComponent {
    fn calculate(&self) -> HwidResult<Option<Vec<u8>>>;
}

pub struct ProcessorId;

impl HardwareIdentifierComponent for ProcessorId {
    /// **Note:** This is not a unique value. It represents the enabled CPU features. CPUID does not return a serial number
    /// on any modern processor, hence this value is used to represent the CPU model. It is similar to the WMI ProcessorId value.
    fn calculate(&self) -> HwidResult<Option<Vec<u8>>> {
        let processor_id = {
            let CpuidResult { eax, edx, .. } = unsafe { cpuid(CPUID_LEAF_FEATURES, 0) };

            let mut bytes = Vec::with_capacity(8);
            bytes.extend_from_slice(&eax.to_le_bytes());
            bytes.extend_from_slice(&edx.to_le_bytes());

            bytes
        };

        let manufacturer = {
            let CpuidResult { ebx, ecx, edx, .. } = unsafe { cpuid(CPUID_LEAF_MANUFACTURER, 0) };

            let mut bytes = Vec::with_capacity(12);
            bytes.extend_from_slice(&ebx.to_le_bytes());
            bytes.extend_from_slice(&edx.to_le_bytes());
            bytes.extend_from_slice(&ecx.to_le_bytes());

            bytes
        };

        let mut final_id = Vec::with_capacity(20);

        final_id.extend(processor_id);
        final_id.extend(manufacturer);

        Ok(Some(final_id))
    }
}

pub struct SmbiosUuid;

impl HardwareIdentifierComponent for SmbiosUuid {
    fn calculate(&self) -> HwidResult<Option<Vec<u8>>> {
        let smbios = smbios::table_load_from_device().map_err(HwidError::Smbios)?;
        let system_information = smbios
            .find_map(|table: SMBiosSystemInformation| Some(table))
            .ok_or(HwidError::MissingSmbiosTable(
                SMBiosSystemInformation::STRUCT_TYPE,
            ))?;

        let uuid = match system_information
            .uuid()
            .unwrap_or(SystemUuidData::IdNotPresent)
        {
            SystemUuidData::Uuid(uuid) => Some(uuid.raw.to_vec()),
            _ => None,
        };

        Ok(uuid)
    }
}

pub struct BaseboardSerial;

impl HardwareIdentifierComponent for BaseboardSerial {
    fn calculate(&self) -> HwidResult<Option<Vec<u8>>> {
        let smbios = smbios::table_load_from_device().map_err(HwidError::Smbios)?;
        let baseboard_information = smbios
            .find_map(|table: SMBiosBaseboardInformation| Some(table))
            .ok_or(HwidError::MissingSmbiosTable(
                SMBiosBaseboardInformation::STRUCT_TYPE,
            ))?;

        let serial_number = baseboard_information
            .serial_number()
            .ok()
            .ok_or(HwidError::SmbiosString)?;

        Ok(Some(serial_number.as_bytes().to_vec()))
    }
}

pub struct MacAddress;

impl HardwareIdentifierComponent for MacAddress {
    fn calculate(&self) -> HwidResult<Option<Vec<u8>>>
    where
        Self: Sized,
    {
        let mut buffer_length: u32 = 0;

        let error = unsafe {
            GetAdaptersAddresses(
                AF_UNSPEC.0 as u32,
                GAA_FLAG_SKIP_ANYCAST
                    | GAA_FLAG_SKIP_MULTICAST
                    | GAA_FLAG_SKIP_DNS_SERVER
                    | GAA_FLAG_SKIP_FRIENDLY_NAME,
                None,
                None,
                &mut buffer_length,
            )
        };

        // yes, if its ERROR_BUFFER_OVERFLOW that means it worked. don't ask, just windows api shit
        if error != ERROR_BUFFER_OVERFLOW.0 {
            return Err(windows::core::Error::from_hresult(WIN32_ERROR(error).to_hresult()).into());
        }

        let mut buffer = vec![0u8; buffer_length as usize];
        let adapter_addresses_ptr = buffer.as_mut_ptr() as *mut IP_ADAPTER_ADDRESSES_LH;

        let error = unsafe {
            GetAdaptersAddresses(
                AF_UNSPEC.0 as u32,
                GAA_FLAG_SKIP_ANYCAST
                    | GAA_FLAG_SKIP_MULTICAST
                    | GAA_FLAG_SKIP_DNS_SERVER
                    | GAA_FLAG_SKIP_FRIENDLY_NAME,
                None,
                Some(adapter_addresses_ptr),
                &mut buffer_length,
            )
        };

        if error != ERROR_SUCCESS.0 {
            return Err(windows::core::Error::from_hresult(WIN32_ERROR(error).to_hresult()).into());
        }

        let mut adapter = adapter_addresses_ptr;

        while !adapter.is_null() {
            unsafe {
                let addr = &*adapter;

                if addr.PhysicalAddressLength > 0 {
                    let mac = addr.PhysicalAddress[..addr.PhysicalAddressLength as usize].to_vec();

                    return Ok(Some(mac));
                }

                adapter = addr.Next;
            }
        }

        return Ok(None);
    }
}

type HmacSha256 = Hmac<Sha256>;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct HardwareIdentifier(Vec<u8>);

impl AsRef<[u8]> for HardwareIdentifier {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl AsMut<[u8]> for HardwareIdentifier {
    fn as_mut(&mut self) -> &mut [u8] {
        &mut self.0
    }
}

impl ToString for HardwareIdentifier {
    fn to_string(&self) -> String {
        self.encode_hex()
    }
}

pub struct HardwareIdentifierBuilder {
    components: Vec<Box<dyn HardwareIdentifierComponent>>,
}

impl HardwareIdentifierBuilder {
    /// Creates a new hardware identifier builder.
    ///
    /// # Examples
    /// ```
    /// let mut builder = HwidBuilder::new();
    /// ```
    pub fn new() -> Self {
        Self {
            components: Vec::new(),
        }
    }

    /// Adds a hardware identifier component to the builder.
    ///
    /// Each component represents a single piece of hardware-identifying data, and must implement
    /// the `HardwareIdentifierComponent` trait (for things like motherboard serial number, MAC, etc).
    ///
    /// Components are **stored in the order they are added**. If you change the order of the components, the
    /// resulting hardware identifier will have a different value.
    ///
    /// # Examples
    /// ```
    /// HwidBuilder::new()
    ///     .component(ProcessorId)
    ///     .component(SmbiosUuid)
    /// ```
    pub fn component<T>(&mut self, component: T) -> &mut Self
    where
        T: HardwareIdentifierComponent + 'static,
    {
        self.components.push(Box::new(component));
        self
    }

    /// Adds multiple (boxed) hardware identifier components to the builder.
    ///
    /// See [`HardwareIdentifierBuilder::component`] for information about ordering.
    ///
    /// # Examples
    /// ```
    /// HardwareIdentifierBuilder::new()
    ///     .components(vec![
    ///         Box::new(SmbiosUuid),
    ///         Box::new(ProcessorId),
    ///     ])
    /// ```
    pub fn components(
        &mut self,
        components: Vec<Box<dyn HardwareIdentifierComponent>>,
    ) -> &mut Self {
        self.components.extend(components);
        self
    }

    /// Builds the unhashed hardware identifier as a byte vector, concatenating the outputs of each component.
    ///
    /// # Errors
    /// Returns an error if any component's `calculate()` method returns an error.
    ///
    /// # Panics
    /// Panics if no components have been added to the builder.
    ///
    /// # Returns
    /// Returns a byte vector representing the concatenation of all component values, separated by a 0x00 byte.
    fn build_unhashed(&self) -> HwidResult<Vec<u8>> {
        assert!(
            self.components.len() > 0,
            "No components have been added to the hardware identifier."
        );

        let mut buffer = Vec::new();

        for (i, component) in self.components.iter().enumerate() {
            let part_opt = component.calculate()?;

            if let Some(part) = part_opt {
                if i > 0 {
                    buffer.push(0x00); // seperator
                }
                buffer.extend_from_slice(&part);
            }
        }

        Ok(buffer)
    }

    /// Builds the hashed hardware identifier as a [`HardwareIdentifier`] using SHA-256.
    ///
    /// # Errors
    /// Returns an error if one of the components cannot be calculated.
    ///
    /// # Panics
    /// Panics if no components have been added to the builder.
    pub fn build(&self) -> HwidResult<HardwareIdentifier> {
        let unhashed = self.build_unhashed()?;
        let digest = Sha256::digest(unhashed);

        Ok(HardwareIdentifier(digest.to_vec()))
    }

    /// Builds a signed hardware identifier as a [`HardwareIdentifier`] using HMAC-SHA256 and a provided secret.
    ///
    /// # Errors
    /// Returns an error if one of the components cannot be calculated or HMAC initialization fails.
    ///
    /// # Panics
    /// Panics if no components have been added to the builder.
    /// # Examples
    /// ```
    /// let hwid = builder.build_signed(b"my-secret")?;
    /// ```
    pub fn build_signed(&self, secret: &[u8]) -> HwidResult<HardwareIdentifier> {
        let unhashed = self.build_unhashed()?;
        let mut mac = HmacSha256::new_from_slice(secret).map_err(|_| HwidError::InitHmac)?;

        mac.update(&unhashed);

        let result = mac.finalize().into_bytes();

        Ok(HardwareIdentifier(result.to_vec()))
    }
}

/// Generates a signed hardware identifier string using the default components and settings.
///
/// See [`HardwareIdentifierBuilder::build_signed`]
///
/// # Returns
/// The string representation of a [`HardwareIdentifier`] created using [`HardwareIdentifier::to_string`]
pub fn get_hwid(hmac_secret: &[u8]) -> HwidResult<String> {
    let hwid = HardwareIdentifierBuilder::new()
        .components(vec![
            Box::new(ProcessorId),
            Box::new(SmbiosUuid),
            Box::new(BaseboardSerial),
            Box::new(MacAddress),
        ])
        .build_signed(hmac_secret)?
        .to_string();

    Ok(hwid)
}

#[cfg(test)]
mod tests {
    use crate::HardwareIdentifierComponent;

    #[test]
    fn processor_id_consistent() {
        let mut prev = crate::ProcessorId.calculate().unwrap();
        for _ in 0..10000 {
            let curr = crate::ProcessorId.calculate().unwrap();
            assert_eq!(curr, prev);
            prev = curr;
        }
    }
}