miasm/loader/new_cstruct.py
#! /usr/bin/env python
from __future__ import print_function
import re
import struct
from miasm.core.utils import force_bytes
from future.utils import PY3, viewitems, with_metaclass
type2realtype = {}
size2type = {}
size2type_s = {}
for t in 'B', 'H', 'I', 'Q':
s = struct.calcsize(t)
type2realtype[t] = s * 8
size2type[s * 8] = t
for t in 'b', 'h', 'i', 'q':
s = struct.calcsize(t)
type2realtype[t] = s * 8
size2type_s[s * 8] = t
type2realtype['u08'] = size2type[8]
type2realtype['u16'] = size2type[16]
type2realtype['u32'] = size2type[32]
type2realtype['u64'] = size2type[64]
type2realtype['s08'] = size2type_s[8]
type2realtype['s16'] = size2type_s[16]
type2realtype['s32'] = size2type_s[32]
type2realtype['s64'] = size2type_s[64]
type2realtype['d'] = 'd'
type2realtype['f'] = 'f'
type2realtype['q'] = 'q'
type2realtype['ptr'] = 'ptr'
sex_types = {0: '<', 1: '>'}
def fix_size(fields, wsize):
out = []
for name, v in fields:
if v.endswith("s"):
pass
elif v == "ptr":
v = size2type[wsize]
elif not v in type2realtype:
raise ValueError("unknown Cstruct type", v)
else:
v = type2realtype[v]
out.append((name, v))
fields = out
return fields
def real_fmt(fmt, wsize):
if fmt == "ptr":
v = size2type[wsize]
elif fmt in type2realtype:
v = type2realtype[fmt]
else:
v = fmt
return v
all_cstructs = {}
class Cstruct_Metaclass(type):
field_suffix = "_value"
def __new__(cls, name, bases, dct):
for fields in dct['_fields']:
fname = fields[0]
if fname in ['parent', 'parent_head']:
raise ValueError('field name will confuse internal structs',
repr(fname))
dct[fname] = property(dct.pop("get_" + fname,
lambda self, fname=fname: getattr(
self, fname + self.__class__.field_suffix)),
dct.pop("set_" + fname,
lambda self, v, fname=fname: setattr(
self, fname + self.__class__.field_suffix, v)),
dct.pop("del_" + fname, None))
o = super(Cstruct_Metaclass, cls).__new__(cls, name, bases, dct)
if name != "CStruct":
all_cstructs[name] = o
return o
def unpack_l(cls, s, off=0, parent_head=None, _sex=None, _wsize=None):
if _sex is None and _wsize is None:
# get sex and size from parent
if parent_head is not None:
_sex = parent_head._sex
_wsize = parent_head._wsize
else:
_sex = 0
_wsize = 32
c = cls(_sex=_sex, _wsize=_wsize)
if parent_head is None:
parent_head = c
c.parent_head = parent_head
of1 = off
for field in c._fields:
cpt = None
if len(field) == 2:
fname, ffmt = field
elif len(field) == 3:
fname, ffmt, cpt = field
if ffmt in type2realtype or (isinstance(ffmt, str) and re.match(r'\d+s', ffmt)):
# basic types
if cpt:
value = []
i = 0
while i < cpt(c):
fmt = real_fmt(ffmt, _wsize)
of2 = of1 + struct.calcsize(fmt)
value.append(struct.unpack(c.sex + fmt, s[of1:of2])[0])
of1 = of2
i += 1
else:
fmt = real_fmt(ffmt, _wsize)
of2 = of1 + struct.calcsize(fmt)
if not (0 <= of1 < len(s) and 0 <= of2 < len(s)):
raise RuntimeError("not enough data")
value = struct.unpack(c.sex + fmt, s[of1:of2])[0]
elif ffmt == "sz": # null terminated special case
of2 = s.find(b'\x00', of1)
if of2 == -1:
raise ValueError('no null char in string!')
of2 += 1
value = s[of1:of2 - 1]
elif ffmt in all_cstructs:
of2 = of1
# sub structures
if cpt:
value = []
i = 0
while i < cpt(c):
v, l = all_cstructs[ffmt].unpack_l(
s, of1, parent_head, _sex, _wsize)
v.parent = c
value.append(v)
of2 = of1 + l
of1 = of2
i += 1
else:
value, l = all_cstructs[ffmt].unpack_l(
s, of1, parent_head, _sex, _wsize)
value.parent = c
of2 = of1 + l
elif isinstance(ffmt, tuple):
f_get, f_set = ffmt
value, of2 = f_get(c, s, of1)
else:
raise ValueError('unknown class', ffmt)
of1 = of2
setattr(c, fname + c.__class__.field_suffix, value)
return c, of2 - off
def unpack(cls, s, off=0, parent_head=None, _sex=None, _wsize=None):
c, l = cls.unpack_l(s, off=off,
parent_head=parent_head, _sex=_sex, _wsize=_wsize)
return c
class CStruct(with_metaclass(Cstruct_Metaclass, object)):
_packformat = ""
_fields = []
def __init__(self, parent_head=None, _sex=None, _wsize=None, **kargs):
self.parent_head = parent_head
self._size = None
kargs = dict(kargs)
# if not sex or size: get the one of the parent
if _sex == None and _wsize == None:
if parent_head:
_sex = parent_head._sex
_wsize = parent_head._wsize
else:
# else default sex & size
_sex = 0
_wsize = 32
# _sex is 0 or 1, sex is '<' or '>'
self._sex = _sex
self._wsize = _wsize
if self._packformat:
self.sex = self._packformat
else:
self.sex = sex_types[_sex]
for f in self._fields:
setattr(self, f[0] + self.__class__.field_suffix, None)
if kargs:
for k, v in viewitems(kargs):
self.__dict__[k + self.__class__.field_suffix] = v
def pack(self):
out = b''
for field in self._fields:
cpt = None
if len(field) == 2:
fname, ffmt = field
elif len(field) == 3:
fname, ffmt, cpt = field
value = getattr(self, fname + self.__class__.field_suffix)
if ffmt in type2realtype or (isinstance(ffmt, str) and re.match(r'\d+s', ffmt)):
# basic types
fmt = real_fmt(ffmt, self._wsize)
if cpt == None:
if value == None:
o = struct.calcsize(fmt) * b"\x00"
elif ffmt.endswith('s'):
new_value = force_bytes(value)
o = struct.pack(self.sex + fmt, new_value)
else:
o = struct.pack(self.sex + fmt, value)
else:
o = b""
for v in value:
if value == None:
o += struct.calcsize(fmt) * b"\x00"
else:
o += struct.pack(self.sex + fmt, v)
elif ffmt == "sz": # null terminated special case
o = value + b'\x00'
elif ffmt in all_cstructs:
# sub structures
if cpt == None:
o = bytes(value)
else:
o = b""
for v in value:
o += bytes(v)
elif isinstance(ffmt, tuple):
f_get, f_set = ffmt
o = f_set(self, value)
else:
raise ValueError('unknown class', ffmt)
out += o
return out
def __bytes__(self):
return self.pack()
def __str__(self):
if PY3:
return repr(self)
return self.__bytes__()
def __len__(self):
return len(self.pack())
def __repr__(self):
return "<%s=%s>" % (self.__class__.__name__, "/".join(
repr(getattr(self, x[0])) for x in self._fields)
)
def __getitem__(self, item): # to work with format strings
return getattr(self, item)