Imported Upstream version 7.59.0
[platform/upstream/curl.git] / tests / python_dependencies / impacket / structure.py
1 # Copyright (c) 2003-2016 CORE Security Technologies
2 #
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.
6 #
7
8 from struct import pack, unpack, calcsize
9
10 class Structure:
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]
14         
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.
17
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)
20
21         format specifiers:
22           specifiers from module pack can be used with the same format 
23           see struct.__doc__ (pack/unpack is finally called)
24             x       [padding byte]
25             c       [character]
26             b       [signed byte]
27             B       [unsigned byte]
28             h       [signed short]
29             H       [unsigned short]
30             l       [signed long]
31             L       [unsigned long]
32             i       [signed integer]
33             I       [unsigned integer]
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]
38             f       [float]
39             d       [double]
40             =       [native byte ordering, size and alignment]
41             @       [native byte ordering, standard size and alignment]
42             !       [network byte ordering]
43             <       [little endian]
44             >       [big endian]
45
46           usual printf like specifiers can be used (if started with %) 
47           [not recommeneded, there is no why to unpack this]
48
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)
53             ...
54
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.
69             
70     """
71     commonHdr = ()
72     structure = ()
73     debug = 0
74
75     def __init__(self, data = None, alignment = 0):
76         if not hasattr(self, 'alignment'):
77             self.alignment = alignment
78
79         self.fields    = {}
80         self.rawData   = data
81         if data is not None:
82             self.fromString(data)
83         else:
84             self.data = None
85
86     @classmethod
87     def fromFile(self, file):
88         answer = self()
89         answer.fromString(file.read(len(answer)))
90         return answer
91
92     def setAlignment(self, alignment):
93         self.alignment = alignment
94
95     def setData(self, data):
96         self.data = data
97
98     def packField(self, fieldName, format = None):
99         if self.debug:
100             print "packField( %s | %s )" % (fieldName, format)
101
102         if format is None:
103             format = self.formatForField(fieldName)
104
105         if self.fields.has_key(fieldName):
106             ans = self.pack(format, self.fields[fieldName], field = fieldName)
107         else:
108             ans = self.pack(format, None, field = fieldName)
109
110         if self.debug:
111             print "\tanswer %r" % ans
112
113         return ans
114
115     def getData(self):
116         if self.data is not None:
117             return self.data
118         data = ''
119         for field in self.commonHdr+self.structure:
120             try:
121                 data += self.packField(field[0], field[1])
122             except Exception, e:
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__),)
125                 else:
126                     e.args += ("When packing field '%s | %s' in %s" % (field[0], field[1], self.__class__),)
127                 raise
128             if self.alignment:
129                 if len(data) % self.alignment:
130                     data += ('\x00'*self.alignment)[:-(len(data) % self.alignment)]
131             
132         #if len(data) % self.alignment: data += ('\x00'*self.alignment)[:-(len(data) % self.alignment)]
133         return data
134
135     def fromString(self, data):
136         self.rawData = data
137         for field in self.commonHdr+self.structure:
138             if self.debug:
139                 print "fromString( %s | %s | %r )" % (field[0], field[1], data)
140             size = self.calcUnpackSize(field[1], data, field[0])
141             if self.debug:
142                 print "  size = %d" % size
143             dataClassOrCode = str
144             if len(field) > 2:
145                 dataClassOrCode = field[2]
146             try:
147                 self[field[0]] = self.unpack(field[1], data[:size], dataClassOrCode = dataClassOrCode, field = field[0])
148             except Exception,e:
149                 e.args += ("When unpacking field '%s | %s | %r[:%d]'" % (field[0], field[1], data, size),)
150                 raise
151
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)
155             data = data[size:]
156
157         return self
158         
159     def __setitem__(self, key, value):
160         self.fields[key] = value
161         self.data = None        # force recompute
162
163     def __getitem__(self, key):
164         return self.fields[key]
165
166     def __delitem__(self, key):
167         del self.fields[key]
168         
169     def __str__(self):
170         return self.getData()
171
172     def __len__(self):
173         # XXX: improve
174         return len(self.getData())
175
176     def pack(self, format, data, field = None):
177         if self.debug:
178             print "  pack( %s | %r | %s)" %  (format, data, field)
179
180         if field:
181             addressField = self.findAddressFieldFor(field)
182             if (addressField is not None) and (data is None):
183                 return ''
184
185         # void specifier
186         if format[:1] == '_':
187             return ''
188
189         # quote specifier
190         if format[:1] == "'" or format[:1] == '"':
191             return format[1:]
192
193         # code specifier
194         two = format.split('=')
195         if len(two) >= 2:
196             try:
197                 return self.pack(two[0], data)
198             except:
199                 fields = {'self':self}
200                 fields.update(self.fields)
201                 return self.pack(two[0], eval(two[1], {}, fields))
202
203         # address specifier
204         two = format.split('&')
205         if len(two) == 2:
206             try:
207                 return self.pack(two[0], data)
208             except:
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) )
211                 else:
212                     return self.pack(two[0], 0)
213
214         # length specifier
215         two = format.split('-')
216         if len(two) == 2:
217             try:
218                 return self.pack(two[0],data)
219             except:
220                 return self.pack(two[0], self.calcPackFieldSize(two[1]))
221
222         # array specifier
223         two = format.split('*')
224         if len(two) == 2:
225             answer = ''
226             for each in data:
227                 answer += self.pack(two[1], each)
228             if two[0]:
229                 if two[0].isdigit():
230                     if int(two[0]) != len(data):
231                         raise Exception, "Array field has a constant size, and it doesn't match the actual value"
232                 else:
233                     return self.pack(two[0], len(data))+answer
234             return answer
235
236         # "printf" string specifier
237         if format[:1] == '%':
238             # format string like specifier
239             return format % data
240
241         # asciiz specifier
242         if format[:1] == 'z':
243             return str(data)+'\0'
244
245         # unicode specifier
246         if format[:1] == 'u':
247             return str(data)+'\0\0' + (len(data) & 1 and '\0' or '')
248
249         # DCE-RPC/NDR string specifier
250         if format[:1] == 'w':
251             if len(data) == 0:
252                 data = '\0\0'
253             elif len(data) % 2:
254                 data += '\0'
255             l = pack('<L', len(data)/2)
256             return '%s\0\0\0\0%s%s' % (l,l,data)
257                     
258         if data is None:
259             raise Exception, "Trying to pack None"
260         
261         # literal specifier
262         if format[:1] == ':':
263             return str(data)
264
265         # struct like specifier
266         return pack(format, data)
267
268     def unpack(self, format, data, dataClassOrCode = str, field = None):
269         if self.debug:
270             print "  unpack( %s | %r )" %  (format, data)
271
272         if field:
273             addressField = self.findAddressFieldFor(field)
274             if addressField is not None:
275                 if not self[addressField]:
276                     return
277
278         # void specifier
279         if format[:1] == '_':
280             if dataClassOrCode != str:
281                 fields = {'self':self, 'inputDataLeft':data}
282                 fields.update(self.fields)
283                 return eval(dataClassOrCode, {}, fields)
284             else:
285                 return None
286
287         # quote specifier
288         if format[:1] == "'" or format[:1] == '"':
289             answer = format[1:]
290             if answer != data:
291                 raise Exception, "Unpacked data doesn't match constant value '%r' should be '%r'" % (data, answer)
292             return answer
293
294         # address specifier
295         two = format.split('&')
296         if len(two) == 2:
297             return self.unpack(two[0],data)
298
299         # code specifier
300         two = format.split('=')
301         if len(two) >= 2:
302             return self.unpack(two[0],data)
303
304         # length specifier
305         two = format.split('-')
306         if len(two) == 2:
307             return self.unpack(two[0],data)
308
309         # array specifier
310         two = format.split('*')
311         if len(two) == 2:
312             answer = []
313             sofar = 0
314             if two[0].isdigit():
315                 number = int(two[0])
316             elif two[0]:
317                 sofar += self.calcUnpackSize(two[0], data)
318                 number = self.unpack(two[0], data[:sofar])
319             else:
320                 number = -1
321
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))
325                 number -= 1
326                 sofar = nsofar
327             return answer
328
329         # "printf" string specifier
330         if format[:1] == '%':
331             # format string like specifier
332             return format % data
333
334         # asciiz specifier
335         if format == 'z':
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
339
340         # unicode specifier
341         if format == 'u':
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
345
346         # DCE-RPC/NDR string specifier
347         if format == 'w':
348             l = unpack('<L', data[:4])[0]
349             return data[12:12+l*2]
350
351         # literal specifier
352         if format == ':':
353             return dataClassOrCode(data)
354
355         # struct like specifier
356         return unpack(format, data)[0]
357
358     def calcPackSize(self, format, data, field = None):
359 #        # print "  calcPackSize  %s:%r" %  (format, data)
360         if field:
361             addressField = self.findAddressFieldFor(field)
362             if addressField is not None:
363                 if not self[addressField]:
364                     return 0
365
366         # void specifier
367         if format[:1] == '_':
368             return 0
369
370         # quote specifier
371         if format[:1] == "'" or format[:1] == '"':
372             return len(format)-1
373
374         # address specifier
375         two = format.split('&')
376         if len(two) == 2:
377             return self.calcPackSize(two[0], data)
378
379         # code specifier
380         two = format.split('=')
381         if len(two) >= 2:
382             return self.calcPackSize(two[0], data)
383
384         # length specifier
385         two = format.split('-')
386         if len(two) == 2:
387             return self.calcPackSize(two[0], data)
388
389         # array specifier
390         two = format.split('*')
391         if len(two) == 2:
392             answer = 0
393             if two[0].isdigit():
394                     if int(two[0]) != len(data):
395                         raise Exception, "Array field has a constant size, and it doesn't match the actual value"
396             elif two[0]:
397                 answer += self.calcPackSize(two[0], len(data))
398
399             for each in data:
400                 answer += self.calcPackSize(two[1], each)
401             return answer
402
403         # "printf" string specifier
404         if format[:1] == '%':
405             # format string like specifier
406             return len(format % data)
407
408         # asciiz specifier
409         if format[:1] == 'z':
410             return len(data)+1
411
412         # asciiz specifier
413         if format[:1] == 'u':
414             l = len(data)
415             return l + (l & 1 and 3 or 2)
416
417         # DCE-RPC/NDR string specifier
418         if format[:1] == 'w':
419             l = len(data)
420             return 12+l+l % 2
421
422         # literal specifier
423         if format[:1] == ':':
424             return len(data)
425
426         # struct like specifier
427         return calcsize(format)
428
429     def calcUnpackSize(self, format, data, field = None):
430         if self.debug:
431             print "  calcUnpackSize( %s | %s | %r)" %  (field, format, data)
432
433         # void specifier
434         if format[:1] == '_':
435             return 0
436
437         addressField = self.findAddressFieldFor(field)
438         if addressField is not None:
439             if not self[addressField]:
440                 return 0
441
442         try:
443             lengthField = self.findLengthFieldFor(field)
444             return self[lengthField]
445         except:
446             pass
447
448         # XXX: Try to match to actual values, raise if no match
449         
450         # quote specifier
451         if format[:1] == "'" or format[:1] == '"':
452             return len(format)-1
453
454         # address specifier
455         two = format.split('&')
456         if len(two) == 2:
457             return self.calcUnpackSize(two[0], data)
458
459         # code specifier
460         two = format.split('=')
461         if len(two) >= 2:
462             return self.calcUnpackSize(two[0], data)
463
464         # length specifier
465         two = format.split('-')
466         if len(two) == 2:
467             return self.calcUnpackSize(two[0], data)
468
469         # array specifier
470         two = format.split('*')
471         if len(two) == 2:
472             answer = 0
473             if two[0]:
474                 if two[0].isdigit():
475                     number = int(two[0])
476                 else:
477                     answer += self.calcUnpackSize(two[0], data)
478                     number = self.unpack(two[0], data[:answer])
479
480                 while number:
481                     number -= 1
482                     answer += self.calcUnpackSize(two[1], data[answer:])
483             else:
484                 while answer < len(data):
485                     answer += self.calcUnpackSize(two[1], data[answer:])
486             return answer
487
488         # "printf" string specifier
489         if format[:1] == '%':
490             raise Exception, "Can't guess the size of a printf like specifier for unpacking"
491
492         # asciiz specifier
493         if format[:1] == 'z':
494             return data.index('\x00')+1
495
496         # asciiz specifier
497         if format[:1] == 'u':
498             l = data.index('\x00\x00')
499             return l + (l & 1 and 3 or 2)
500
501         # DCE-RPC/NDR string specifier
502         if format[:1] == 'w':
503             l = unpack('<L', data[:4])[0]
504             return 12+l*2
505
506         # literal specifier
507         if format[:1] == ':':
508             return len(data)
509
510         # struct like specifier
511         return calcsize(format)
512
513     def calcPackFieldSize(self, fieldName, format = None):
514         if format is None:
515             format = self.formatForField(fieldName)
516
517         return self.calcPackSize(format, self[fieldName])
518
519     def formatForField(self, fieldName):
520         for field in self.commonHdr+self.structure:
521             if field[0] == fieldName:
522                 return field[1]
523         raise Exception, ("Field %s not found" % fieldName)
524
525     def findAddressFieldFor(self, fieldName):
526         descriptor = '&%s' % fieldName
527         l = len(descriptor)
528         for field in self.commonHdr+self.structure:
529             if field[1][-l:] == descriptor:
530                 return field[0]
531         return None
532         
533     def findLengthFieldFor(self, fieldName):
534         descriptor = '-%s' % fieldName
535         l = len(descriptor)
536         for field in self.commonHdr+self.structure:
537             if field[1][-l:] == descriptor:
538                 return field[0]
539         return None
540         
541     def zeroValue(self, format):
542         two = format.split('*')
543         if len(two) == 2:
544             if two[0].isdigit():
545                 return (self.zeroValue(two[1]),)*int(two[0])
546                         
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'
551
552         return 0
553
554     def clear(self):
555         for field in self.commonHdr + self.structure:
556             self[field[0]] = self.zeroValue(field[1])
557
558     def dump(self, msg = None, indent = 0):
559         if msg is None: msg = self.__class__.__name__
560         ind = ' '*indent
561         print "\n%s" % msg
562         fixedFields = []
563         for field in self.commonHdr+self.structure:
564             i = field[0] 
565             if i in self.fields:
566                 fixedFields.append(i)
567                 if isinstance(self[i], Structure):
568                     self[i].dump('%s%s:{' % (ind,i), indent = indent + 4)
569                     print "%s}" % ind
570                 else:
571                     print "%s%s: {%r}" % (ind,i,self[i])
572         # Do we have remaining fields not defined in the structures? let's 
573         # print them
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)
578                 print "%s}" % ind
579             else:
580                 print "%s%s: {%r}" % (ind,i,self[i])
581
582
583 class _StructureTest:
584     alignment = 0
585     def create(self,data = None):
586         if data is not None:
587             return self.theClass(data, alignment = self.alignment)
588         else:
589             return self.theClass(alignment = self.alignment)
590
591     def run(self):
592         print
593         print "-"*70
594         testName = self.__class__.__name__
595         print "starting test: %s....." % testName
596         a = self.create()
597         self.populate(a)
598         a.dump("packing.....")
599         a_str = str(a)
600         print "packed: %r" % a_str
601         print "unpacking....."
602         b = self.create(a_str)
603         b.dump("unpacked.....")
604         print "repacking....."
605         b_str = str(b)
606         if b_str != a_str:
607             print "ERROR: original packed and repacked don't match"
608             print "packed: %r" % b_str
609
610 class _Test_simple(_StructureTest):
611     class theClass(Structure):
612         commonHdr = ()
613         structure = (
614                 ('int1', '!L'),
615                 ('len1','!L-z1'),
616                 ('arr1','B*<L'),
617                 ('z1', 'z'),
618                 ('u1','u'),
619                 ('', '"COCA'),
620                 ('len2','!H-:1'),
621                 ('', '"COCA'),
622                 (':1', ':'),
623                 ('int3','>L'),
624                 ('code1','>L=len(arr1)*2+0x1000'),
625                 )
626
627     def populate(self, a):
628         a['default'] = 'hola'
629         a['int1'] = 0x3131
630         a['int3'] = 0x45444342
631         a['z1']   = 'hola'
632         a['u1']   = 'hola'.encode('utf_16_le')
633         a[':1']   = ':1234:'
634         a['arr1'] = (0x12341234,0x88990077,0x41414141)
635         # a['len1'] = 0x42424242
636
637 class _Test_fixedLength(_Test_simple):
638     def populate(self, a):
639         _Test_simple.populate(self, a)
640         a['len1'] = 0x42424242
641
642 class _Test_simple_aligned4(_Test_simple):
643     alignment = 4
644
645 class _Test_nested(_StructureTest):
646     class theClass(Structure):
647         class _Inner(Structure):
648             structure = (('data', 'z'),)
649
650         structure = (
651             ('nest1', ':', _Inner),
652             ('nest2', ':', _Inner),
653             ('int', '<L'),
654         )
655
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
662     
663 class _Test_Optional(_StructureTest):
664     class theClass(Structure):
665         structure = (
666                 ('pName','<L&Name'),
667                 ('pList','<L&List'),
668                 ('Name','w'),
669                 ('List','<H*<L'),
670             )
671             
672     def populate(self, a):
673         a['Name'] = 'Optional test'
674         a['List'] = (1,2,3,4)
675         
676 class _Test_Optional_sparse(_Test_Optional):
677     def populate(self, a):
678         _Test_Optional.populate(self, a)
679         del a['Name']
680
681 class _Test_AsciiZArray(_StructureTest):
682     class theClass(Structure):
683         structure = (
684             ('head','<L'),
685             ('array','B*z'),
686             ('tail','<L'),
687         )
688
689     def populate(self, a):
690         a['head'] = 0x1234
691         a['tail'] = 0xabcd
692         a['array'] = ('hola','manola','te traje')
693         
694 class _Test_UnpackCode(_StructureTest):
695     class theClass(Structure):
696         structure = (
697             ('leni','<L=len(uno)*2'),
698             ('cuchi','_-uno','leni/2'),
699             ('uno',':'),
700             ('dos',':'),
701         )
702
703     def populate(self, a):
704         a['uno'] = 'soy un loco!'
705         a['dos'] = 'que haces fiera'
706
707 class _Test_AAA(_StructureTest):
708     class theClass(Structure):
709         commonHdr = ()
710         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'),
716           ('data',':'),
717           ('icv','>L'),
718         )
719
720     def populate(self, a):
721         a['init_vector']=0x01020304
722         #a['pad']=int('01010101',2)
723         a['pad']=int('010101',2)
724         a['keyid']=0x07
725         a['data']="\xA0\xA1\xA2\xA3\xA4\xA5\xA6\xA7\xA8\xA9"
726         a['icv'] = 0x05060708
727         #a['iv'] = 0x01020304
728         
729 if __name__ == '__main__':
730     _Test_simple().run()
731
732     try:
733         _Test_fixedLength().run()
734     except:
735         print "cannot repack because length is bogus"
736
737     _Test_simple_aligned4().run()
738     _Test_nested().run()
739     _Test_Optional().run()
740     _Test_Optional_sparse().run()
741     _Test_AsciiZArray().run()
742     _Test_UnpackCode().run()
743     _Test_AAA().run()