/*

  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) 2010-2012 Michael Möller <mmoeller@openhardwaremonitor.org>

*/

using System;

using System.IO;

using System.Reflection;

using System.Runtime.InteropServices;

using System.Security.AccessControl;

using System.Threading;

using System.Text;

namespace OpenHardwareMonitor.Hardware {

  internal static class Ring0 {

    private static KernelDriver driver;

    private static string fileName;

    private static Mutex isaBusMutex;

    private static readonly StringBuilder report = new StringBuilder();

    private const uint OLS_TYPE = 40000;

    private static IOControlCode

      IOCTL_OLS_GET_REFCOUNT = new IOControlCode(OLS_TYPE, 0x801,

        IOControlCode.Access.Any),

      IOCTL_OLS_GET_DRIVER_VERSION = new IOControlCode(OLS_TYPE, 0x800,

        IOControlCode.Access.Any),

      IOCTL_OLS_READ_MSR = new IOControlCode(OLS_TYPE, 0x821,

        IOControlCode.Access.Any),

      IOCTL_OLS_WRITE_MSR = new IOControlCode(OLS_TYPE, 0x822, 

        IOControlCode.Access.Any),

      IOCTL_OLS_READ_IO_PORT_BYTE = new IOControlCode(OLS_TYPE, 0x833,

        IOControlCode.Access.Read),

      IOCTL_OLS_WRITE_IO_PORT_BYTE = new IOControlCode(OLS_TYPE, 0x836, 

        IOControlCode.Access.Write),

      IOCTL_OLS_READ_PCI_CONFIG = new IOControlCode(OLS_TYPE, 0x851, 

        IOControlCode.Access.Read),

      IOCTL_OLS_WRITE_PCI_CONFIG = new IOControlCode(OLS_TYPE, 0x852,

        IOControlCode.Access.Write),

      IOCTL_OLS_READ_MEMORY = new IOControlCode(OLS_TYPE, 0x841,

        IOControlCode.Access.Read);

    private static string GetTempFileName() {

      // try to create one in the application folder

      string fileName = Path.ChangeExtension(

        Assembly.GetEntryAssembly().Location, ".sys");

      try {

        using (FileStream stream = File.Create(fileName)) {

          return fileName;

        }        

      } catch (IOException) { } 

        catch (UnauthorizedAccessException) { }

      // if this failed, try to get a file in the temporary folder

      try {

        return Path.GetTempFileName();        

      } catch (IOException) { 

          // some I/O exception

      } 

      catch (UnauthorizedAccessException) { 

        // we do not have the right to create a file in the temp folder

      }

      catch (NotSupportedException) {

        // invalid path format of the TMP system environment variable

      }

      return null;

    }

    private static bool ExtractDriver(string fileName) {

      string resourceName = "OpenHardwareMonitor.Hardware." +

        (OperatingSystem.Is64BitOperatingSystem() ? "WinRing0x64.sys" : 

        "WinRing0.sys");

      string[] names =

        Assembly.GetExecutingAssembly().GetManifestResourceNames();

      byte[] buffer = null;

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

        if (names[i].Replace('\\', '.') == resourceName) {

          using (Stream stream = Assembly.GetExecutingAssembly().

            GetManifestResourceStream(names[i])) 

          {

              buffer = new byte[stream.Length];

              stream.Read(buffer, 0, buffer.Length);

          }

        }

      }

      if (buffer == null)

        return false;

      try {

        using (FileStream target = new FileStream(fileName, FileMode.Create)) {

          target.Write(buffer, 0, buffer.Length);

          target.Flush();

        }

      } catch (IOException) { 

        // for example there is not enough space on the disk

        return false; 

      }

      // make sure the file is actually writen to the file system

      for (int i = 0; i < 20; i++) {

        try {

          if (File.Exists(fileName) &&

            new FileInfo(fileName).Length == buffer.Length) 

          {

            return true;

          }

          Thread.Sleep(100);

        } catch (IOException) {

          Thread.Sleep(10);

        }

      }

      // file still has not the right size, something is wrong

      return false;

    }

    public static void Open() {

      // no implementation for unix systems

      int p = (int)Environment.OSVersion.Platform;

      if ((p == 4) || (p == 128))

        return;  

      if (driver != null)

        return;

      // clear the current report

      report.Length = 0;

      driver = new KernelDriver("WinRing0_1_2_0");

      driver.Open();

      if (!driver.IsOpen) {

        // driver is not loaded, try to install and open

        fileName = GetTempFileName();

        if (fileName != null && ExtractDriver(fileName)) {

          string installError;

          if (driver.Install(fileName, out installError)) {

            driver.Open();

            if (!driver.IsOpen) {

              driver.Delete();

              report.AppendLine("Status: Opening driver failed after install");

            }

          } else {

            string errorFirstInstall = installError;

            // install failed, try to delete and reinstall

            driver.Delete();

            // wait a short moment to give the OS a chance to remove the driver

            Thread.Sleep(2000);

            string errorSecondInstall;

            if (driver.Install(fileName, out errorSecondInstall)) {

              driver.Open();

              if (!driver.IsOpen) {

                driver.Delete();

                report.AppendLine(

                  "Status: Opening driver failed after reinstall");

              }

            } else {

              report.AppendLine("Status: Installing driver \"" +

                fileName + "\" failed" +

                (File.Exists(fileName) ? " and file exists" : ""));

              report.AppendLine("First Exception: " + errorFirstInstall);

              report.AppendLine("Second Exception: " + errorSecondInstall);

            }

          }

        } else {

          report.AppendLine("Status: Extracting driver failed");

        }

        try {

          // try to delte the driver file

          if (File.Exists(fileName))

            File.Delete(fileName);

          fileName = null;

        } catch (IOException) { } 

          catch (UnauthorizedAccessException) { }

      }

      if (!driver.IsOpen) 

        driver = null;

      string mutexName = "Global\\Access_ISABUS.HTP.Method";

      try {

        isaBusMutex = new Mutex(false, mutexName);

      } catch (UnauthorizedAccessException) {

        try {

          isaBusMutex = Mutex.OpenExisting(mutexName, MutexRights.Synchronize);

        } catch { }

      }

    }

    public static bool IsOpen {

      get { return driver != null; }

    }

    public static void Close() {

      if (driver == null)

        return;

      uint refCount = 0;

      driver.DeviceIOControl(IOCTL_OLS_GET_REFCOUNT, null, ref refCount);

      driver.Close();

      if (refCount <= 1)

        driver.Delete();

      driver = null;

      if (isaBusMutex != null) {

        isaBusMutex.Close();

        isaBusMutex = null;

      }

      // try to delete temporary driver file again if failed during open

      if (fileName != null && File.Exists(fileName)) {

        try {

          File.Delete(fileName);

          fileName = null;

        } catch (IOException) { } 

          catch (UnauthorizedAccessException) { }

      }

    }

    public static string GetReport() {

      if (report.Length > 0) {

        StringBuilder r = new StringBuilder();

        r.AppendLine("Ring0");

        r.AppendLine();

        r.Append(report);

        r.AppendLine();

        return r.ToString();

      } else

        return null;

    }

    public static bool WaitIsaBusMutex(int millisecondsTimeout) {

      if (isaBusMutex == null)

        return true;

      try {

        return isaBusMutex.WaitOne(millisecondsTimeout, false);

      } catch (AbandonedMutexException) { return false; } 

        catch (InvalidOperationException) { return false; }

    }

    public static void ReleaseIsaBusMutex() {

      if (isaBusMutex == null)

        return;

      isaBusMutex.ReleaseMutex();

    }

    public static bool Rdmsr(uint index, out uint eax, out uint edx) {

      if (driver == null) {

        eax = 0;

        edx = 0;

        return false;

      }

      ulong buffer = 0;

      bool result = driver.DeviceIOControl(IOCTL_OLS_READ_MSR, index,

        ref buffer);

      edx = (uint)((buffer >> 32) & 0xFFFFFFFF);

      eax = (uint)(buffer & 0xFFFFFFFF);

      return result;

    }

    public static bool RdmsrTx(uint index, out uint eax, out uint edx,

      ulong threadAffinityMask) 

    {

      ulong mask = ThreadAffinity.Set(threadAffinityMask);

      bool result = Rdmsr(index, out eax, out edx);

      ThreadAffinity.Set(mask);

      return result;

    }

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

    private struct WrmsrInput {

      public uint Register;

      public ulong Value;

    }

    public static bool Wrmsr(uint index, uint eax, uint edx) {

      if (driver == null)

        return false;

      WrmsrInput input = new WrmsrInput();

      input.Register = index;

      input.Value = ((ulong)edx << 32) | eax;

      return driver.DeviceIOControl(IOCTL_OLS_WRITE_MSR, input);

    }

    public static byte ReadIoPort(uint port) {

      if (driver == null)

        return 0;

      uint value = 0;

      driver.DeviceIOControl(IOCTL_OLS_READ_IO_PORT_BYTE, port, ref value);

      return (byte)(value & 0xFF);

    }

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

    private struct WriteIoPortInput {

      public uint PortNumber;

      public byte Value;

    }

    public static void WriteIoPort(uint port, byte value) {

      if (driver == null)

        return;

      WriteIoPortInput input = new WriteIoPortInput();

      input.PortNumber = port;

      input.Value = value;

      driver.DeviceIOControl(IOCTL_OLS_WRITE_IO_PORT_BYTE, input);

    }

    public const uint InvalidPciAddress = 0xFFFFFFFF;

    public static uint GetPciAddress(byte bus, byte device, byte function) {

      return

        (uint)(((bus & 0xFF) << 8) | ((device & 0x1F) << 3) | (function & 7));

    }

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

    private struct ReadPciConfigInput {

      public uint PciAddress;

      public uint RegAddress;

    }

    public static bool ReadPciConfig(uint pciAddress, uint regAddress, 

      out uint value) 

    {

      if (driver == null || (regAddress & 3) != 0) {

        value = 0;

        return false;

      }

      ReadPciConfigInput input = new ReadPciConfigInput();

      input.PciAddress = pciAddress;

      input.RegAddress = regAddress;

      value = 0;

      return driver.DeviceIOControl(IOCTL_OLS_READ_PCI_CONFIG, input, 

        ref value);

    }

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

    private struct WritePciConfigInput {

      public uint PciAddress;

      public uint RegAddress;

      public uint Value;

    }

    public static bool WritePciConfig(uint pciAddress, uint regAddress, 

      uint value) 

    {

      if (driver == null || (regAddress & 3) != 0)

        return false;

      WritePciConfigInput input = new WritePciConfigInput();

      input.PciAddress = pciAddress;

      input.RegAddress = regAddress;

      input.Value = value;

      return driver.DeviceIOControl(IOCTL_OLS_WRITE_PCI_CONFIG, input);

    }

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

    private struct ReadMemoryInput {

      public ulong address;

      public uint unitSize;

      public uint count;

    }

    public static bool ReadMemory<T>(ulong address, ref T buffer) {

      if (driver == null) {

        return false;

      }

      ReadMemoryInput input = new ReadMemoryInput();

      input.address = address;

      input.unitSize = 1;

      input.count = (uint)Marshal.SizeOf(buffer);

      return driver.DeviceIOControl(IOCTL_OLS_READ_MEMORY, input,

        ref buffer);

    }

  }

}