/* 
 * Copyright (c) 2010, 2012, Oracle and/or its affiliates. All rights reserved.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; version 2 of the
 * License.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301  USA
 */

#include "stdafx.h"

#include <iostream>
#include <sstream>

#include "base/string_utilities.h"
#include "wmi.h"
#include "base/threading.h"

using namespace wmi;
using namespace grt;

//--------------------------------------------------------------------------------------------------

/**
 * Converts the given variant to an UTF-8 encoded standard string. The variant content is converted to
 * a string if it is not already one.
 */
std::string variant2string(VARIANT& value) 
{
  if (value.vt == VT_NULL)
    return "NULL";

  VariantChangeType(&value, &value, VARIANT_ALPHABOOL, VT_BSTR);
  CW2A valueString(V_BSTR(&value), CP_UTF8);
  return (LPSTR) valueString;
}

//--------------------------------------------------------------------------------------------------

/**
 * Converts the given string (which must be UTF-8 encoded) to a BSTR, encapsulated by CComBSTR
 * which frees us from taking care to deallocate the result.
 */
CComBSTR string2Bstr(const std::string& value)
{
  CComBSTR result = CA2W(value.c_str(), CP_UTF8);
  return result;
}

//--------------------------------------------------------------------------------------------------

#ifdef _DEBUG

void dumpObject(IWbemClassObject* object)
{
  if (object != NULL)
  {
    BSTR objectText;
    object->GetObjectText(0, &objectText);
    std::wcout << objectText << std::endl;
    SysFreeString(objectText);
  }
  else
    std::wcout << L"Object is NULL" << std::endl;
}

//--------------------------------------------------------------------------------------------------

#endif

std::string wmiResultToString(HRESULT wmiResult)
{
  std::string result;
  IWbemStatusCodeText* status = NULL;

  HRESULT hres = CoCreateInstance(CLSID_WbemStatusCodeText, 0, CLSCTX_INPROC_SERVER,
    IID_IWbemStatusCodeText, (LPVOID*)&status);

  if (SUCCEEDED(hres))
  {
    CComBSTR errorString;
    hres = status->GetErrorCodeText(wmiResult, 0, 0, &errorString);

    if (FAILED(hres))
      result = "Internal error: WMI error description retrieval failed.";

    CW2A converted_text(errorString, CP_UTF8);
    result = converted_text;
    status->Release();
  }
  else
    result = "Internal error: WMI status code text creation failed.";

  return result;
}

//--------------------------------------------------------------------------------------------------

std::string serviceResultToString(unsigned int result)
{
  // Error codes according to http://msdn.microsoft.com/en-us/library/aa393660%28VS.85%29.aspx.
  static std::string code2String[] = {
    "Success",
    "Not Supported",
    "Access Denied",
    "Dependent Services Running",
    "Invalid Service Control",
    "Service Cannot Accept Control",
    "Service Not Active",
    "Service Request Timeout",
    "Unknown Failure",
    "Path Not Found",
    "Service Already Running",
    "Service Database Locked",
    "Service Dependency Deleted",
    "Service Dependency Failure",
    "Service Disabled",
    "Service Logon Failure",
    "Service Marked For Deletion",
    "Service No Thread",
    "Status Circular Dependency",
    "Status Duplicate Name",
    "Status Invalid Name",
    "Status Invalid Parameter",
    "Status Invalid Service Account",
    "Status Service Exists",
    "Service Already Paused"
  };


  if (result < 25)
    return code2String[result];

  return "Unknown error";
}

//----------------- WmiMonitor ---------------------------------------------------------------------

WmiMonitor::WmiMonitor(IWbemServices* services, const std::string& parameter)
  : _services(services), _propertyHandle(0), _namePropertyHandle(0), _findTotal(false)
{
  CComBSTR instancePath;

  HRESULT hr = CoCreateInstance(CLSID_WbemRefresher, NULL, CLSCTX_INPROC_SERVER, IID_IWbemRefresher, 
    (void**) &_refresher);
  if (FAILED(hr))
    throw std::runtime_error(_("WMI - Could not create monitor object.\n\n") + wmiResultToString(hr));

  CComPtr<IWbemConfigureRefresher> config;
  hr = _refresher->QueryInterface(IID_IWbemConfigureRefresher, (void**) &config);
  if (FAILED(hr))
  {
    // _refersher is a smart pointer and is automatically freed.
    throw std::runtime_error(_("WMI - Could not create monitor object.\n\n") + wmiResultToString(hr));
  }

  const std::vector<std::string> args = base::split(parameter, ".");
  
  if (args.size() < 2)
  {
	throw std::runtime_error(_("Wrong monitor format. Got '") + parameter + _("', expected '<WMIClass.Name>'"));
  }

  CComBSTR path = args[0].c_str();
  _propertyName = args[1].c_str();

  if (path.Length() == 0)
    throw std::runtime_error(_("WMI - Invalid monitoring parameter specified."));

  // Add an enumerator to the refresher. This is what actually gets the value to monitor.
  hr = config->AddEnum(_services, path, 0, NULL, &_enumerator, &_enumeratorId);
  if (FAILED(hr))
  {
	throw std::runtime_error(_("WMI - Could not register monitoring enumerator.\n\n") + wmiResultToString(hr));
  }
}

//--------------------------------------------------------------------------------------------------

WmiMonitor::~WmiMonitor()
{
  CComPtr<IWbemConfigureRefresher> config;
  HRESULT hr = _refresher->QueryInterface(IID_IWbemConfigureRefresher, (void**) &config);
  if (SUCCEEDED(hr))
    config->Remove(_enumeratorId, WBEM_FLAG_REFRESH_NO_AUTO_RECONNECT);
  // Ignore query interface errors here, as we cannot do anything about them anymore anyway.
}

//--------------------------------------------------------------------------------------------------

std::string WmiMonitor::readValue()
{
  // TODO: error processing is very rudimentary here.
  //       Though raising exceptions for errors is not a good option here,
  //       especially when considering border crossings like C++ <-> Python, C++ <-> Lua.
  // Refresh the enumerator so we actually get values.
  _refresher->Refresh(0L);

  // Read the results.
  ULONG returnCount = 0;
  IWbemObjectAccess** accessors = NULL; // List of accessors.

  // Determine required space. We have to provide it to the getter.
  HRESULT hr = _enumerator->GetObjects(0L, 0, accessors, &returnCount); 
  if (hr == WBEM_E_BUFFER_TOO_SMALL && returnCount > 0)
  {
    accessors = new IWbemObjectAccess*[returnCount];
    SecureZeroMemory(accessors, returnCount * sizeof(IWbemObjectAccess*));

    hr = _enumerator->GetObjects(0L, returnCount, accessors,  &returnCount);
    if (FAILED(hr))
    {
      delete [] accessors;
      return "0";
    }
  }

  std::string result = "0";

  // First time we access the property we have to get a handle for it, which
  // can be reused for any subsequent query.
  CIMTYPE propertyType;
  if (_propertyHandle == 0)
    hr = accessors[0]->GetPropertyHandle(_propertyName, &propertyType, &_propertyHandle);

  CIMTYPE namePropType;
  if (_findTotal && _namePropertyHandle == 0)
    hr = accessors[0]->GetPropertyHandle(L"Name", &namePropType, &_namePropertyHandle);

  if (_propertyHandle != 0)
  {
    DWORD value = 0;
    if (_namePropertyHandle != 0)
    {
      // For processor time queries a value is returned for every processor and a total entry
      // that comprises all processors.
      for (ULONG i = 0; i < returnCount; i++)
      {
        long byteCount = 0;
        byte buffer[20];
        hr = accessors[i]->ReadPropertyValue(_namePropertyHandle, 20, &byteCount, buffer);

        if (StrCmpW((LPCWSTR) buffer, L"_Total") == 0)
        {
          hr = accessors[i]->ReadDWORD(_propertyHandle, &value);
          break;
        }
      }
    }
    else
    {
      // No need to search for a specific value. Simply return the first (and mostly only) property value
      // we got.
      hr = accessors[0]->ReadDWORD(_propertyHandle, &value);
    }

    std::stringstream ss;
    ss << value;
    result = ss.str();
  }

  for (ULONG i = 0; i < returnCount; i++)
    accessors[i]->Release();
  delete [] accessors;

  return result;
}

//----------------- WmiServices --------------------------------------------------------------------

static GStaticMutex _locatorMutex_ = G_STATIC_MUTEX_INIT;
static IWbemLocator* _locator_ = NULL;
static int _locator_refcount_ = 0;

WmiServices::WmiServices(const std::string& server, const std::string& user, const std::string& password)
{
  allocate_locator();
  if (!_locator_)
    throw std::runtime_error("Could not allocate WMI locator object");

  // Connect to the given target node using the provided credentials.
  HRESULT hr;

  // If node is empty then we are connecting to the local box. In this case don't use the given credentials.
  if (server.size() > 0)
  {
    std::string unc = "\\\\" + server + "\\root\\cimv2";

    CComBSTR nodeString = string2Bstr(unc);
    CComBSTR userString = string2Bstr(user);
    CComBSTR passwordString = string2Bstr(password);
    hr = _locator_->ConnectServer(nodeString, userString, passwordString, NULL,
      WBEM_FLAG_CONNECT_USE_MAX_WAIT, NULL, NULL, &_services);
  }
  else
    hr = _locator_->ConnectServer(L"root\\cimv2", NULL, NULL, NULL, WBEM_FLAG_CONNECT_USE_MAX_WAIT,
      NULL, NULL, &_services);

  if (SUCCEEDED(hr))
  {
    /*
    // Set security levels on a WMI connection
    COAUTHIDENTITY cID;

    cID.User           = (USHORT*) CA2W(user.c_str(), CP_UTF8).m_psz;
    cID.UserLength     = user.size();
    cID.Password       = (USHORT*) CA2W(password.c_str(), CP_UTF8).m_psz;
    cID.PasswordLength = password.size();
    cID.Domain         = NULL;
    cID.DomainLength   = 0;
    cID.Flags          = SEC_WINNT_AUTH_IDENTITY_UNICODE;
*/
    // Set a blanket to our service proxy to establish a security context.
    // Values as recommended (http://msdn.microsoft.com/en-us/library/windows/desktop/aa393620%28v=vs.85%29.aspx).
    hr = CoSetProxyBlanket(
      _services,                   // Indicates the proxy to set
      RPC_C_AUTHN_DEFAULT,         // Authentication service
      RPC_C_AUTHZ_DEFAULT,         // Authorization service
      COLE_DEFAULT_PRINCIPAL,      // Server principal name
      RPC_C_AUTHN_LEVEL_DEFAULT,   // Authentication level
      RPC_C_IMP_LEVEL_IMPERSONATE, // Impersonation level
      COLE_DEFAULT_AUTHINFO,       // client identity
      EOAC_DEFAULT                 // proxy capabilities
    );

    if (FAILED(hr))
    {
      throw std::runtime_error(_("WMI setting security blanket failed.\n\n") + wmiResultToString(hr));
    }
  }
  else
    throw std::runtime_error(_("Could not connect to target machine.\n\n") + wmiResultToString(hr));
}

//--------------------------------------------------------------------------------------------------

WmiServices::~WmiServices()
{
  deallocate_locator();
}

//--------------------------------------------------------------------------------------------------

void WmiServices::allocate_locator()
{
  base::GStaticMutexLock lock(_locatorMutex_);
  if (_locator_refcount_ == 0)
  {
    HRESULT hr = CoCreateInstance(CLSID_WbemLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator,
      reinterpret_cast<void**>(&_locator_));

    if (FAILED(hr))
      throw std::runtime_error("Internal error: Instantiation of IWbemLocator failed.\n\n" + wmiResultToString(hr));
  }

  _locator_refcount_++;
}

//--------------------------------------------------------------------------------------------------

void WmiServices::deallocate_locator()
{
  base::GStaticMutexLock lock(_locatorMutex_);
  if (_locator_refcount_ > 0)
  {
    _locator_refcount_--;
    if (_locator_refcount_ == 0)
    {
      _locator_->Release();
      _locator_ = NULL;
    }
  }
}

//--------------------------------------------------------------------------------------------------

/**
 * Executes a wmi query and returns the queried objects as grt dicts.
 *
 * @param query The query to send. It must be of type WQL.
 * @param node The target machine where to execute the query on. Leave blank for localhost.
 * @param user The user name for authentication on a remote box. Ignored for localhost (the current 
 *             user is used in this case).
 * @param password The password for the given user. Also ignore for localhost.
 * @return A list of dictionaries containing an entry for each object returned by the query, with
 *         name/value pairs of object properties.
 */
grt::DictListRef WmiServices::query(grt::GRT* grt, const std::string& query)
{
  // Making this function explicitly thread-safe might be unnecessary as we don't have
  // any data which is allocated/deallocated concurrently. But since we know we will be called
  // from different threads we play safe here, as it does not harm either.
  base::GStaticMutexLock lock(_locatorMutex_);

  grt::DictListRef queryResult(grt);

  // Execute the given query.
  CComPtr<IEnumWbemClassObject> enumerator;
  CComBSTR converted_query = string2Bstr(query);
  HRESULT hr = _services->ExecQuery(L"WQL", converted_query, WBEM_FLAG_FORWARD_ONLY, NULL, &enumerator);

  if (SUCCEEDED(hr))
  {
    // We also need to set a security context for the enumerator or the next code will fail
    // with "access denied" on remote boxes.
    hr = CoSetProxyBlanket(enumerator, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL,
      RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);

    // Read the returned results and create our own response.
    ULONG count = 0;
    while (true)
    {
      CComPtr<IWbemClassObject> wbemObject = NULL;
      hr = enumerator->Next(WBEM_INFINITE, 1L, &wbemObject, &count);
      if (FAILED(hr))
        throw std::runtime_error(_("WMI enumeration failed.\n\n") + wmiResultToString(hr));

      if (count == 0)
        break;

      SAFEARRAY* names;
      wbemObject->GetNames(NULL, WBEM_FLAG_NONSYSTEM_ONLY, NULL, &names);

      long lowBound, highBound;
      SafeArrayGetLBound(names, 1, &lowBound);
      SafeArrayGetUBound(names, 1, &highBound);

      DictRef resultNames(grt);
      for (long i = lowBound; i <= highBound; i++)
      {
        CComBSTR name;
        SafeArrayGetElement(names, &i, &name);
        CW2A nameString(name, CP_UTF8);

        // Read the actual value for the given property.
        VARIANT value;
        hr = wbemObject->Get(name, 0, &value, NULL, NULL);

        if (SUCCEEDED(hr))
        {
          // In order to ease the value transport everything is converted to a string.
          resultNames.gset((LPSTR)nameString, variant2string(value));

          VariantClear(&value);
        }
      }
      SafeArrayDestroy(names);

      queryResult.insert(resultNames);
    }
  }
  else
  {
    throw std::runtime_error("WMI query execution failed with this error:" + wmiResultToString(hr));
  }

  return queryResult;
}

//--------------------------------------------------------------------------------------------------

/**
 * Allows to work with a Windows service locally or remotely. The result depends on the action.
 *
 * @param service The display name of the service (as stored in the server instance) to act on.
 * @param action The action to execute. Supported are:
 *   - status: return the status of the service
 *   - start: start the service if it is not running.
 *   - stop: stop a running service.
 * @return A string describing the outcome of the action.
 *   - for status: either running, stopping, starting, stopped, unknown
 *   - for start: either completed, already-running, already-starting, stopping, error
 *   - for stop: either completed, already-stopped, already-stopping, starting, error
 */
std::string WmiServices::serviceControl(const std::string& service, const std::string& action)
{
  base::GStaticMutexLock lock(_locatorMutex_);

  CComBSTR instancePath = string2Bstr("Win32_Service.Name='" + service + "'");

  CComPtr<IWbemClassObject> serviceInstance;
  HRESULT hr = _services->GetObject(instancePath, 0, NULL, &serviceInstance, NULL);
  if (FAILED(hr))
    return "error: " + wmiResultToString(hr);

  VARIANT value;
  hr = serviceInstance->Get(L"State", 0, &value, NULL, NULL);
  if (FAILED(hr))
    return "unknown";

  std::string state = base::tolower(variant2string(value));
  VariantClear(&value);

  if (action == "status")
    return state;
  else
  {
    CComBSTR methodName;
    std::string waitState;
    std::string expectedState;
    if (action == "start")
    {
      if (state == "running")
        return "already-running";
      if (state == "stop pending")
        return "stopping";
      if (state == "start pending")
        return "already-starting";
      if (state != "stopped")
        return "error";

      methodName = "StartService";
      waitState = "start pending";
      expectedState = "running";
    }
    else
      if (action == "stop")
      {
        if (state == "stopped")
          return "already-stopped";
        if (state == "stop pending")
          return "stopping";
        if (state == "start pending")
          return "starting";
        if (state != "running")
          return "error";

        methodName = "StopService";
        waitState = "stop pending";
        expectedState = "stopped";
      }

    if (methodName.Length() > 0)
    {
      CComBSTR serviceClassPath = "Win32_Service";
      CComPtr<IWbemClassObject> serviceClass;
      CComPtr<IWbemClassObject> outInstance;
      CComPtr<IWbemClassObject> inClass;
      CComPtr<IWbemClassObject> inInstance;
      CComPtr<IWbemClassObject> outClass;

      if (SUCCEEDED(hr))
        hr = _services->GetObject(serviceClassPath, 0, NULL, &serviceClass, NULL);
      if (SUCCEEDED(hr))
        hr = serviceClass->GetMethod(methodName, 0, &inClass, &outClass);
      if (SUCCEEDED(hr) && inClass)
        hr = inClass->SpawnInstance(0, &inInstance);

      if (SUCCEEDED(hr))
        hr = _services->ExecMethod(instancePath, methodName, 0, NULL, inInstance, &outInstance, NULL);

      if (FAILED(hr))
        return "error: " + wmiResultToString(hr);

      // Try to figure out what happened if the call as such succeeded but returned an error code.
      // Service methods return uint32 error codes.
      VARIANT returnValue;
      hr = outInstance->Get(L"ReturnValue", 0, &returnValue, NULL, NULL);

      if (FAILED(hr))
        return "error: " + wmiResultToString(hr);

      if (returnValue.vt != VT_EMPTY)
      {
        unsigned int result = V_UI4(&returnValue);
        VariantClear(&returnValue);
        if (result != 0)
          return "error: " + serviceResultToString(result);
      }

      // So far everything worked fine. Now check the service if it changed as we intended.
      // Wait in a loop with increasing time steps for the result. If there was not the expected
      // change after at most 34 secs then assume an error and return.
      int i = 10;
      int sleepTime = 250; // Start with 1/4 sec steps, to return as quick as possible if things are fine.
      do
      {
        Sleep(sleepTime);

        HRESULT hr = _services->GetObject(instancePath, 0, NULL, &serviceInstance, NULL);
        if (FAILED(hr))
          return "error: " + wmiResultToString(hr);

        VARIANT value;
        hr = serviceInstance->Get(L"State", 0, &value, NULL, NULL);
        if (FAILED(hr))
          return "error: " + wmiResultToString(hr);

        state = base::tolower(variant2string(value));
        VariantClear(&value);

        if (state != waitState)
          break;

        if (--i == 0)
        {
          i = 10;
          sleepTime *= 2; // Double the time we wait for the next check of the service.
          if (sleepTime >= 4000)
            break;
        }
      } while (true);

      if (state != expectedState)
        return "error";

      return "completed";
    }
  }

  return "unknown";
}

//--------------------------------------------------------------------------------------------------

/**
 * Retrieves operating system values like total available memory, OS description, boot device etc.
 *
 * @param what Specifies which value to query and return. Supported are all properties which belong
 * to the Win32_OperatingSystem class (see http://msdn.microsoft.com/en-us/library/aa394239%28VS.85%29.aspx).
 */
std::string WmiServices::systemStat(const std::string& what)
{
  base::GStaticMutexLock lock(_locatorMutex_);

  CComBSTR instancePath = string2Bstr("Win32_OperatingSystem");

  CComPtr<IEnumWbemClassObject> enumerator;
  HRESULT hr = _services->ExecQuery(L"WQL", L"select * from Win32_OperatingSystem", WBEM_FLAG_FORWARD_ONLY, NULL, &enumerator);

  std::string result = "-1";
  if (SUCCEEDED(hr))
  {
    hr = CoSetProxyBlanket(enumerator, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL,
      RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE);

    CComPtr<IWbemClassObject> wbemObject = NULL;
    ULONG count = 0;
    hr = enumerator->Next(WBEM_INFINITE, 1L, &wbemObject, &count);
    if (FAILED(hr) || count == 0)
      return "-1";

    CComBSTR propertyName = string2Bstr(what);
    VARIANT value;
    hr = wbemObject->Get(propertyName, 0, &value, NULL, NULL);

    if (SUCCEEDED(hr))
    {
      result = base::tolower(variant2string(value));
      VariantClear(&value);
    };
  };

  return result;
}

//--------------------------------------------------------------------------------------------------

WmiMonitor* WmiServices::startMonitoring(const std::string& parameter)
{
  return new WmiMonitor(_services, parameter);
}

//--------------------------------------------------------------------------------------------------

void wmi::WmiServices::stopMonitoring(WmiMonitor* monitor)
{
  delete monitor;
}

//--------------------------------------------------------------------------------------------------
