Peter M. Groen / oletools

Commit 470d0806adfd231d2cd1736d71f79d659499a970

Authored by Christian Herdtweck
1 parent 5c9b328c

record_base: make compatible with container substreams

Inline Side-by-side
Showing 1 changed file with 24 additions and 12 deletions
oletools/record_base.py
  View file @470d080
... ... @@ -73,6 +73,7 @@ from oletools.thirdparty import olefile
73 73 # Helpers
74 74 ###############################################################################
75 75  
  76 +STGTY_SUBSTREAM = 10
76 77  
77 78 ENTRY_TYPE2STR = {
78 79 olefile.STGTY_EMPTY: 'empty',
... ... @@ -80,7 +81,8 @@ ENTRY_TYPE2STR = {
80 81 olefile.STGTY_STREAM: 'stream',
81 82 olefile.STGTY_LOCKBYTES: 'lock-bytes',
82 83 olefile.STGTY_PROPERTY: 'property',
83   - olefile.STGTY_ROOT: 'root'
  84 + olefile.STGTY_ROOT: 'root',
  85 + STGTY_SUBSTREAM: 'substream'
84 86 }
85 87  
86 88  
... ... @@ -131,7 +133,9 @@ class OleRecordFile(olefile.OleFileIO):
131 133 else:
132 134 clz = self.stream_class_for_name(direntry.name)
133 135 yield clz(self._open(direntry.isectStart, direntry.size),
134   - None if is_orphan else direntry.name)
  136 + direntry.size,
  137 + None if is_orphan else direntry.name,
  138 + direntry.entry_type)
135 139  
136 140  
137 141 class OleRecordStream(object):
... ... @@ -142,10 +146,13 @@ class OleRecordStream(object):
142 146 abstract base class
143 147 """
144 148  
145   - def __init__(self, stream, name):
  149 + def __init__(self, stream, size, name, stream_type):
146 150 self.stream = stream
  151 + self.size = size
147 152 self.name = name
148   - self.size = stream.size
  153 + if stream_type not in ENTRY_TYPE2STR:
  154 + raise ValueError('Unknown stream type: {0}'.format(stream_type))
  155 + self.stream_type = stream_type
149 156  
150 157 def read_record_head(self):
151 158 """ read first few bytes of record to determine size and type
... ... @@ -191,7 +198,7 @@ class OleRecordStream(object):
191 198 if fill_data or force_read:
192 199 data = self.stream.read(rec_size)
193 200 if len(data) != rec_size:
194   - raise IOError('Not enough data in stream ({0} < {1})'
  201 + raise IOError('Unexpected end of stream ({0} < {1})'
195 202 .format(len(data), rec_size))
196 203 else:
197 204 self.stream.seek(rec_size, SEEK_CUR)
... ... @@ -199,9 +206,11 @@ class OleRecordStream(object):
199 206 yield rec_clz(rec_type, rec_size, other, pos, data)
200 207  
201 208 def __str__(self):
202   - return '[{2} {0} (size {1})' \
203   - .format(self.name or '[orphan]', self.size,
204   - self.__class__.__name__)
  209 + return '[{0} {1} (type {2}, size {3})' \
  210 + .format(self.__class__.__name__,
  211 + self.name or '[orphan]',
  212 + ENTRY_TYPE2STR[self.stream_type],
  213 + self.size)
205 214  
206 215  
207 216 class OleSummaryInformationStream(OleRecordStream):
... ... @@ -272,13 +281,15 @@ class OleRecordBase(object):
272 281  
273 282  
274 283 def test(filenames, ole_file_class=OleRecordFile,
275   - must_parse=None):
  284 + must_parse=None, do_per_record=None):
276 285 """ parse all given file names and print rough structure
277 286  
278 287 if an error occurs while parsing a stream of type in must_parse, the error
279 288 will be raised. Otherwise a message is printed
280 289 """
281   - logging.basicConfig(level=logging.DEBUG)
  290 + logging.basicConfig(level=logging.INFO)
  291 + if do_per_record is None:
  292 + do_per_record = lambda record: None
282 293 if not filenames:
283 294 logging.info('need file name[s]')
284 295 return 2
... ... @@ -290,10 +301,11 @@ def test(filenames, ole_file_class=OleRecordFile,
290 301 ole = ole_file_class(filename)
291 302  
292 303 for stream in ole.iter_streams():
293   - logging.info(stream)
  304 + logging.info(' parse ' + str(stream))
294 305 try:
295 306 for record in stream.iter_records():
296   - logging.info(' {0}'.format(record))
  307 + logging.info(' found ' + str(record))
  308 + do_per_record(record)
297 309 except Exception:
298 310 if not must_parse:
299 311 raise
... ...