import struct
from inspect import isclass
from StringIO import StringIO
import vstruct.primitives as vs_prims
def isVstructType(x):
[docs] return isinstance(x, vs_prims.v_base)
class VStruct(vs_prims.v_base):
[docs]
'''
The VStruct class is the bases for all groups of primitive fields which define a "structure".
Fields may be added with vsAddField() or simply added as attributes (provided you use a VStruct
or one of the vstruct.primitives in the initial assignment.)
Example:
import vstruct
from vstruct.primitives import *
vs = vstruct.VStruct()
vs.fieldone = v_uint32()
vs.fieldtwo = v_str(size=30)
bytes = vs.vsEmit()
'''
def __init__(self):
# A tiny bit of evil...
object.__setattr__(self, "_vs_values", {})
vs_prims.v_base.__init__(self)
self._vs_name = self.__class__.__name__
self._vs_fields = []
self._vs_field_align = False # To toggle visual studio style packing
self._vs_padnum = 0
self._vs_pcallbacks = {}
def vsAddParseCallback(self, fieldname, callback):
[docs] '''
Register a callback which will be triggered when the field with the given name
is set by the parser. This can be used to simplify auto-parsing to change fields
sizes or whatnot during parsing.
(You may also name a method pcb_<FieldName> to get a callback for your struct.)
Example:
def updateLengthTarget(vs):
dostuff()
v.vsAddParseCallback('lenfield', updateLengthTarget)
'''
if self._vs_values.get(fieldname) == None:
raise Exception('Invalid Field: %s' % fieldname)
cblist = self._vs_pcallbacks.get(fieldname)
if cblist == None:
cblist = []
self._vs_pcallbacks[fieldname] = cblist
cblist.append(callback)
def vsGetClassPath(self):
[docs] '''
Return the entire class name (including module path).
'''
return '%s.%s' % (self.__module__, self._vs_name)
def vsParseFd(self, fd):
[docs] '''
Parse from the given file like object as input.
'''
for fname in self._vs_fields:
fobj = self._vs_values.get(fname)
fobj.vsParseFd(fd)
callback = getattr(self, 'pcb_%s' % fname, None)
if callback != None:
callback()
cblist = self._vs_pcallbacks.get(fname)
if cblist != None:
for callback in cblist:
callback(self)
def vsParse(self, sbytes, offset=0):
[docs] """
For all the primitives contained within, allow them
an opportunity to parse the given data and return the
total offset...
Any method named pcb_<FieldName> will be called back when the specified
field is set by the parser.
"""
# In order for callbacks to change fields, we can't use vsGetFields()
for fname in self._vs_fields:
#if fname == "Magic": import pdb; pdb.set_trace()
#if fname == "MajorOperatingSystemVersion": import pdb; pdb.set_trace()
fobj = self._vs_values.get(fname)
#print fname + " " + str(offset) + " " + str(type(fobj))
offset = fobj.vsParse(sbytes, offset=offset)
callback = getattr(self, 'pcb_%s' % fname, None)
if callback != None:
callback()
cblist = self._vs_pcallbacks.get(fname)
if cblist != None:
for callback in cblist:
callback(self)
return offset
def vsEmit(self):
[docs] """
Get back the byte sequence associated with this structure.
"""
# FIXME....
ret = ''
for fname, fobj in self.vsGetFields():
ret += fobj.vsEmit()
return ret
def vsCalculate(self):
[docs] '''
Calculate fields which need correction before emitting bytes etc...
(VStruct extenders may call this, then modify fields internally)
'''
for fname, fobj in self.vsGetFields():
fobj.vsCalculate()
def vsIsPrim(self):
[docs] return False
def vsGetFields(self):
[docs] '''
Get a list of (fieldname, fieldobj) tuples for all the kids
in this VStruct (non-recursive)
Example:
for kidname, kidobj in x.vsGetFields():
print kidname
'''
ret = []
for fname in self._vs_fields:
fobj = self._vs_values.get(fname)
ret.append((fname,fobj))
return ret
def vsGetField(self, name):
[docs] x = self._vs_values.get(name)
if x == None:
raise Exception("Invalid field: %s" % name)
return x
def vsHasField(self, name):
[docs] '''
Test weather this structure contains a field with the
given name....
Example:
if x.vsHasField('woot'):
print 'STRUCT HAS WOOT FIELD!'
'''
return self._vs_values.get(name) != None
def vsSetField(self, name, value):
[docs] '''
Mostly for internal use...
'''
if isVstructType(value):
self._vs_values[name] = value
return
x = self._vs_values.get(name)
return x.vsSetValue(value)
# FIXME implement more arithmetic for structs...
def __ixor__(self, other):
for name,value in other._vs_values.items():
self._vs_values[name] ^= value
return self
def vsClearFields(self):
[docs] '''
Clear all fields from the current vstruct object. This may be useful
in specialized parsers which populate their structure on vsParse()
'''
self.__init__()
def vsAddField(self, name, value):
[docs] if not isVstructType(value):
raise Exception("Added fields MUST be vstruct types!")
# Do optional field alignment...
if self._vs_field_align:
# If it's a primitive, all is well, if not, pad to size of
# the first element of the VStruct/VArray...
if value.vsIsPrim():
align = value._vs_align
if align == None:
align = len(value)
else:
fname = value._vs_fields[0]
field = value._vs_values.get(fname)
align = field._vs_align
if align == None:
align = len(field)
delta = len(self) % align
if delta != 0:
pname = "_pad%d" % self._vs_padnum
self._vs_padnum += 1
self._vs_fields.append(pname)
self._vs_values[pname] = vs_prims.v_bytes(align-delta)
self._vs_fields.append(name)
self._vs_values[name] = value
def vsInsertField(self, name, value, befname):
[docs] '''
WARNING: vsInsertField does NOT honor field alignment! # FIXME
(AND CAN MESS UP OTHER FIELDS ALIGNMENT!)
'''
if not isVstructType(value):
raise Exception("Added fields MUST be vstruct types!")
idx = self._vs_fields.index(befname)
self._vs_fields.insert(idx, name)
self._vs_values[name] = value
def vsGetPrims(self):
[docs] """
return an order'd list of the primitive fields in this
structure definition. This is recursive and will return
the sub fields of all nested structures.
"""
ret = []
for name, field in self.vsGetFields():
if field.vsIsPrim():
ret.append(field)
else:
ret.extend(field.vsGetPrims())
return ret
def vsGetTypeName(self):
[docs] return self._vs_name
def vsGetOffset(self, name):
[docs] """
Return the offset of a member (by name):
"""
offset = 0
for fname in self._vs_fields:
if name == fname:
return offset
x = self._vs_values.get(fname)
offset += len(x)
raise Exception("Invalid Field Specified!")
def vsGetPrintInfo(self, offset=0, indent=0, top=True):
[docs] ret = []
if top:
ret.append((offset, indent, self._vs_name, self))
off = offset
indent += 1
for fname in self._vs_fields:
x = self._vs_values.get(fname)
#off = offset + self.vsGetOffset(fname)
if isinstance(x, VStruct):
ret.append((off, indent, fname, x))
ret.extend(x.vsGetPrintInfo(offset=off, indent=indent, top=False))
else:
ret.append((off, indent, fname, x))
off += len(x)
return ret
def __len__(self):
ret = 0
for fname, fobj in self.vsGetFields():
ret += len(fobj)
return ret
def __getattr__(self, name):
# Gotta do this for pickle issues...
vsvals = self.__dict__.get("_vs_values")
if vsvals == None:
vsvals = {}
self.__dict__["_vs_values"] = vsvals
r = vsvals.get(name)
if r is None:
raise AttributeError(name)
if isinstance(r, vs_prims.v_prim):
return r.vsGetValue()
return r
def __setattr__(self, name, value):
# If we have this field, asign to it
x = self._vs_values.get(name, None)
if x != None:
return self.vsSetField(name, value)
# If it's a vstruct type, create a new field
if isVstructType(value):
return self.vsAddField(name, value)
# Fail over to standard object attribute behavior
return object.__setattr__(self, name, value)
def __iter__(self):
# Our iteration returns name,field pairs
ret = []
for name in self._vs_fields:
ret.append((name, self._vs_values.get(name)))
return iter(ret)
def __repr__(self):
return '<[%s] %s>' % ("VStruct", self._vs_name)
def tree(self, va=0, reprmax=None):
[docs] ret = ""
for off, indent, name, field in self.vsGetPrintInfo():
rstr = field.vsGetTypeName()
if isinstance(field, vs_prims.v_number):
val = field.vsGetValue()
rstr = '0x%.8x (%d)' % (val,val)
elif isinstance(field, vs_prims.v_prim):
rstr = repr(field)
if reprmax != None and len(rstr) > reprmax:
rstr = rstr[:reprmax] + '...'
ret += "%.8x (%.2d)%s %s: %s\n" % (va+off, len(field), " "*(indent*2),name,rstr)
return ret
class VArray(VStruct):
[docs]
def __init__(self, elems=()):
VStruct.__init__(self)
for e in elems:
self.vsAddElement(e)
def vsAddElement(self, elem):
[docs] """
Used to add elements to an array
"""
idx = len(self._vs_fields)
self.vsAddField("%d" % idx, elem)
def __getitem__(self, index):
return self.vsGetField("%d" % index)
#FIXME slice asignment
class VUnion(VStruct):
[docs]
def vsEmit(self):
[docs] raise Exception('VUnion is only for parse right now!')
def vsParse(self, sbytes, offset=0):
[docs] """
For all the primitives contained within, allow them
an opportunity to parse the given data and return the
total offset...
Any method named pcb_<FieldName> will be called back when the specified
field is set by the parser.
"""
# In order for callbacks to change fields, we can't use vsGetFields()
ret = offset
for fname in self._vs_fields:
fobj = self._vs_values.get(fname)
ret = max(offset, fobj.vsParse(sbytes, offset=offset))
callback = getattr(self, 'pcb_%s' % fname, None)
if callback != None:
callback()
cblist = self._vs_pcallbacks.get(fname)
if cblist != None:
for callback in cblist:
callback(self)
return ret
def __len__(self):
ret = 0
for fname, fobj in self.vsGetFields():
ret = max(ret, len(fobj))
return ret
def vsGetPrintInfo(self, offset=0, indent=0, top=True):
[docs] ret = []
if top:
ret.append((offset, indent, self._vs_name, self))
indent += 1
for fname in self._vs_fields:
x = self._vs_values.get(fname)
if isinstance(x, VStruct):
ret.append((offset, indent, fname, x))
ret.extend(x.vsGetPrintInfo(offset=offset, indent=indent, top=False))
else:
ret.append((offset, indent, fname, x))
return ret
def resolve(impmod, nameparts):
[docs] """
Resolve the given (potentially nested) object
from within a module.
"""
if not nameparts:
return None
m = impmod
for nname in nameparts:
m = getattr(m, nname, None)
if m == None:
break
return m
def resolvepath(impmod, pathstr):
[docs] '''
Resolve an object/module from within the given module
by path name (ie. 'foo.bar.baz')
Example: x = resolvepath(vstruct.defs, 'win32.SEH_SCOPETABLE')
'''
nameparts = pathstr.split('.')
return resolve(impmod, nameparts)
# NOTE: Gotta import this *after* VStruct/VSArray defined
import vstruct.defs as vs_defs
def getStructure(sname):
[docs] """
Return an instance of the specified structure. The
structure name may be a definition that was added with
addStructure() or a python path (ie. win32.TEB) of a
definition from within vstruct.defs.
"""
x = resolve(vs_defs, sname.split("."))
if x != None:
return x()
return None
def getModuleNames():
[docs] return [x for x in dir(vs_defs) if not x.startswith("__")]
def getStructNames(modname):
[docs] ret = []
mod = resolve(vs_defs, modname)
if mod == None:
return ret
for n in dir(mod):
x = getattr(mod, n)
if isclass(x) and issubclass(x, VStruct):
ret.append(n)
return ret