wireshark/epan/wspython/wspy_dissector.py

# wspy_dissector.py
#
# $Id: $
#
# Wireshark Protocol Python Binding
#
# Copyright (c) 2009 by Sebastien Tandel <sebastien [AT] tandel [dot] be>
# Copyright (c) 2001 by Gerald Combs <gerald@wireshark.org>
#
# 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; either version 2
# of the License, or (at your option) any later version.
#
# 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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

from ctypes import c_int, pointer, POINTER, py_object
from wspy_libws import get_libws_handle

# From epan/proto.h
# ? STA ? : is there a better way to include/define these constants?
# (duplicating definition is not a good thing)
(BASE_NONE,
BASE_DEC,
BASE_HEX,
BASE_OCT,
BASE_DEC_HEX,
BASE_HEX_DEC,
BASE_CUSTOM) = map(int, xrange(7))

# field types, see epan/ftypes/ftypes.h
(FT_NONE,
FT_PROTOCOL,
FT_BOOLEAN,
FT_UINT8,
FT_UINT16,
FT_UINT24,
FT_UINT32,
FT_UINT64,
FT_INT8,
FT_INT16,
FT_INT24,
FT_INT32,
FT_INT64,
FT_FLOAT,
FT_DOUBLE,
FT_ABSOLUTE_TIME,
FT_RELATIVE_TIME,
FT_STRING,
FT_STRINGZ,
FT_EBCDIC,
FT_UINT_STRING,
FT_ETHER,
FT_BYTES,
FT_UINT_BYTES,
FT_IPv4,
FT_IPv6,
FT_IPXNET,
FT_FRAMENUM,
FT_PCRE,
FT_GUID,
FT_OID) = map(int, xrange(31))

# hf_register_info from usual dissectors
class register_info(object):
  def __init__(self, wsl):
    self.__protocol = None
    self.__wsl = wsl
    self.__hf_register = None
    self.__registers = []

  def add(self, name, short_desc, \
          type=FT_UINT32, display=BASE_DEC, \
          strings=None, bitmask=0x0, desc=None):
    if not desc:
      desc = name
    self.__registers.append( (name, short_desc, \
          type, display, strings, bitmask, desc) )

  def register(self, protocol):
    self.__protocol = protocol
    hf = self.__registers
    lr = len(hf)
    if not lr:
      return None

    self.__hf_register = self.__wsl.hf_register_info_create(lr)
    chf = self.__hf_register
    if not self.__hf_register:
      return None

    for i in xrange(lr):
      n, sd, t, d, st, bm, ld = hf[i]
      sdn = sd.replace('.', '_')
      self.__dict__[sdn] = c_int(-1)
      p_id = pointer(self.__dict__[sdn])
      self.__wsl.hf_register_info_add(chf, i, p_id, n , sd, t, d, st, bm, ld)

    self.__wsl.proto_register_field_array(self.__protocol, chf, lr)

  def display(self):
    self.__wsl.hf_register_info_print(self.__hf_register, \
                                    len(self.__registers))

  def get(self):
    return self.__hf_register, len(self.__registers)

  def __del__(self):
    self.__wsl.hf_register_info_destroy(self.__hf_register)

#Subtrees definition
#Every subtree added can be accesses as an attribute after having been
#registered
class Subtree(object):
  def __init__(self, wsl, protocol):
    self.__wsl = wsl
    self.__protocol = protocol
    self.__st = {}
    self.__user_defined_protocol_tree = False

  def add(self, name):
    if name == self.__protocol:
      self.__user_defined_protocol_tree = True

    self.__st[name] = c_int(-1)

  def has_user_defined_protocol_tree(self):
    return self.__user_defined_protocol_tree

  def register(self):
    if not self.__user_defined_protocol_tree:
      self.__st[self.__protocol] = c_int(-1)

    ls = len(self.__st)
    if not ls:
      return

    CSubtrees = POINTER(c_int) * ls
    p_sts = CSubtrees()
    k = self.__st.keys()
    for i in xrange(ls):
      p_sts[i] = pointer(self.__st[k[i]])

    self.__wsl.proto_register_subtree_array(p_sts, ls)

  def __getattr__(self, name):
    if self.__st.has_key(name):
      return self.__st[name]
    #raise KeyError

#Dissector class : base class to write a dissector in python
class Dissector(object):
  def __init__(self, protocol_name, short_desc, short):
    self.__protocol_name = protocol_name
    self.__short_desc = short_desc
    self.__short = short

    self.__tvb = None
    self.__pinfo = None
    self.__tree = None

    self.__Tree = None

    self.__offset = 0

    self.__wsl = get_libws_handle()
    self.__hf = None
    self.__subtree = None

  def fields(self):
    if not self.__hf:
      self.__hf = register_info(self.__wsl)
    return self.__hf

  def subtrees(self):
    if not self.__subtree:
      self.__subtree = Subtree(self.__wsl, self.__short)
    return self.__subtree

  def tree(self):
    if not self.__Tree:
      self.__Tree = Tree(self.__tree, self)
    return self.__Tree

  def display(self):
    print self.__short

  def __str__(self):
    # STA TODO : keep with short_desc because used in the hash table of
    # dissectors in C code. If it is modified, it won't work anymore
    return self.__short_desc

  def __unicode__(self):
    return self.__short

  def __hash__(self):
    return hash(self.__short)

  def protocol(self):
    return self.__protocol

  def register_protocol(self):
    self.__protocol = \
      self.__wsl.proto_register_protocol( \
        self.__protocol_name, self.__short_desc, \
        self.__short)
    self.__hf.register(self.__protocol)
    #self.__hf.display()
    self.__subtree.register()

  def dissect(self):
    pass

  def pre_dissect(self):
    self.__tvb = self.__wsl.py_tvbuff()
    self.__pinfo = self.__wsl.py_pinfo()
    self.__tree = self.__wsl.py_tree()

    #self.__wsl.print_current_proto(py_object(pinfo))
    subt = self.subtrees()
    try:
      if not subt.has_user_defined_protocol_tree():
        tree = self.tree()
        p_tree = tree.add_item(self.protocol())
        self.__Tree = p_tree.add_subtree(subt.homeplug)
    except:
      print e
    self.dissect()

  def protocol_ids(self):
    return [ (None, 0, None) ]

  def create_dissector_handle(self, protocol=None):
    gdissector = self.__wsl.py_generic_dissector()
    if not protocol:
      protocol = self.__protocol
    return self.__wsl.create_dissector_handle(gdissector, protocol)

  def find_dissector(self, protocol):
    return self.__wsl.find_dissector(protocol)

  def register_handoff(self):
    #TODO STA : think how we would use dissector_add in an easy way *and* with
    #the possibility to add the same dissector for TCP and UDP (extend
    #py_generic_dissector)
    private_handle = None
    try:
      ids = self.protocol_ids()
      for type, protocol_id, handle in self.protocol_ids():
        if not type:
          continue
        if not handle:
          if not private_handle:
            handle = \
              self.create_dissector_handle(self.__protocol)
          else:
            handle = private_handle
        self.__wsl.dissector_add(type, protocol_id, handle)
    except Exception, e:
      print "creating dissector failed", e
      raise

  def advance(self, step):
    self.__offset += step

  def offset(self):
    return self.__offset

  def _wsl(self):
    return self.__wsl

  def _tvb(self):
    return self.__tvb

#Tree class implementation
#see proto.h
class Tree(object):
  def __init__(self, tree, dissector):
    self.__dissector = dissector
    self.__tree = tree
    self.__wsl = dissector._wsl()
    self.__tvb = dissector._tvb()

  def add_item(self, field, offset=0, length=-1, little_endian=False, adv=True):
    '''add an item to the tree'''
    try:
      tree = self.__wsl.proto_tree_add_item(self.__tree, \
        field, self.__tvb, self.__dissector.offset(), length, \
        little_endian)
    except Exception, e:
      print e
    else:
      if length > 0 and adv:
        self.__dissector.advance(length)
      return Tree(tree, self.__dissector)

  def add_uint(self, field, value, offset=0, length=4, adv=True):
    '''add unsigned integer to the tree'''
    try:
      tree = self.__wsl.proto_tree_add_uint(self.__tree, field, self.__tvb, self.__dissector.offset(), length, value)
    except Exception, e:
      print e
    else:
      if adv:
        self.__dissector.advance(length)
      return Tree(tree, self.__dissector)

  def add_text(self, string, offset=0, length=-1, adv=True):
    '''add text to the tree'''
    try:
      tree = self.__wsl.proto_tree_add_text(self.__tree, self.__tvb, self.__dissector.offset(), length, string)
    except Exception, e:
      print e
    else:
      if length > 0 and adv:
        self.__dissector.advance(length)
      return Tree(tree, self.__dissector)

  def add_subtree(self, subtree):
    '''add a subtree to the tree'''
    try:
      tree = self.__wsl.proto_item_add_subtree(self.__tree, subtree)
    except Exception, e:
      print e
    else:
      return Tree(tree, self.__dissector)

#tvb class implementation
#see proto.h
class TVB(object):
  def __init__(self, wsl, tvb, dissector):
    self.__tvb = tvb
    self.__wsl = wsl
    self.__dissector = dissector

  def length(self):
    return self.__wsl.length(self.__wsl)

  def length_remaining(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_length_remaining(self.__tvb, offset)

  def reported_length(self):
    return self.__wsl.tvb_reported_length(self.__tvb)

  def reported_length_remaining(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_length_remaining(self.__tvb, offset)

  def get_guint8(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_guint8(self.__tvb)

  def get_ntohs(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_ntohs(self.__tvb, offset)

  def get_ntohl(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_ntohl(self.__tvb, offset)

  def get_letohl(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_letohl(self.__tvb, offset)

  def get_letohs(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_letohs(self.__tvb, offset)

  #STA TODO : check that we can do that
  def get_ptr(self, offset=-1):
    if offset < 0:
      offset = self.__dissector.offset()
    return self.__wsl.tvb_get_ptr(self.__tvb, offset)

  #how to get this working ??? check how application uses this!
  #def new_subset(self, offset=0):
  #  return self.__wsl.tvb_get_new_subset(self.tvb, offset)