1 # Copyright (c) 2003-2016 CORE Security Technologies
3 # This software is provided under under a slightly modified version
4 # of the Apache Software License. See the accompanying LICENSE file
5 # for more information.
8 from struct import pack, unpack, calcsize
11 """ sublcasses can define commonHdr and/or structure.
12 each of them is an tuple of either two: (fieldName, format) or three: (fieldName, ':', class) fields.
13 [it can't be a dictionary, because order is important]
15 where format specifies how the data in the field will be converted to/from bytes (string)
16 class is the class to use when unpacking ':' fields.
18 each field can only contain one value (or an array of values for *)
19 i.e. struct.pack('Hl',1,2) is valid, but format specifier 'Hl' is not (you must use 2 dfferent fields)
22 specifiers from module pack can be used with the same format
23 see struct.__doc__ (pack/unpack is finally called)
34 q [signed long long (quad)]
35 Q [unsigned long long (quad)]
36 s [string (array of chars), must be preceded with length in format specifier, padded with zeros]
37 p [pascal string (includes byte count), must be preceded with length in format specifier, padded with zeros]
40 = [native byte ordering, size and alignment]
41 @ [native byte ordering, standard size and alignment]
42 ! [network byte ordering]
46 usual printf like specifiers can be used (if started with %)
47 [not recommeneded, there is no why to unpack this]
49 %08x will output an 8 bytes hex
50 %s will output a string
51 %s\\x00 will output a NUL terminated string
52 %d%d will output 2 decimal digits (against the very same specification of Structure)
55 some additional format specifiers:
56 : just copy the bytes from the field into the output string (input may be string, other structure, or anything responding to __str__()) (for unpacking, all what's left is returned)
57 z same as :, but adds a NUL byte at the end (asciiz) (for unpacking the first NUL byte is used as terminator) [asciiz string]
58 u same as z, but adds two NUL bytes at the end (after padding to an even size with NULs). (same for unpacking) [unicode string]
59 w DCE-RPC/NDR string (it's a macro for [ '<L=(len(field)+1)/2','"\\x00\\x00\\x00\\x00','<L=(len(field)+1)/2',':' ]
60 ?-field length of field named 'field', formated as specified with ? ('?' may be '!H' for example). The input value overrides the real length
61 ?1*?2 array of elements. Each formated as '?2', the number of elements in the array is stored as specified by '?1' (?1 is optional, or can also be a constant (number), for unpacking)
62 'xxxx literal xxxx (field's value doesn't change the output. quotes must not be closed or escaped)
63 "xxxx literal xxxx (field's value doesn't change the output. quotes must not be closed or escaped)
64 _ will not pack the field. Accepts a third argument, which is an unpack code. See _Test_UnpackCode for an example
65 ?=packcode will evaluate packcode in the context of the structure, and pack the result as specified by ?. Unpacking is made plain
66 ?&fieldname "Address of field fieldname".
67 For packing it will simply pack the id() of fieldname. Or use 0 if fieldname doesn't exists.
68 For unpacking, it's used to know weather fieldname has to be unpacked or not, i.e. by adding a & field you turn another field (fieldname) in an optional field.
75 def __init__(self, data = None, alignment = 0):
76 if not hasattr(self, 'alignment'):
77 self.alignment = alignment
87 def fromFile(self, file):
89 answer.fromString(file.read(len(answer)))
92 def setAlignment(self, alignment):
93 self.alignment = alignment
95 def setData(self, data):
98 def packField(self, fieldName, format = None):
100 print "packField( %s | %s )" % (fieldName, format)
103 format = self.formatForField(fieldName)
105 if self.fields.has_key(fieldName):
106 ans = self.pack(format, self.fields[fieldName], field = fieldName)
108 ans = self.pack(format, None, field = fieldName)
111 print "\tanswer %r" % ans
116 if self.data is not None:
119 for field in self.commonHdr+self.structure:
121 data += self.packField(field[0], field[1])
123 if self.fields.has_key(field[0]):
124 e.args += ("When packing field '%s | %s | %r' in %s" % (field[0], field[1], self[field[0]], self.__class__),)
126 e.args += ("When packing field '%s | %s' in %s" % (field[0], field[1], self.__class__),)
129 if len(data) % self.alignment:
130 data += ('\x00'*self.alignment)[:-(len(data) % self.alignment)]
132 #if len(data) % self.alignment: data += ('\x00'*self.alignment)[:-(len(data) % self.alignment)]
135 def fromString(self, data):
137 for field in self.commonHdr+self.structure:
139 print "fromString( %s | %s | %r )" % (field[0], field[1], data)
140 size = self.calcUnpackSize(field[1], data, field[0])
142 print " size = %d" % size
143 dataClassOrCode = str
145 dataClassOrCode = field[2]
147 self[field[0]] = self.unpack(field[1], data[:size], dataClassOrCode = dataClassOrCode, field = field[0])
149 e.args += ("When unpacking field '%s | %s | %r[:%d]'" % (field[0], field[1], data, size),)
152 size = self.calcPackSize(field[1], self[field[0]], field[0])
153 if self.alignment and size % self.alignment:
154 size += self.alignment - (size % self.alignment)
159 def __setitem__(self, key, value):
160 self.fields[key] = value
161 self.data = None # force recompute
163 def __getitem__(self, key):
164 return self.fields[key]
166 def __delitem__(self, key):
170 return self.getData()
174 return len(self.getData())
176 def pack(self, format, data, field = None):
178 print " pack( %s | %r | %s)" % (format, data, field)
181 addressField = self.findAddressFieldFor(field)
182 if (addressField is not None) and (data is None):
186 if format[:1] == '_':
190 if format[:1] == "'" or format[:1] == '"':
194 two = format.split('=')
197 return self.pack(two[0], data)
199 fields = {'self':self}
200 fields.update(self.fields)
201 return self.pack(two[0], eval(two[1], {}, fields))
204 two = format.split('&')
207 return self.pack(two[0], data)
209 if (self.fields.has_key(two[1])) and (self[two[1]] is not None):
210 return self.pack(two[0], id(self[two[1]]) & ((1<<(calcsize(two[0])*8))-1) )
212 return self.pack(two[0], 0)
215 two = format.split('-')
218 return self.pack(two[0],data)
220 return self.pack(two[0], self.calcPackFieldSize(two[1]))
223 two = format.split('*')
227 answer += self.pack(two[1], each)
230 if int(two[0]) != len(data):
231 raise Exception, "Array field has a constant size, and it doesn't match the actual value"
233 return self.pack(two[0], len(data))+answer
236 # "printf" string specifier
237 if format[:1] == '%':
238 # format string like specifier
242 if format[:1] == 'z':
243 return str(data)+'\0'
246 if format[:1] == 'u':
247 return str(data)+'\0\0' + (len(data) & 1 and '\0' or '')
249 # DCE-RPC/NDR string specifier
250 if format[:1] == 'w':
255 l = pack('<L', len(data)/2)
256 return '%s\0\0\0\0%s%s' % (l,l,data)
259 raise Exception, "Trying to pack None"
262 if format[:1] == ':':
265 # struct like specifier
266 return pack(format, data)
268 def unpack(self, format, data, dataClassOrCode = str, field = None):
270 print " unpack( %s | %r )" % (format, data)
273 addressField = self.findAddressFieldFor(field)
274 if addressField is not None:
275 if not self[addressField]:
279 if format[:1] == '_':
280 if dataClassOrCode != str:
281 fields = {'self':self, 'inputDataLeft':data}
282 fields.update(self.fields)
283 return eval(dataClassOrCode, {}, fields)
288 if format[:1] == "'" or format[:1] == '"':
291 raise Exception, "Unpacked data doesn't match constant value '%r' should be '%r'" % (data, answer)
295 two = format.split('&')
297 return self.unpack(two[0],data)
300 two = format.split('=')
302 return self.unpack(two[0],data)
305 two = format.split('-')
307 return self.unpack(two[0],data)
310 two = format.split('*')
317 sofar += self.calcUnpackSize(two[0], data)
318 number = self.unpack(two[0], data[:sofar])
322 while number and sofar < len(data):
323 nsofar = sofar + self.calcUnpackSize(two[1],data[sofar:])
324 answer.append(self.unpack(two[1], data[sofar:nsofar], dataClassOrCode))
329 # "printf" string specifier
330 if format[:1] == '%':
331 # format string like specifier
336 if data[-1] != '\x00':
337 raise Exception, ("%s 'z' field is not NUL terminated: %r" % (field, data))
338 return data[:-1] # remove trailing NUL
342 if data[-2:] != '\x00\x00':
343 raise Exception, ("%s 'u' field is not NUL-NUL terminated: %r" % (field, data))
344 return data[:-2] # remove trailing NUL
346 # DCE-RPC/NDR string specifier
348 l = unpack('<L', data[:4])[0]
349 return data[12:12+l*2]
353 return dataClassOrCode(data)
355 # struct like specifier
356 return unpack(format, data)[0]
358 def calcPackSize(self, format, data, field = None):
359 # # print " calcPackSize %s:%r" % (format, data)
361 addressField = self.findAddressFieldFor(field)
362 if addressField is not None:
363 if not self[addressField]:
367 if format[:1] == '_':
371 if format[:1] == "'" or format[:1] == '"':
375 two = format.split('&')
377 return self.calcPackSize(two[0], data)
380 two = format.split('=')
382 return self.calcPackSize(two[0], data)
385 two = format.split('-')
387 return self.calcPackSize(two[0], data)
390 two = format.split('*')
394 if int(two[0]) != len(data):
395 raise Exception, "Array field has a constant size, and it doesn't match the actual value"
397 answer += self.calcPackSize(two[0], len(data))
400 answer += self.calcPackSize(two[1], each)
403 # "printf" string specifier
404 if format[:1] == '%':
405 # format string like specifier
406 return len(format % data)
409 if format[:1] == 'z':
413 if format[:1] == 'u':
415 return l + (l & 1 and 3 or 2)
417 # DCE-RPC/NDR string specifier
418 if format[:1] == 'w':
423 if format[:1] == ':':
426 # struct like specifier
427 return calcsize(format)
429 def calcUnpackSize(self, format, data, field = None):
431 print " calcUnpackSize( %s | %s | %r)" % (field, format, data)
434 if format[:1] == '_':
437 addressField = self.findAddressFieldFor(field)
438 if addressField is not None:
439 if not self[addressField]:
443 lengthField = self.findLengthFieldFor(field)
444 return self[lengthField]
448 # XXX: Try to match to actual values, raise if no match
451 if format[:1] == "'" or format[:1] == '"':
455 two = format.split('&')
457 return self.calcUnpackSize(two[0], data)
460 two = format.split('=')
462 return self.calcUnpackSize(two[0], data)
465 two = format.split('-')
467 return self.calcUnpackSize(two[0], data)
470 two = format.split('*')
477 answer += self.calcUnpackSize(two[0], data)
478 number = self.unpack(two[0], data[:answer])
482 answer += self.calcUnpackSize(two[1], data[answer:])
484 while answer < len(data):
485 answer += self.calcUnpackSize(two[1], data[answer:])
488 # "printf" string specifier
489 if format[:1] == '%':
490 raise Exception, "Can't guess the size of a printf like specifier for unpacking"
493 if format[:1] == 'z':
494 return data.index('\x00')+1
497 if format[:1] == 'u':
498 l = data.index('\x00\x00')
499 return l + (l & 1 and 3 or 2)
501 # DCE-RPC/NDR string specifier
502 if format[:1] == 'w':
503 l = unpack('<L', data[:4])[0]
507 if format[:1] == ':':
510 # struct like specifier
511 return calcsize(format)
513 def calcPackFieldSize(self, fieldName, format = None):
515 format = self.formatForField(fieldName)
517 return self.calcPackSize(format, self[fieldName])
519 def formatForField(self, fieldName):
520 for field in self.commonHdr+self.structure:
521 if field[0] == fieldName:
523 raise Exception, ("Field %s not found" % fieldName)
525 def findAddressFieldFor(self, fieldName):
526 descriptor = '&%s' % fieldName
528 for field in self.commonHdr+self.structure:
529 if field[1][-l:] == descriptor:
533 def findLengthFieldFor(self, fieldName):
534 descriptor = '-%s' % fieldName
536 for field in self.commonHdr+self.structure:
537 if field[1][-l:] == descriptor:
541 def zeroValue(self, format):
542 two = format.split('*')
545 return (self.zeroValue(two[1]),)*int(two[0])
547 if not format.find('*') == -1: return ()
548 if 's' in format: return ''
549 if format in ['z',':','u']: return ''
550 if format == 'w': return '\x00\x00'
555 for field in self.commonHdr + self.structure:
556 self[field[0]] = self.zeroValue(field[1])
558 def dump(self, msg = None, indent = 0):
559 if msg is None: msg = self.__class__.__name__
563 for field in self.commonHdr+self.structure:
566 fixedFields.append(i)
567 if isinstance(self[i], Structure):
568 self[i].dump('%s%s:{' % (ind,i), indent = indent + 4)
571 print "%s%s: {%r}" % (ind,i,self[i])
572 # Do we have remaining fields not defined in the structures? let's
574 remainingFields = list(set(self.fields) - set(fixedFields))
575 for i in remainingFields:
576 if isinstance(self[i], Structure):
577 self[i].dump('%s%s:{' % (ind,i), indent = indent + 4)
580 print "%s%s: {%r}" % (ind,i,self[i])
583 class _StructureTest:
585 def create(self,data = None):
587 return self.theClass(data, alignment = self.alignment)
589 return self.theClass(alignment = self.alignment)
594 testName = self.__class__.__name__
595 print "starting test: %s....." % testName
598 a.dump("packing.....")
600 print "packed: %r" % a_str
601 print "unpacking....."
602 b = self.create(a_str)
603 b.dump("unpacked.....")
604 print "repacking....."
607 print "ERROR: original packed and repacked don't match"
608 print "packed: %r" % b_str
610 class _Test_simple(_StructureTest):
611 class theClass(Structure):
624 ('code1','>L=len(arr1)*2+0x1000'),
627 def populate(self, a):
628 a['default'] = 'hola'
630 a['int3'] = 0x45444342
632 a['u1'] = 'hola'.encode('utf_16_le')
634 a['arr1'] = (0x12341234,0x88990077,0x41414141)
635 # a['len1'] = 0x42424242
637 class _Test_fixedLength(_Test_simple):
638 def populate(self, a):
639 _Test_simple.populate(self, a)
640 a['len1'] = 0x42424242
642 class _Test_simple_aligned4(_Test_simple):
645 class _Test_nested(_StructureTest):
646 class theClass(Structure):
647 class _Inner(Structure):
648 structure = (('data', 'z'),)
651 ('nest1', ':', _Inner),
652 ('nest2', ':', _Inner),
656 def populate(self, a):
657 a['nest1'] = _Test_nested.theClass._Inner()
658 a['nest2'] = _Test_nested.theClass._Inner()
659 a['nest1']['data'] = 'hola manola'
660 a['nest2']['data'] = 'chau loco'
661 a['int'] = 0x12345678
663 class _Test_Optional(_StructureTest):
664 class theClass(Structure):
672 def populate(self, a):
673 a['Name'] = 'Optional test'
674 a['List'] = (1,2,3,4)
676 class _Test_Optional_sparse(_Test_Optional):
677 def populate(self, a):
678 _Test_Optional.populate(self, a)
681 class _Test_AsciiZArray(_StructureTest):
682 class theClass(Structure):
689 def populate(self, a):
692 a['array'] = ('hola','manola','te traje')
694 class _Test_UnpackCode(_StructureTest):
695 class theClass(Structure):
697 ('leni','<L=len(uno)*2'),
698 ('cuchi','_-uno','leni/2'),
703 def populate(self, a):
704 a['uno'] = 'soy un loco!'
705 a['dos'] = 'que haces fiera'
707 class _Test_AAA(_StructureTest):
708 class theClass(Structure):
711 ('iv', '!L=((init_vector & 0xFFFFFF) << 8) | ((pad & 0x3f) << 2) | (keyid & 3)'),
712 ('init_vector', '_','(iv >> 8)'),
713 ('pad', '_','((iv >>2) & 0x3F)'),
714 ('keyid', '_','( iv & 0x03 )'),
715 ('dataLen', '_-data', 'len(inputDataLeft)-4'),
720 def populate(self, a):
721 a['init_vector']=0x01020304
722 #a['pad']=int('01010101',2)
723 a['pad']=int('010101',2)
725 a['data']="\xA0\xA1\xA2\xA3\xA4\xA5\xA6\xA7\xA8\xA9"
726 a['icv'] = 0x05060708
727 #a['iv'] = 0x01020304
729 if __name__ == '__main__':
733 _Test_fixedLength().run()
735 print "cannot repack because length is bogus"
737 _Test_simple_aligned4().run()
739 _Test_Optional().run()
740 _Test_Optional_sparse().run()
741 _Test_AsciiZArray().run()
742 _Test_UnpackCode().run()