github.com/n00py/Slackor@v0.0.0-20200610224921-d007fcea1740/impacket/tests/misc/test_structure.py (about) 1 from __future__ import print_function 2 import unittest 3 from binascii import hexlify, unhexlify 4 5 from impacket.structure import Structure 6 7 def hexl(b): 8 hexstr = str(hexlify(b).decode('ascii')) 9 return ' '.join([hexstr[i:i+8] for i in range(0,len(hexstr),8)]) 10 11 class _StructureTest(unittest.TestCase): 12 # Subclass: 13 # - must define theClass 14 # - may override alignment 15 alignment = 0 16 17 def create(self,data = None): 18 if data is not None: 19 return self.theClass(data, alignment = self.alignment) 20 else: 21 return self.theClass(alignment = self.alignment) 22 23 def test_structure(self): 24 #print() 25 #print("-"*70) 26 #testName = self.__class__.__name__ 27 #print("starting test: %s....." % testName) 28 # Create blank structure and fill its fields 29 a = self.create() 30 self.populate(a) 31 #a.dump("packing.....") 32 # Get its binary representation 33 a_str = a.getData() 34 self.check_data(a_str) 35 #print("packed: %r" % a_str) 36 #print("unpacking.....") 37 b = self.create(a_str) 38 #b.dump("unpacked.....") 39 #print("repacking.....") 40 b_str = b.getData() 41 self.assertEqual(b_str, a_str, 42 "ERROR: original packed and repacked don't match") 43 44 def check_data(self, a_str): 45 if hasattr(self, 'hexData'): 46 # Regression check 47 self.assertEqual(hexl(a_str), self.hexData) 48 else: 49 # Show result, to aid adding regression check 50 print(self.__class__.__name__, hexl(a_str)) 51 52 if not hasattr(unittest.TestCase,'assertRaisesRegex'): 53 if hasattr(unittest.TestCase,'assertRaisesRegexp'): # PY2.7, PY3.1 54 assertRaisesRegex = unittest.TestCase.assertRaisesRegexp 55 else: # PY2.6 56 def assertRaisesRegex(self,ex,rx,*args): 57 # Just ignore the regex 58 return self.assertRaises(ex,*args) 59 60 class Test_simple(_StructureTest): 61 class theClass(Structure): 62 commonHdr = () 63 structure = ( 64 ('int1', '!L'), 65 ('len1','!L-z1'), 66 ('arr1','B*<L'), 67 ('z1', 'z'), 68 ('u1','u'), 69 ('', '"COCA'), 70 ('len2','!H-:1'), 71 ('', '"COCA'), 72 (':1', ':'), 73 ('int3','>L'), 74 ('code1','>L=len(arr1)*2+0x1000'), 75 ) 76 77 def populate(self, a): 78 a['default'] = 'hola' 79 a['int1'] = 0x3131 80 a['int3'] = 0x45444342 81 a['z1'] = 'hola' 82 a['u1'] = 'hola'.encode('utf_16_le') 83 a[':1'] = ':1234:' 84 a['arr1'] = (0x12341234,0x88990077,0x41414141) 85 # a['len1'] = 0x42424242 86 87 hexData = '00003131 00000005 03341234 12770099 88414141 41686f6c 61006800 6f006c00 61000000 434f4341 0006434f 43413a31 3233343a 45444342 00001006' 88 89 class Test_fixedLength(Test_simple): 90 def test_structure(self): 91 a = self.create() 92 self.populate(a) 93 # Set a bogus length... 94 a['len1'] = 0x42424242 95 a_str = a.getData() 96 if hasattr(self, 'hexData'): 97 # Regression check 98 self.assertEqual(hexl(a_str), self.hexData) 99 else: 100 print(hexl(a_str)) 101 # ... so that unpacking will now fail 102 with self.assertRaisesRegex(Exception, r'not NUL terminated'): 103 self.create(a_str) 104 105 hexData = '00003131 42424242 03341234 12770099 88414141 41686f6c 61006800 6f006c00 61000000 434f4341 0006434f 43413a31 3233343a 45444342 00001006' 106 107 class Test_simple_aligned4(Test_simple): 108 alignment = 4 109 hexData = '00003131 00000005 03341234 12770099 88414141 41000000 686f6c61 00000000 68006f00 6c006100 00000000 434f4341 00060000 434f4341 3a313233 343a0000 45444342 00001006' 110 111 class Test_nested(_StructureTest): 112 class theClass(Structure): 113 class _Inner(Structure): 114 structure = (('data', 'z'),) 115 116 structure = ( 117 ('nest1', ':', _Inner), 118 ('nest2', ':', _Inner), 119 ('int', '<L'), 120 ) 121 122 def populate(self, a): 123 a['nest1'] = Test_nested.theClass._Inner() 124 a['nest2'] = Test_nested.theClass._Inner() 125 a['nest1']['data'] = 'hola manola' 126 a['nest2']['data'] = 'chau loco' 127 a['int'] = 0x12345678 128 129 hexData = '686f6c61 206d616e 6f6c6100 63686175 206c6f63 6f007856 3412' 130 131 class Test_Optional(_StructureTest): 132 class theClass(Structure): 133 structure = ( 134 ('pName','<L&Name'), 135 ('pList','<L&List'), 136 ('Name','w'), 137 ('List','<H*<L'), 138 ) 139 140 def populate(self, a): 141 a['Name'] = 'Optional test' 142 a['List'] = (1,2,3,4) 143 144 def check_data(self, a_str): 145 # Pointer values change between runs, so ignore them 146 filtered = ''.join(['-' if h=='-' else a 147 for a, h in zip(hexl(a_str), self.hexData)]) 148 self.assertEqual(filtered, self.hexData) 149 150 hexData = '-------- -------- 07000000 07000000 00000000 4f707469 6f6e616c 20746573 74000400 01000000 02000000 03000000 04000000' 151 152 class Test_Optional_sparse(Test_Optional): 153 def populate(self, a): 154 Test_Optional.populate(self, a) 155 del a['Name'] 156 157 hexData = '00000000 -------- 04000100 00000200 00000300 00000400 0000' 158 159 class Test_AsciiZArray(_StructureTest): 160 class theClass(Structure): 161 structure = ( 162 ('head','<L'), 163 ('array','B*z'), 164 ('tail','<L'), 165 ) 166 167 def populate(self, a): 168 a['head'] = 0x1234 169 a['tail'] = 0xabcd 170 a['array'] = ('hola','manola','te traje') 171 172 hexData = '34120000 03686f6c 61006d61 6e6f6c61 00746520 7472616a 6500cdab 0000' 173 class Test_UnpackCode(_StructureTest): 174 class theClass(Structure): 175 structure = ( 176 ('leni','<L=len(uno)*2'), 177 ('cuchi','_-uno','leni//2'), 178 ('uno',':'), 179 ('dos',':'), 180 ) 181 182 def populate(self, a): 183 a['uno'] = 'soy un loco!' 184 a['dos'] = 'que haces fiera' 185 186 hexData = '18000000 736f7920 756e206c 6f636f21 71756520 68616365 73206669 657261' 187 188 class Test_AAA(_StructureTest): 189 class theClass(Structure): 190 commonHdr = () 191 structure = ( 192 ('iv', '!L=((init_vector & 0xFFFFFF) << 8) | ((pad & 0x3f) << 2) | (keyid & 3)'), 193 ('init_vector', '_','(iv >> 8)'), 194 ('pad', '_','((iv >>2) & 0x3F)'), 195 ('keyid', '_','( iv & 0x03 )'), 196 ('dataLen', '_-data', 'len(inputDataLeft)-4'), 197 ('data',':'), 198 ('icv','>L'), 199 ) 200 201 def populate(self, a): 202 a['init_vector']=0x01020304 203 #a['pad']=int('01010101',2) 204 a['pad']=int('010101',2) 205 a['keyid']=0x07 206 a['data']="\xA0\xA1\xA2\xA3\xA4\xA5\xA6\xA7\xA8\xA9" 207 a['icv'] = 0x05060708 208 #a['iv'] = 0x01020304 209 210 hexData = '02030457 a0a1a2a3 a4a5a6a7 a8a90506 0708' 211 212 if __name__ == "__main__": 213 # Hide base class so that unittest.main() will not try to load it 214 del _StructureTest 215 unittest.main(verbosity=1)