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)