/*

  This Source Code Form is subject to the terms of the Mozilla Public

  License, v. 2.0. If a copy of the MPL was not distributed with this

  file, You can obtain one at http://mozilla.org/MPL/2.0/.

  Copyright (C) 2009-2012 Michael Möller <mmoeller@openhardwaremonitor.org>

	Copyright (C) 2010 Paul Werelds

  Copyright (C) 2011 Roland Reinl <roland-reinl@gmx.de>

*/

using System;

using System.Collections.Generic;

using System.Management;

using System.Runtime.InteropServices;

namespace OpenHardwareMonitor.Hardware.HDD {

  internal class WindowsSmart : ISmart {

    [Flags]

    protected enum AccessMode : uint {     

      Read = 0x80000000,    

      Write = 0x40000000,     

      Execute = 0x20000000,     

      All = 0x10000000

    }

    [Flags]

    protected enum ShareMode : uint {

      None = 0,     

      Read = 1,     

      Write = 2,    

      Delete = 4

    }

    protected enum CreationMode : uint {

      New = 1,

      CreateAlways = 2,    

      OpenExisting = 3,    

      OpenAlways = 4,    

      TruncateExisting = 5

    }

    [Flags]

    protected enum FileAttribute : uint {

      Readonly = 0x00000001,

      Hidden = 0x00000002,

      System = 0x00000004,

      Directory = 0x00000010,

      Archive = 0x00000020,

      Device = 0x00000040,

      Normal = 0x00000080,

      Temporary = 0x00000100,

      SparseFile = 0x00000200,

      ReparsePoint = 0x00000400,

      Compressed = 0x00000800,

      Offline = 0x00001000,

      NotContentIndexed = 0x00002000,

      Encrypted = 0x00004000,

    }

    protected enum DriveCommand : uint {

      GetVersion = 0x00074080,

      SendDriveCommand = 0x0007c084,

      ReceiveDriveData = 0x0007c088

    }

    protected enum RegisterCommand : byte {

      /// <summary>

      /// SMART data requested.

      /// </summary>

      SmartCmd = 0xB0,

      /// <summary>

      /// Identify data is requested.

      /// </summary>

      IdCmd = 0xEC,

    }

    protected enum RegisterFeature : byte {

      /// <summary>

      /// Read SMART data.

      /// </summary>

      SmartReadData = 0xD0,

      /// <summary>

      /// Read SMART thresholds.

      /// </summary>

      SmartReadThresholds = 0xD1, /* obsolete */

      /// <summary>

      /// Autosave SMART data.

      /// </summary>

      SmartAutosave = 0xD2,

      /// <summary>

      /// Save SMART attributes.

      /// </summary>

      SmartSaveAttr = 0xD3,

      /// <summary>

      /// Set SMART to offline immediately.

      /// </summary>

      SmartImmediateOffline = 0xD4,

      /// <summary>

      /// Read SMART log.

      /// </summary>

      SmartReadLog = 0xD5,

      /// <summary>

      /// Write SMART log.

      /// </summary>

      SmartWriteLog = 0xD6,

      /// <summary>

      /// Write SMART thresholds.

      /// </summary>

      SmartWriteThresholds = 0xD7, /* obsolete */

      /// <summary>

      /// Enable SMART.

      /// </summary>

      SmartEnableOperations = 0xD8,

      /// <summary>

      /// Disable SMART.

      /// </summary>

      SmartDisableOperations = 0xD9,

      /// <summary>

      /// Get SMART status.

      /// </summary>

      SmartStatus = 0xDA,

      /// <summary>

      /// Set SMART to offline automatically.

      /// </summary>

      SmartAutoOffline = 0xDB, /* obsolete */

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct CommandBlockRegisters {

      public RegisterFeature Features;         

      public byte SectorCount;      

      public byte LBALow;       

      public byte LBAMid;           

      public byte LBAHigh;        

      public byte Device;

      public RegisterCommand Command;           

      public byte Reserved;                  

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriveCommandParameter {

      public uint BufferSize;           

      public CommandBlockRegisters Registers;           

      public byte DriveNumber;   

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]

      public byte[] Reserved;                                

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriverStatus {

      public byte DriverError;   

      public byte IDEError;             

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 10)]

      public byte[] Reserved;               

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriveCommandResult {

      public uint BufferSize;

      public DriverStatus DriverStatus;

    } 

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriveSmartReadDataResult {

      public uint BufferSize;           

      public DriverStatus DriverStatus;

      public byte Version;

      public byte Reserved;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = MAX_DRIVE_ATTRIBUTES)]

      public DriveAttributeValue[] Attributes;                                                                                       

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriveSmartReadThresholdsResult {

      public uint BufferSize;

      public DriverStatus DriverStatus;

      public byte Version;

      public byte Reserved;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = MAX_DRIVE_ATTRIBUTES)]

      public DriveThresholdValue[] Thresholds;

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct Identify {

      public ushort GeneralConfiguration;

      public ushort NumberOfCylinders;

      public ushort Reserved;

      public ushort NumberOfHeads;

      public ushort UnformattedBytesPerTrack;

      public ushort UnformattedBytesPerSector;

      public ushort SectorsPerTrack;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 3)]

      public ushort[] VendorUnique;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]

      public byte[] SerialNumber;

      public ushort BufferType;

      public ushort BufferSectorSize;

      public ushort NumberOfEccBytes;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 8)]

      public byte[] FirmwareRevision;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 40)]

      public byte[] ModelNumber;

      public ushort MoreVendorUnique;

      public ushort DoubleWordIo;

      public ushort Capabilities;

      public ushort MoreReserved;

      public ushort PioCycleTimingMode;

      public ushort DmaCycleTimingMode;

      [MarshalAs(UnmanagedType.ByValArray, SizeConst = 406)]

      public byte[] More;

    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]

    protected struct DriveIdentifyResult {

      public uint BufferSize;

      public DriverStatus DriverStatus;

      public Identify Identify;

    }

    public IntPtr InvalidHandle { get { return (IntPtr)(-1); } }

    private const byte SMART_LBA_MID = 0x4F;

    private const byte SMART_LBA_HI = 0xC2;

    private const int MAX_DRIVE_ATTRIBUTES = 512;

    public IntPtr OpenDrive(int driveNumber) {

      return NativeMethods.CreateFile(@"\\.\PhysicalDrive" + driveNumber,

        AccessMode.Read | AccessMode.Write, ShareMode.Read | ShareMode.Write,

        IntPtr.Zero, CreationMode.OpenExisting, FileAttribute.Device,

        IntPtr.Zero);

    }

    public bool EnableSmart(IntPtr handle, int driveNumber) {

      DriveCommandParameter parameter = new DriveCommandParameter();

      DriveCommandResult result;

      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;

      parameter.Registers.Features = RegisterFeature.SmartEnableOperations;

      parameter.Registers.LBAMid = SMART_LBA_MID;

      parameter.Registers.LBAHigh = SMART_LBA_HI;

      parameter.Registers.Command = RegisterCommand.SmartCmd;

      return NativeMethods.DeviceIoControl(handle, DriveCommand.SendDriveCommand, 

        ref parameter, Marshal.SizeOf(typeof(DriveCommandParameter)), out result,

        Marshal.SizeOf(typeof(DriveCommandResult)), out bytesReturned, 

        IntPtr.Zero);

    }

    public DriveAttributeValue[] ReadSmartData(IntPtr handle, int driveNumber) {

      DriveCommandParameter parameter = new DriveCommandParameter();

      DriveSmartReadDataResult result;

      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;

      parameter.Registers.Features = RegisterFeature.SmartReadData;

      parameter.Registers.LBAMid = SMART_LBA_MID;

      parameter.Registers.LBAHigh = SMART_LBA_HI;

      parameter.Registers.Command = RegisterCommand.SmartCmd;

      bool isValid = NativeMethods.DeviceIoControl(handle, 

        DriveCommand.ReceiveDriveData, ref parameter, Marshal.SizeOf(parameter), 

        out result, Marshal.SizeOf(typeof(DriveSmartReadDataResult)), 

        out bytesReturned, IntPtr.Zero);

      return (isValid) ? result.Attributes : new DriveAttributeValue[0];

    }

    public DriveThresholdValue[] ReadSmartThresholds(IntPtr handle,

      int driveNumber) 

    {

      DriveCommandParameter parameter = new DriveCommandParameter();

      DriveSmartReadThresholdsResult result;

      uint bytesReturned = 0;

      parameter.DriveNumber = (byte)driveNumber;

      parameter.Registers.Features = RegisterFeature.SmartReadThresholds;

      parameter.Registers.LBAMid = SMART_LBA_MID;

      parameter.Registers.LBAHigh = SMART_LBA_HI;

      parameter.Registers.Command = RegisterCommand.SmartCmd;

      bool isValid = NativeMethods.DeviceIoControl(handle,

        DriveCommand.ReceiveDriveData, ref parameter, Marshal.SizeOf(parameter),

        out result, Marshal.SizeOf(typeof(DriveSmartReadThresholdsResult)), 

        out bytesReturned, IntPtr.Zero); 

      return (isValid) ? result.Thresholds : new DriveThresholdValue[0];

    }

    private string GetString(byte[] bytes) {   

      char[] chars = new char[bytes.Length];

      for (int i = 0; i < bytes.Length; i += 2) {

        chars[i] = (char)bytes[i + 1];

        chars[i + 1] = (char)bytes[i];

      }

      return new string(chars).Trim(new char[] { ' ', '\0' });

    }

    public bool ReadNameAndFirmwareRevision(IntPtr handle, int driveNumber, 

      out string name, out string firmwareRevision) 

    {

      DriveCommandParameter parameter = new DriveCommandParameter();

      DriveIdentifyResult result;

      uint bytesReturned;

      parameter.DriveNumber = (byte)driveNumber;

      parameter.Registers.Command = RegisterCommand.IdCmd;

      bool valid = NativeMethods.DeviceIoControl(handle, 

        DriveCommand.ReceiveDriveData, ref parameter, Marshal.SizeOf(parameter), 

        out result, Marshal.SizeOf(typeof(DriveIdentifyResult)), 

        out bytesReturned, IntPtr.Zero);

      if (!valid) {

        name = null;

        firmwareRevision = null;

        return false;

      }

      name = GetString(result.Identify.ModelNumber);

      firmwareRevision = GetString(result.Identify.FirmwareRevision);

      return true;

    }

    public void CloseHandle(IntPtr handle) {

      NativeMethods.CloseHandle(handle);

    }

    public string[] GetLogicalDrives(int driveIndex) {

      List<string> list = new List<string>();

      try {

        using (ManagementObjectSearcher s = new ManagementObjectSearcher(

            "root\\CIMV2",

            "SELECT * FROM Win32_DiskPartition " +

            "WHERE DiskIndex = " + driveIndex))

        using (ManagementObjectCollection dpc = s.Get())

        foreach (ManagementObject dp in dpc) 

          using (ManagementObjectCollection ldc = 

            dp.GetRelated("Win32_LogicalDisk"))

          foreach (ManagementBaseObject ld in ldc) 

            list.Add(((string)ld["Name"]).TrimEnd(':')); 

      } catch { }

      return list.ToArray();

    }

    protected static class NativeMethods {

      private const string KERNEL = "kernel32.dll";

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi,

        CharSet = CharSet.Unicode)]

      public static extern IntPtr CreateFile(string fileName,

        AccessMode desiredAccess, ShareMode shareMode, IntPtr securityAttributes,

        CreationMode creationDisposition, FileAttribute flagsAndAttributes,

        IntPtr templateFilehandle);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]

      public static extern int CloseHandle(IntPtr handle);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]

      [return: MarshalAsAttribute(UnmanagedType.Bool)]

      public static extern bool DeviceIoControl(IntPtr handle,

        DriveCommand command, ref DriveCommandParameter parameter,

        int parameterSize, out DriveSmartReadDataResult result, int resultSize,

        out uint bytesReturned, IntPtr overlapped);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]

      [return: MarshalAsAttribute(UnmanagedType.Bool)]

      public static extern bool DeviceIoControl(IntPtr handle,

        DriveCommand command, ref DriveCommandParameter parameter,

        int parameterSize, out DriveSmartReadThresholdsResult result, 

        int resultSize, out uint bytesReturned, IntPtr overlapped);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]

      [return: MarshalAsAttribute(UnmanagedType.Bool)]

      public static extern bool DeviceIoControl(IntPtr handle,

        DriveCommand command, ref DriveCommandParameter parameter,

        int parameterSize, out DriveCommandResult result, int resultSize,

        out uint bytesReturned, IntPtr overlapped);

      [DllImport(KERNEL, CallingConvention = CallingConvention.Winapi)]

      [return: MarshalAsAttribute(UnmanagedType.Bool)]

      public static extern bool DeviceIoControl(IntPtr handle,

        DriveCommand command, ref DriveCommandParameter parameter,

        int parameterSize, out DriveIdentifyResult result, int resultSize,

        out uint bytesReturned, IntPtr overlapped);

    }    

  }

}