generate_descriptor.py

#!/usr/bin/python2.6
# Copyright 2010 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A script to generate ETW event descriptors for ETW MOF events."""
import optparse
import os
import pywintypes
import win32com.client

LICENSE_HEADER = """\
#!/usr/bin/python2.6
# Copyright 2010 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License."""

DO_NOT_EDIT_HEADER = """\
\"\"\"Generated event descriptor file for a MOF event class.

DO NOT EDIT. This is an ETW event descriptor file generated by
sawbuck/py/etw/generate_descriptor.py. It contains event descriptions for
MOF GUID %s.
\"\"\"
"""

# Generate symbols for the WbemScripting module so that we can have symbols
# for debugging and use constants throughout the file.
win32com.client.gencache.EnsureModule('{565783C6-CB41-11D1-8B02-00600806D9B6}',
                                      0, 1, 1)


class DescriptorGenerator(object):
  """A class used to generate an ETW event descriptor."""

  def __init__(self):
    self._locator = win32com.client.Dispatch('WbemScripting.SWbemLocator')
    self._services = self._locator.ConnectServer('localhost', r'root\wmi')

  def _GetCategories(self, guid):
    """Returns the event classes (categories) that match the given GUID.

    Args:
      guid: The guid string of the categories to retrieve.

    Returns:
      A list of SWbemObject category classes.
    """
    categories = self._services.SubclassesOf(
        'EventTrace', win32com.client.constants.wbemFlagUseAmendedQualifiers)

    ret = []
    for category in categories:
      if self._GetQualifier(category, 'Guid') == guid:
        ret.append(category)

    return ret

  def _GetEvents(self, category):
    """Returns the event subclasses of the given category class.

    Args:
      category: The SWbemObject category class to get event subclasses for.

    Returns:
      A list of SWbemObject event classes.
    """
    return self._services.SubclassesOf(
        category.Path_.Class,
        win32com.client.constants.wbemFlagUseAmendedQualifiers)

  def _GetQualifier(self, obj, name):
    """Helper function to get a qualifer that may not exist.

    Args:
      obj: The SWbemObject class on which to query for a qualifier.
      name: The name of the qualifier.

    Returns:
      The value (mixed type) of the named qualifier or None if the named
      qualifier doesn't exist.
    """
    try:
      return obj.Qualifiers_.Item(name).Value
    except pywintypes.com_error:
      return None

  def Generate(self, guid, output_dir):
    """Generates a descriptor file for guid and stores it in output_dir.

    Args:
      guid: The string guid of the MOF event category class.
      output_dir: The output directory to save the generated descriptors in.
    """
    guid = guid.lower()

    lines = []
    lines.append(LICENSE_HEADER)
    lines.append(DO_NOT_EDIT_HEADER % guid)
    lines.append('')
    lines.append('from etw.descriptors import event')
    lines.append('from etw.descriptors import field')

    categories = self._GetCategories(guid)
    lines.append('')
    lines.append('')
    lines.append(self._GenerateEventTypeClass(guid, categories))

    for category in categories:
      lines.append('')
      lines.append('')
      lines.append(self._GenerateCategoryClass(category))
    lines.append('')

    # All categories of the same GUID have the same display name.
    name = categories[0].Qualifiers_.Item('DisplayName').Value.lower()
    f = open(os.path.join(output_dir, '%s.py' % name), 'wb')
    f.write('\n'.join(lines))
    f.close()

  def _GenerateEventTypeClass(self, guid, categories):
    """Generates the event type class definition for the given categories.

    Args:
      guid: The string guid of the event category.
      categories: A list of SWbemObject category classes that match the guid.

    Returns:
      A string of the generated code for the event type class.
    """
    lines = []
    lines.append('class Event(object):')
    lines.append('  GUID = \'%s\'' % guid)

    event_types = dict()
    for category in categories:
      for event in self._GetEvents(category):
        event_types.update(self._GetEventTypes(event))

    keys = event_types.keys()
    keys.sort()
    for key in keys:
      lines.append('  %s = (GUID, %d)' % (event_types[key], key))

    return '\n'.join(lines)

  def _GetEventTypes(self, event):
    """Gets a dict of event type to name for the given event.

    Args:
      event: The SWbemObject event class to get event types for.

    Returns:
      A dict of event type (number) to event type name (string).
    """
    event_types = event.Qualifiers_.Item('EventType').Value
    event_type_names = event.Qualifiers_.Item('EventTypeName').Value

    # The above qualifiers can be int or list, so we convert ints to list.
    if type(event_types) is int:
      event_types = [event_types]
      event_type_names = [event_type_names]

    return dict(zip(event_types, event_type_names))

  def _GenerateCategoryClass(self, category):
    """Generates the category class definition for the given category.

    Args:
      category: The SWbemObject category class to generate a definition for.

    Returns:
      A string of the generated event category class definition.
    """
    class_name = category.Path_.Class
    lines = []
    lines.append('class %s(event.EventCategory):' % class_name)
    lines.append('  GUID = Event.GUID')
    lines.append('  VERSION = %d' %
                 self._GetQualifier(category, 'EventVersion') or 0)

    for event in self._GetEvents(category):
      lines.append('')
      lines.append(self._GenerateEventClass(class_name, event))

    return '\n'.join(lines)

  def _GenerateEventClass(self, category_name, event):
    """Generates the event class definition for the given event.

    Args:
      category_name: The string name of the parent category class.
      event: The SWbemObject event class to generate a definition for.

    Returns:
      A string of the generated event class definition.
    """
    # Class may be named "EventCategoryName_EventClassName" or "EventClassName".
    # We just want the event class name, so we strip off the category name if
    # it's there.
    class_name = event.Path_.Class
    index = class_name.find(category_name + '_')
    if index == 0:
      class_name = class_name[len(category_name) + 1:]

    lines = []
    lines.append('  class %s(event.EventClass):' % class_name)
    lines.append('    _event_types_ = [%s]' % self._GenerateEventTypes(event))
    lines.append('    _fields_ = [%s]' % self._GenerateFields(event))

    return '\n'.join(lines)

  def _GenerateEventTypes(self, event):
    """Generates the event types as a comma delimited string for event.

    Args:
      event: The SWbemObject event class to generate an event types definition
        for.

    Returns:
      A string of the generated event types definition.
    """
    event_types = self._GetEventTypes(event).values()
    event_types.sort()
    event_types = ['Event.%s' % str(event_type) for event_type in event_types]
    return ',\n                     '.join(event_types)

  def _GenerateFields(self, event):
    """Generates the fields as comma delimited tuples for event.

    Args:
      event: The SWbemObject event class to generate a fields definition for.

    Returns:
      A string of the generated fields definition.
    """
    lines = []
    fields = self._GetFields(event)
    keys = fields.keys()
    keys.sort()
    for key in keys:
      field = fields[key]
      lines.append('(\'%s\', field.%s)' %
                   (field.Name, self._GetFieldType(field)))

    return ',\n                '.join(lines)

  def _GetFields(self, event):
    """Gets the fields of the given event.

    Args:
      event: The SWbemObject event class to get the fields for.

    Returns:
      A dict of numerical index to SWbemProperty field object.
    """
    # Normally we could just return a list here, but the properties don't come
    # out in the right order. We return a dict with the WmiDataId property
    # as the key instead.
    fields = dict()
    for prop in event.Properties_:
      index = self._GetQualifier(prop, 'WmiDataId')
      if index:
        fields[index] = prop
    return fields

  def _GetFieldType(self, field):
    """Gets the buffer reader function to call for the given property.

    Args:
      field: The SWbemProperty field object to get the field type for.

    Returns:
      The type of the field as a string.

    Raises:
      ValueError: Raised if the field's CIMType is not handled.
    """
    const = win32com.client.constants
    if field.CIMType == const.wbemCimtypeBoolean:
      return 'Boolean'
    elif field.CIMType == const.wbemCimtypeSint8:
      return 'Int8'
    elif field.CIMType == const.wbemCimtypeUint8:
      return 'UInt8'
    elif field.CIMType == const.wbemCimtypeSint16:
      return 'Int16'
    elif field.CIMType == const.wbemCimtypeUint16:
      return 'UInt16'
    elif field.CIMType == const.wbemCimtypeSint32:
      return 'Int32'
    elif field.CIMType == const.wbemCimtypeUint32:
      pointer = self._GetQualifier(field, 'pointer')
      if pointer:
        return 'Pointer'
      else:
        return 'UInt32'
    elif field.CIMType == const.wbemCimtypeSint64:
      return 'Int64'
    elif field.CIMType == const.wbemCimtypeUint64:
      return 'UInt64'
    elif field.CIMType == const.wbemCimtypeString:
      str_format = self._GetQualifier(field, 'format')
      if str_format == 'w':
        return 'WString'
      else:
        return 'String'
    elif field.CIMType == const.wbemCimtypeObject:
      ext = field.Qualifiers_.Item('extension').Value
      if ext == 'Sid':
        return 'Sid'
      elif ext == 'SizeT':
        return 'Int32'
      elif ext == 'WmiTime':
        return 'WmiTime'
      raise ValueError('Field %s is of unhandled object type: %s' %
                       (field.Name, ext))

    raise ValueError('Field %s is of unhandled type: %s' %
                     (field.Name, field.CIMType))


def main():
  parser = optparse.OptionParser()
  parser.add_option('-g', '--guid', dest='guid',
                    help='GUID of the event category to generate a paser for.')
  parser.add_option('-o', '--output_dir', dest='output_dir', default='.',
                    help='Output directory for the generated file.')
  options, unused_args = parser.parse_args()

  if not options.guid:
    print 'No guid specified.'
    return

  generator = DescriptorGenerator()
  generator.Generate(options.guid, options.output_dir)


if __name__ == '__main__':
  main()