Skip to content

Commit ba269d3

Browse files
committed
add/fix support for struct "slop" and allow for self-defined alignements
1 parent 162ec00 commit ba269d3

File tree

4 files changed

+150
-65
lines changed

4 files changed

+150
-65
lines changed

amoco/system/imx6.py

+11-11
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@
8686
B : tag
8787
H :> length
8888
B : version
89-
""")
89+
""",packed=True)
9090
class HAB_Header(StructFormatter):
9191
order = '<'
9292
def __init__(self,data="",offset=0):
@@ -109,7 +109,7 @@ def token_ver_format(k,x,cls=None):
109109
I : self
110110
I : csf
111111
I : reserved2
112-
""")
112+
""",packed=True)
113113
class IVT(StructFormatter):
114114
order = '<'
115115
def __init__(self,data="",offset=0):
@@ -127,7 +127,7 @@ def __init__(self,data="",offset=0):
127127
I : start
128128
I : size
129129
I : plugin
130-
""")
130+
""",packed=True)
131131
class BootData(StructFormatter):
132132
order = '<'
133133
def __init__(self,data="",offset=0):
@@ -146,7 +146,7 @@ def __init__(self,data="",offset=0):
146146
I : unk_flags
147147
H :> nlen
148148
H :> elen
149-
""")
149+
""",packed=True)
150150
class PublicKey(StructFormatter):
151151
order = '<'
152152
def __init__(self,data="",offset=0):
@@ -169,7 +169,7 @@ def size(self):
169169

170170
@StructDefine("""
171171
HAB_Header : header
172-
""")
172+
""",packed=True)
173173
class CRT(StructFormatter):
174174
def __init__(self,data="",offset=0):
175175
self.keys = []
@@ -200,7 +200,7 @@ def size(self):
200200

201201
@StructDefine("""
202202
HAB_Header : header
203-
""")
203+
""",packed=True)
204204
class CSF(StructFormatter):
205205
def __init__(self,data="",offset=0):
206206
self.cmds = []
@@ -248,7 +248,7 @@ def CMD(data,offset):
248248
B : par
249249
I : address
250250
I : mask
251-
""")
251+
""",packed=True)
252252
class CheckData(StructFormatter):
253253
order = '<'
254254
def __init__(self,data="",offset=0):
@@ -279,7 +279,7 @@ def size(self):
279279
B : cmd
280280
H :> len
281281
B : und
282-
""")
282+
""",packed=True)
283283
class NOP(StructFormatter):
284284
order = '<'
285285
def __init__(self,data="",offset=0):
@@ -291,7 +291,7 @@ def __init__(self,data="",offset=0):
291291
B : cmd
292292
H :> len
293293
B : eng
294-
""")
294+
""",packed=True)
295295
class Unlock(StructFormatter):
296296
order = '<'
297297
def __init__(self,data="",offset=0):
@@ -325,7 +325,7 @@ def size(self):
325325
B : src
326326
B : tgt
327327
I :> key_dat
328-
""")
328+
""",packed=True)
329329
class InstallKey(StructFormatter):
330330
order = '<'
331331
def __init__(self,data="",offset=0):
@@ -362,7 +362,7 @@ def size(self):
362362
B : eng
363363
B : cfg
364364
I :> aut_start
365-
""")
365+
""",packed=True)
366366
class Authenticate(StructFormatter):
367367
order = '<'
368368
def __init__(self,data="",offset=0):

amoco/system/structs.py

+118-49
Original file line numberDiff line numberDiff line change
@@ -164,27 +164,33 @@ class Field(object):
164164
pack (value,order='<') : packs the value with the given order and returns the
165165
byte string according to type typename.
166166
"""
167-
def __init__(self,ftype,fcount=0,fname=None,forder=None,fcomment=''):
167+
def __init__(self,ftype,fcount=0,fname=None,forder=None,falign=0,fcomment=''):
168168
self.typename = ftype
169169
self.type_private = isinstance(ftype,(StructCore))
170170
self.count = fcount
171171
self.name = fname
172-
self.order = forder
172+
self.order = forder or '<'
173+
self._align_value = None
174+
if falign:
175+
self.align_value = falign
173176
self.comment = fcomment
174177
@property
175178
def type(self):
176179
if self.type_private:
177180
return self.typename
178-
cls = StructDefine.All[self.typename]
179-
return cls()
181+
try:
182+
cls = StructDefine.All[self.typename]
183+
except KeyError:
184+
return None
185+
else:
186+
return cls()
180187
def format(self):
181-
fmt = self.type.format()
182-
sz = struct.calcsize(fmt)
183-
if self.count>0: sz = sz*self.count
188+
sz = self.size()
184189
return '%ds'%sz
185-
@property
186190
def size(self):
187-
return struct.calcsize(self.format())
191+
sz = self.type.size()
192+
if self.count>0: sz = sz*self.count
193+
return sz
188194
@property
189195
def source(self):
190196
res = "%s"%self.typename
@@ -193,29 +199,39 @@ def source(self):
193199
if self.comment: res += " ;%s"%self.comment
194200
return res
195201
def __len__(self):
196-
return self.size
197-
def unpack(self,data,offset=0,order='<'):
198-
if self.order: order=self.order
199-
blob = self.type.unpack(data,offset,order)
200-
sz = len(self.type)
202+
return self.size()
203+
@property
204+
def align_value(self):
205+
if self._align_value: return self._align_value
206+
if isinstance(self.type,Field): return self.type.align_value
207+
return self.type.align_value()
208+
@align_value.setter
209+
def align_value(self,val):
210+
self._align_value = val
211+
def align(self,offset):
212+
A = self.align_value
213+
r = offset%A
214+
if r==0: return offset
215+
return offset+(A-r)
216+
def unpack(self,data,offset=0):
217+
blob = self.type.unpack(data,offset)
218+
sz = self.type.size()
201219
count = self.count
202220
if count>0:
203221
blob = [blob]
204222
count -= 1
205223
offset += sz
206224
while count>0:
207-
blob.append(self.type.unpack(data,offset,order))
225+
blob.append(self.type.unpack(data,offset))
208226
offset += sz
209227
count -= 1
210228
return blob
211-
def get(self,data,offset=0,order='<'):
212-
if self.order: order=self.order
213-
return (self.name,self.unpack(data,offset,order))
214-
def pack(self,value,order='<'):
215-
if self.order: order=self.order
229+
def get(self,data,offset=0):
230+
return (self.name,self.unpack(data,offset))
231+
def pack(self,value):
216232
if self.count>0:
217-
return b''.join([struct.pack(order+self.format(),v) for v in value])
218-
return struct.pack(order+self.format(),value)
233+
return b''.join([self.type.pack(v) for v in value])
234+
return self.type.pack(value)
219235
def __call__(self):
220236
return self
221237
def __repr__(self):
@@ -237,14 +253,20 @@ class RawField(Field):
237253
def format(self):
238254
fmt = self.typename
239255
if self.count==0: return fmt
240-
sz = struct.calcsize(fmt)*self.count
241-
return '%ds'%sz
242-
def unpack(self,data,offset=0,order='<'):
243-
if self.order: order=self.order
256+
return '%d%s'%(sz,fmt)
257+
def size(self):
258+
sz = struct.calcsize(self.typename)
259+
if self.count>0: sz = sz*self.count
260+
return sz
261+
def unpack(self,data,offset=0):
244262
pfx = '%d'%self.count if self.count>0 else ''
245-
res = struct.unpack(order+pfx+self.typename,data[offset:offset+self.size])
263+
res = struct.unpack(self.order+pfx+self.typename,data[offset:offset+self.size()])
246264
if self.count==0 or self.typename=='s': return res[0]
247265
return res
266+
def pack(self,value):
267+
pfx = '%d'%self.count if self.count>0 else ''
268+
res = struct.pack(self.order+pfx+self.typename,value)
269+
return res
248270
def __repr__(self):
249271
fmt = self.typename
250272
r = '<Field %s [%s]'%(self.name,fmt)
@@ -260,6 +282,11 @@ class StructDefine(object):
260282
"""
261283
All = {}
262284
rawtypes = ('x','c','b','B','h','H','i','I','l','L','f','d','s','n','N','p','P','q','Q')
285+
alignments = {'x':1, 'c':1, 'b':1, 'B':1, 's':1,
286+
'h':2, 'H':2,
287+
'i':4, 'I':4, 'l':4, 'L':4, 'f':4,
288+
'q':8, 'Q':8, 'd':8,
289+
'P':8}
263290
integer = pp.Regex(r'[1-9][0-9]*')
264291
number = integer
265292
number.setParseAction(lambda r: int(r[0]))
@@ -273,17 +300,28 @@ class StructDefine(object):
273300
def __init__(self,fmt,**kargs):
274301
self.fields = []
275302
self.source = fmt
303+
self.packed = kargs.get('packed',False)
304+
if 'alignments' in kargs:
305+
self.alignments = kargs['alignments']
276306
for l in self.structfmt.parseString(fmt,True):
277307
f_type,f_name,f_comment = l
278308
f_order,f_name = f_name
279309
f_type,f_count = f_type
280-
f_cls = RawField if f_type in self.rawtypes else Field
281-
if f_type in kargs: f_type = kargs[f_type]
282-
self.fields.append(f_cls(f_type,f_count,f_name,f_order,f_comment))
310+
if f_order is None and 'order' in kargs:
311+
f_order=kargs['order']
312+
if f_type in self.rawtypes:
313+
f_cls = RawField
314+
f_align = self.alignments[f_type]
315+
else:
316+
f_cls = Field
317+
f_type = kargs.get(f_type,f_type)
318+
f_align = 0
319+
self.fields.append(f_cls(f_type,f_count,f_name,f_order,f_align,f_comment))
283320
def __call__(self,cls):
284321
self.All[cls.__name__] = cls
285322
cls.fields = self.fields
286323
cls.source = self.source
324+
cls.packed = self.packed
287325
return cls
288326

289327
class UnionDefine(StructDefine):
@@ -293,13 +331,18 @@ def __call__(self,cls):
293331
self.All[cls.__name__] = cls
294332
cls.fields = self.fields
295333
cls.source = self.source
296-
s = [f.size for f in cls.fields]
334+
s = [f.size() for f in cls.fields]
297335
cls.union = s.index(max(s))
298336
return cls
299337

300-
def TypeDefine(newname, typebase, typecount=0):
301-
f_cls = RawField if typebase in StructDefine.rawtypes else Field
302-
StructDefine.All[newname] = f_cls(typebase,fcount=typecount,fname='typedef')
338+
def TypeDefine(newname, typebase, typecount=0,align_value=0):
339+
if typebase in StructDefine.rawtypes:
340+
f_cls = RawField
341+
f_align = align_value or StructDefine.alignments[typebase]
342+
else:
343+
f_cls = Field
344+
f_align = 0
345+
StructDefine.All[newname] = f_cls(typebase,fcount=typecount,falign=f_align,fname='typedef')
303346

304347
class StructCore(object):
305348
"""StructCore is a ParentClass for all user-defined structures based on a StructDefine format.
@@ -308,42 +351,68 @@ class StructCore(object):
308351
Note: It is mandatory that any class that inherits from StructCore can be instanciated
309352
with no arguments.
310353
"""
311-
order = '@'
354+
packed = False
312355
union = False
313356

314357
@classmethod
315358
def format(cls):
316359
if cls.union is False:
317-
return cls.order+(''.join((f.format() for f in cls.fields)))
360+
return ''.join((f.format() for f in cls.fields))
318361
else:
319-
return cls.order+cls.fields[cls.union].format()
362+
return cls.fields[cls.union].format()
320363
@classmethod
321364
def size(cls):
322-
return struct.calcsize(cls.format())
365+
A = cls.align_value()
366+
sz = 0
367+
for f in cls.fields:
368+
if cls.union is False and not cls.packed:
369+
sz = f.align(sz)
370+
if cls.union is False:
371+
sz += f.size()
372+
elif f.size>sz:
373+
sz = f.size()
374+
r = sz%A
375+
if (not cls.packed) and r>0:
376+
sz += (A-r)
377+
return sz
323378
def __len__(self):
324379
return self.size()
325-
def unpack(self,data,offset=0,order=None):
380+
@classmethod
381+
def align_value(cls):
382+
return max([f.align_value for f in cls.fields])
383+
def unpack(self,data,offset=0):
326384
for f in self.fields:
327-
setattr(self,f.name,f.unpack(data,offset,order or self.order))
385+
if self.union is False and not self.packed:
386+
offset = f.align(offset)
387+
setattr(self,f.name,f.unpack(data,offset))
328388
if self.union is False:
329-
offset += f.size
389+
offset += f.size()
330390
return self
331391
def pack(self,data=None):
332392
if data is None:
333393
data = [getattr(self,f.name) for f in self.fields]
334394
parts = []
395+
offset = 0
335396
for f,v in zip(self.fields,data):
336-
parts.append(f.pack(v,self.order))
397+
p = f.pack(v)
398+
if not self.packed:
399+
pad = f.align(offset)-offset
400+
p = b'\0'*pad + p
401+
parts.append(p)
337402
if self.union is False:
338-
return b''.join(parts)
403+
res = b''.join(parts)
404+
if not self.packed:
405+
res = res.ljust(self.size(),b'\0')
406+
return res
339407
else:
340408
return parts[self.union]
341409
def offset_of(self,name):
410+
if self.union is not False:
411+
return 0
342412
o = 0
343413
for f in self.fields:
344414
if f.name==name: return o
345-
if self.union is False:
346-
o += f.size
415+
o = f.align(o)
347416
raise AttributeError(name)
348417

349418
class StructFormatter(StructCore):
@@ -442,13 +511,13 @@ def define(cls,fmt,offset=-1,atvalue=None,indata=b""):
442511

443512
#------------------------------------------------------------------------------
444513

445-
def StructFactory(name,fmt):
514+
def StructFactory(name,fmt,**kargs):
446515
'Returns a StructFormatter class build with name and format'
447-
return StructDefine(fmt)(type(name,(StructFormatter,),{}))
516+
return StructDefine(fmt,**kargs)(type(name,(StructFormatter,),{}))
448517

449-
def UnionFactory(name,fmt):
518+
def UnionFactory(name,fmt,**kargs):
450519
'Returns a StructFormatter (union) class build with name and format'
451-
return UnionDefine(fmt)(type(name,(StructFormatter,),{}))
520+
return UnionDefine(fmt,**kargs)(type(name,(StructFormatter,),{}))
452521

453522
#------------------------------------------------------------------------------
454523

setup.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030

3131
setup(
3232
name = 'amoco',
33-
version = '2.6.1',
33+
version = '2.6.2',
3434
description = 'yet another binary analysis framework',
3535
long_description = long_descr,
3636
# Metadata

0 commit comments

Comments
 (0)