//{ **************************************************************************** }
//{ FileName............: WMI.CPP                                                }
//{ Project.............: FLEXCOP                                                }
//{ Author(s)...........: M.Majoor                                               }
//{ Original contents...: Walter Oney ('Programming the Windows Driver Model')   }
//{ Version.............: 1.00                                                   }
//{ ---------------------------------------------------------------------------- }
//{ Windows Management Instrumentation handlers for FLEXCOP driver.              }
//{                                                                              }
//{ Version  Date      Comment                                                   }
//{ 1.00     20050410  - Initial release based on SAA7146A driver V2.02          }
//{                                                                              }
//{ Notes:   NOT IMPLEMENTED (just the basics)                                   }
//{ **************************************************************************** }


#include "stddcls.h"
#include <wmilib.h>
#include "driver.h"

#include <initguid.h>
#include <wdmguid.h>
#include "guids.h"


//{ **************************************************************************** }
//{ Descript : Forwards.                                                         }
//{ **************************************************************************** }
NTSTATUS QueryRegInfo   (PDEVICE_OBJECT fdo, PULONG flags, PUNICODE_STRING instname, PUNICODE_STRING* regpath, PUNICODE_STRING resname, PDEVICE_OBJECT* pdo);
NTSTATUS QueryDataBlock (PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG instcount, PULONG instlength, ULONG bufsize, PUCHAR buffer);
NTSTATUS SetDataBlock   (PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG bufsize, PUCHAR buffer);
NTSTATUS SetDataItem    (PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG bufsize, PUCHAR buffer);
NTSTATUS ExecuteMethod  (PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG cbInbuf, ULONG cbOutbuf, PUCHAR buffer);
NTSTATUS FunctionControl(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, WMIENABLEDISABLECONTROL fcn, BOOLEAN enable);


// TODO Add entries to the GUID list for any standard WMI classes
// that your driver supports. The wizard generates entries just
// for classes in your custom schema. If you also support standard
// classes, add them *after* the custom ones because the wizard will
// have generated case labels based on assumption that the first entry
// in the list is for the first custom class.


// The WMILIB_CONTEXT structure declared at file scope describes the class GUIDs
// your driver supports and lists several callback functions that WMILIB uses to
// handle WMI requests in the appropriate device-dependent and driver-dependent way.
WMIGUIDREGINFO guidlist[] =
{
  {&JUST_SOME_WMI, 1, WMIREG_FLAG_INSTANCE_PDO},
};


WMILIB_CONTEXT libinfo =
{
  arraysize(guidlist),                                     // GuidCount
  guidlist,                                                // GuidList
  QueryRegInfo,
  QueryDataBlock,
  SetDataBlock,
  SetDataItem,
  ExecuteMethod,
  FunctionControl,
};


//{ **************************************************************************** }
//{ Params   : <fdo>     Functional device object                                }
//{            <Irp>     I/O request packet                                      }
//{ Return   : <Result>  NTSTATUS                                                }
//{                                                                              }
//{ Descript : Dispatch of WMI.                                                  }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS DispatchWmi(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  // As with other dispatch routines, we acquire and release the remove lock while
  // handling this IRP. The problem we prevent is having the device object underneath
  // us disappear because of a Plug and Play (PnP) event. Our own device object cannot
  // disappear because our call to IoWMIRegistrationControl acquired a reference to it.
  NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
  if (!NT_SUCCESS(status))
    return CompleteRequest(Irp, status, 0);

  #if DBG
    PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
    ULONG fcn = stack->MinorFunction;

    static char* fcnname[] =
    {
      "IRP_MN_QUERY_ALL_DATA",
      "IRP_MN_QUERY_SINGLE_INSTANCE",
      "IRP_MN_CHANGE_SINGLE_INSTANCE",
      "IRP_MN_CHANGE_SINGLE_ITEM",
      "IRP_MN_ENABLE_EVENTS",
      "IRP_MN_DISABLE_EVENTS",
      "IRP_MN_ENABLE_COLLECTION",
      "IRP_MN_DISABLE_COLLECTION",
      "IRP_MN_REGINFO",
      "IRP_MN_EXECUTE_METHOD",
    };

    KdPrint((DRIVERNAME " - WMI Request (%s)\n", fcnname[fcn]));
  #endif

  // Delegate processing to the WMILIB helper library
  SYSCTL_IRP_DISPOSITION disposition;
  // This statement calls WMILIB to handle the IRP. We pass the address of our
  // WMILIB_CONTEXT structure. It's customary to use a static context structure,
  // by the way, because the information in it is unlikely to change from one IRP
  // to the next. WmiSystemControl returns two pieces of information: an NTSTATUS
  // code and a SYSCTL_IRP_DISPOSITION value.
  status = WmiSystemControl(&libinfo, fdo, Irp, &disposition);

  // Depending on the disposition code, we might have additional work to perform
  // on this IRP. If the code is IrpProcessed, the IRP has already been completed
  // and we need do nothing more with it. This case would be the normal one for minor
  // functions other than IRP_MN_REGINFO.
  switch (disposition)
  {                                                        // Finish handling IRP
    case IrpProcessed:
      break;
    case IrpNotCompleted:
      // If the disposition code is IrpNotCompleted, completing the IRP is our responsibility.
      // This case would be the normal one for IRP_MN_REGINFO. WMILIB has already filled in
      // the IoStatus block of the IRP, so we need only call IoCompleteRequest.
      IoCompleteRequest(Irp, IO_NO_INCREMENT);             // WMILIB already filled in IoStatus fields
      break;
    default:
      // The default and IrpNotWmi cases shouldn't arise in Windows 2000. We'd get
      // to the default label if we weren't handling all possible disposition codes;
      // we'd get to the IrpNotWmi case label if we sent an IRP to WMILIB that didn't
      // have one of the minor function codes that specifies WMI functionality.

      // The IrpForward case occurs for system control IRPs that are intended for
      // some other driver. Recall that the ProviderId parameter indicates the driver
      // that is supposed to handle this IRP. WmiSystemControl compares that value
      // to the device object pointer we supply as the second function argument. If
      // they're not the same, it returns IrpForward so that we'll send the IRP down
      // the stack to the next driver.

      IoSkipCurrentIrpStackLocation(Irp);
      status = IoCallDriver(pdx->LowerDeviceObject, Irp);
      break;
  }

  IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
  return status;
}


//{ **************************************************************************** }
//{ Params   : <fdo>  Functional device object                                   }
//{ Return   : -                                                                 }
//{                                                                              }
//{ Descript : Initialize WMI.                                                   }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
void WmiInitialize(PDEVICE_OBJECT fdo)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

  // TODO Do any required initialization of data blocks
  NTSTATUS status = IoWMIRegistrationControl(fdo, WMIREG_ACTION_REGISTER);
  if (!NT_SUCCESS(status))
    KdPrint((DRIVERNAME " - IoWMIRegistrationControl failed - %X\n", status));
}


//{ **************************************************************************** }
//{ Params   : <fdo>  Functional device object                                   }
//{ Return   : -                                                                 }
//{                                                                              }
//{ Descript : Terminate WMI.                                                    }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
void WmiTerminate(PDEVICE_OBJECT fdo)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  IoWMIRegistrationControl(fdo, WMIREG_ACTION_DEREGISTER);
}


//{ **************************************************************************** }
//{ Params   : <fdo>       Functional device object                              }
//{            <flags>     ..                                                    }
//{            <instname>  Instance name                                         }
//{            <regpath>   Path to registry key                                  }
//{            <resname>   Resource name                                         }
//{            <pdo>       Device object                                         }
//{ Return   : <Result>    NTSTATUS                                              }
//{                                                                              }
//{ Descript : Query registry info.                                              }
//{            The first system control IRP we'll receive after making our       }
//{            registration call has the minor function code IRP_MN_REGINFO. When}
//{            we pass this IRP to WmiSystemControl, it turns around and calls   }
//{            the QueryRegInfo function—it finds the function's address in our  }
//{            WMILIB_CONTEXT structure.                                         }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS QueryRegInfo(PDEVICE_OBJECT fdo, PULONG flags, PUNICODE_STRING instname,
                      PUNICODE_STRING* regpath, PUNICODE_STRING resname, PDEVICE_OBJECT* pdo)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  *flags = WMIREG_FLAG_INSTANCE_PDO;
  // We set regpath to the address of a UNICODE_STRING structure that contains the
  // name of the service registry key describing our driver. This key is the one
  // below \System\CurrentControlSet\Services. Our DriverEntry routine received
  // the name of this key as an argument and saved it in the global variable servkey.
  // We set resname to the name we chose to give our schema in our resource script.
  *regpath = &servkey;
  RtlInitUnicodeString(resname, L"MofResource");
  *pdo = pdx->Pdo;
  return STATUS_SUCCESS;
}


//{ **************************************************************************** }
//{ Params   : <fdo>         Functional device object                            }
//{            <Irp>         I/O request packet                                  }
//{            <guidindex>   ID                                                  }
//{            <instindex>   Instance index                                      }
//{            <instcount>   Instance count                                      }
//{            <instlength>  Instance length                                     }
//{            <bufavail>    Available buffer length                             }
//{            <buffer>      Buffer                                              }
//{ Return   : <Result>      NTSTATUS                                            }
//{                                                                              }
//{ Descript : Query data block.                                                 }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS QueryDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex,
                        ULONG instindex, ULONG instcount, PULONG instlength, ULONG bufavail, PUCHAR buffer)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  ULONG bufsize = 0;
  NTSTATUS status = STATUS_SUCCESS;

  // TODO Complete each of the cases in this switch according to the following
  // model

  // bufsize = # bytes needed for data alone
  // if (!instlength || bufavail < bufsize)
  //    status = STATUS_BUFFER_TOO_SMALL;
  // else
  //   <fill in data values starting at "buffer" and instlength[] array

  // TODO Also add case labels for any standard WMI classes you support
  switch (guidindex)
  {                                                        // Provide indicated data
    case 0:
    {                                                      // JUST_SOME_WMI
      // TODO provide data for this guid as explained above
      break;
    }
    default:
      ASSERT(FALSE);                                       // WMILIB should never give us an unexpected GUID
      status = STATUS_WMI_GUID_NOT_FOUND;
  }

  return WmiCompleteRequest(fdo, Irp, status, bufsize, IO_NO_INCREMENT);
}


//{ **************************************************************************** }
//{ Params   : <fdo>        Functional device object                             }
//{            <Irp>        I/O request packet                                   }
//{            <guidindex>  ID                                                   }
//{            <instindex>  Instance index                                       }
//{            <bufavail>   Available buffer length                              }
//{            <buffer>     Buffer                                               }
//{ Return   : <Result>     NTSTATUS                                             }
//{                                                                              }
//{ Descript : Set data block.                                                   }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS SetDataBlock(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG bufavail, PUCHAR buffer)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  NTSTATUS status = STATUS_SUCCESS;
  ULONG bufsize = 0;

  // TODO Complete each of the cases in this switch according to the following
  // model

  // bufsize = # bytes needed for data alone
  // if (bufavail == bufsize)
  //    <set data block for "instindex" instance from data at "buffer">
  // else
  //    status = STATUS_INFO_LENGTH_MISMATCH, bufsize = 0;

  // TODO Also add case labels for any standard WMI classes you support
  switch (guidindex)
  {                                                        // Provide indicated data
    case 0:
    {                                                      // JUST_SOME_WMI
      // TODO set data for this guid and instance as explained above
      break;
    }
    default:
      ASSERT(FALSE);                                       // WMILIB should never give us an unexpected GUID
      status = STATUS_WMI_GUID_NOT_FOUND;
  }

  return WmiCompleteRequest(fdo, Irp, status, bufsize, IO_NO_INCREMENT);
}


//{ **************************************************************************** }
//{ Params   : <fdo>        Functional device object                             }
//{            <Irp>        I/O request packet                                   }
//{            <guidindex>  ID                                                   }
//{            <instindex>  Instance index                                       }
//{            <id>         ID                                                   }
//{            <bufavail>   Available buffer length                              }
//{            <buffer>     Buffer                                               }
//{ Return   : <Result>     NTSTATUS                                             }
//{                                                                              }
//{ Descript : Set data item.                                                    }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS SetDataItem(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG bufavail, PUCHAR buffer)
{
  PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
  NTSTATUS status;
  ULONG bufsize = 0;

  // TODO Complete each of the cases in this switch according to the following
  // model

  // bufsize = # bytes needed for data item alone
  // if (bufavail == bufsize)
  //    <set item # "id" from data at "buffer">
  // else
  //    status = STATUS_INFO_LENGTH_MISMATCH, bufsize = 0;

  // TODO Also add case labels for any standard WMI classes you support
  switch (guidindex)
  {                                                        // Provide indicated data
    case 0:
    {                                                      // JUST_SOME_WMI
      // TODO set data for this guid, instance, and item as explained above
      break;
    }
    default:
      ASSERT(FALSE);                                       // WMILIB should never give us an unexpected GUID
      status = STATUS_WMI_GUID_NOT_FOUND;
  }

  return WmiCompleteRequest(fdo, Irp, status, bufsize, IO_NO_INCREMENT);
}


//{ **************************************************************************** }
//{ Params   : <fdo>        Functional device object                             }
//{            <Irp>        I/O request packet                                   }
//{            <guidindex>  ID                                                   }
//{            <instindex>  Instance index                                       }
//{            <id>         ID                                                   }
//{            <cbInbuf>    Input buffer                                         }
//{            <cbOutbug>   Output buffer                                        }
//{            <buffer>     Buffer                                               }
//{ Return   : <Result>     NTSTATUS                                             }
//{                                                                              }
//{ Descript : Execute method.                                                   }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS ExecuteMethod(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, ULONG instindex, ULONG id, ULONG cbInbuf, ULONG cbOutbuf, PUCHAR buffer)
{
  // TODO Write this function if any of the classes you support have method
  // routines
  return WmiCompleteRequest(fdo, Irp, STATUS_INVALID_DEVICE_REQUEST, 0, IO_NO_INCREMENT);
}


//{ **************************************************************************** }
//{ Params   : <fdo>        Functional device object                             }
//{            <Irp>        I/O request packet                                   }
//{            <guidindex>  ID                                                   }
//{            <fcn>        Control                                              }
//{            <enable>     Enable                                               }
//{ Return   : <Result>     NTSTATUS                                             }
//{                                                                              }
//{ Descript : Function control.                                                 }
//{ Notes    :                                                                   }
//{ **************************************************************************** }
NTSTATUS FunctionControl(PDEVICE_OBJECT fdo, PIRP Irp, ULONG guidindex, WMIENABLEDISABLECONTROL fcn, BOOLEAN enable)
{
  // TODO Write this function if any of the classes you support are "expensive"
  // or involve events.
  return WmiCompleteRequest(fdo, Irp, STATUS_SUCCESS, 0, IO_NO_INCREMENT);
}
