From 82bb3107927ea4f4356789843698cba652713dd6 Mon Sep 17 00:00:00 2001
From: Vincent Brillault <git@lerya.net>
Date: Tue, 31 Jan 2017 19:41:05 +0100
Subject: [PATCH] WIP: Separate object (ExtendedStream) from logic (functions)

---
 oletools/oleform.py | 354 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 oletools/olevba.py  |   4 ++--
 2 files changed, 181 insertions(+), 177 deletions(-)

diff --git a/oletools/oleform.py b/oletools/oleform.py
index 3ad821f..0b9be2e 100644
--- a/oletools/oleform.py
+++ b/oletools/oleform.py
@@ -49,33 +49,43 @@ class MorphDataPropMask(Mask):
               'fBorderColor', 'fSpecialEffect', 'fMouseIcon', 'fPicture', 'fAccelerator',
               'UnusedBits2', 'Reserved', 'fGroupName']
 
-class OleUserFormParser(object):
-    def __init__(self, control_stream, data_stream):
-        self.variables = []
-        self.set_stream(control_stream)
-        self.consume_FormControl()
-        self.set_stream(data_stream)
-        self.consume_stored_data()
-
-    def set_stream(self, stream):
+class ExtendedStream(object):
+    def __init__(self, stream, path):
         self._pos = 0
-        self._frozen_pos = []
+        self._jumps = []
         self._stream = stream
+        self._path = path
+
+    @classmethod
+    def open(cls, ole_file, path):
+        stream = ole_file.openstream(path)
+        return cls(stream, path)
 
     def read(self, size):
         self._pos += size
         return self._stream.read(size)
 
-    def freeze(self):
-        self._frozen_pos.append(self._pos)
+    def will_jump_to(self, size):
+        self._next_jump = (True, size)
+        return self
+
+    def will_pad(self, pad=4):
+        self._next_jump = (False, pad)
+        return self
 
-    def unfreeze(self, size):
-        self.read(self._frozen_pos.pop() - self._pos + size)
+    def __enter__(self):
+        (jump_type, size) = self._next_jump
+        self._jumps.append((self._pos, jump_type, size))
 
-    def unfreeze_pad(self):
-        align_pos = (self._pos - self._frozen_pos.pop()) % 4
-        if align_pos:
-            self.read(4 - align_pos)
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            (start, jump_type, size) = self._jumps.pop()
+            if jump_type:
+                self.read(size - (self._pos - start))
+            else:
+                align = (self._pos - start) % size
+                if align:
+                    self.read(size - align)
 
     def unpacks(self, format, size):
         return struct.unpack(format, self.read(size))
@@ -83,177 +93,174 @@ class OleUserFormParser(object):
     def unpack(self, format, size):
         return self.unpacks(format, size)[0]
 
+    def raise_error(self, reason, back=0):
+        raise OleFormParsingError('{0}:{1}: {2}'.format(self.path, self._pos - back))
+
     def check_values(self, name, format, size, expected):
         value = self.unpacks(format, size)
         if value != expected:
-            raise OleFormParsingError('Invalid {0} at {1}: expected {2} got {3}'.format(name, self._pos - size, str(expected), str(value)))
+            self.raise_error('Invalid {0}: expected {1} got {2}'.format(name, str(expected), str(value)))
 
     def check_value(self, name, format, size, expected):
         self.check_values(name, format, size, (expected,))
 
-    def consume_TextProps(self):
-        # TextProps: [MS-OFORMS] 2.3.1
-        self.check_values('TextProps (versions)', '<BB', 2, (0, 2))
-        cbTextProps = self.unpack('<H', 2)
-        self.read(cbTextProps)
-
-    def consume_GuidAndFont(self):
-        # GuidAndFont: [MS-OFORMS] 2.4.7
-        UUIDS = self.unpacks('<LHH', 8) + self.unpacks('>Q', 8)
-        if UUIDS == (199447043, 36753, 4558, 11376937813817407569L):
-            # UUID == {0BE35203-8F91-11CE-9DE300AA004BB851}
-            # StdFont: [MS-OFORMS] 2.4.12
-            self.check_value('StdFont (version)', '<B', 1, 1)
-            # Skip sCharset, bFlags, sWeight, ulHeight
-            self.read(9)
-            bFaceLen = self.unpack('<B', 1)
-            self.read(bFaceLen)
-        elif UUIDs == (2948729120, 55886, 4558, 13349514450607572916L):
-            # UUID == {AFC20920-DA4E-11CE-B94300AA006887B4}
-            self.consume_TextProps()
-        else:
-            raise OleFormParsingError('Invalid GuidAndFont at {0}: UUID'.format(self._pos - 16))
-
-    def consume_GuidAndPicture(self):
-        # GuidAndPicture: [MS-OFORMS] 2.4.8
-        # UUID == {0BE35204-8F91-11CE-9DE3-00AA004BB851}
-        self.check_values('GuidAndPicture (UUID part 1)', '<LHH', 8, (199447044, 36753, 4558))
-        self.check_value('GuidAndPicture (UUID part 1)', '>Q', 8, 11376937813817407569L)
-        # StdPicture: [MS-OFORMS] 2.4.13
-        self.check_value('StdPicture (Preamble)', '<L', 4, 0x0000746C)
-        size = self.unpack('<L', 4)
-        self.read(size)
-
-    def consume_CountOfBytesWithCompressionFlag(self):
-        # CountOfBytesWithCompressionFlag or CountOfCharsWithCompressionFlag: [MS-OFORMS] 2.4.14.2 or 2.4.14.3
-        count = self.unpack('<L', 4)
-        if not count & 0x80000000 and count != 0:
-            raise OleFormParsingError('Uncompress string length at {0}', self._pos - 4)
-        return count & 0x7FFFFFFF
-
-    def consume_SiteClassInfo(self):
-       # SiteClassInfo: [MS-OFORMS] 2.2.10.10.1
-       self.check_value('SiteClassInfo (version)', '<H', 2, 0)
-       cbClassTable = self.unpack('<H', 2)
-       self.read(cbClassTable)
-
-    def consume_FormObjectDepthTypeCount(self):
-       # FormObjectDepthTypeCount: [MS-OFORMS] 2.2.10.7
-       (depth, mixed) = self.unpacks('<BB', 2)
-       if mixed & 0x80:
-           self.check_value('FormObjectDepthTypeCount (SITE_TYPE)', '<B', 1, 1)
-           return mixed ^ 0x80
-       if mixed != 1:
-           raise OleFormParsingError('Invalid FormObjectDepthTypeCount (SITE_TYPE) at {0}: expected 1 got {3}'.format(self._pos - 2, str(mixed)))
-       return 1
-
-    def consume_OleSiteConcreteControl(self):
-        # OleSiteConcreteControl: [MS-OFORMS] 2.2.10.12.1
-        self.check_value('OleSiteConcreteControl (version)', '<H', 2, 0)
-        cbSite = self.unpack('<H', 2)
-        self.freeze()
-        propmask = SitePropMask(self.unpack('<L', 4))
+
+def consume_TextProps(stream):
+    # TextProps: [MS-OFORMS] 2.3.1
+    stream.check_values('TextProps (versions)', '<BB', 2, (0, 2))
+    cbTextProps = stream.unpack('<H', 2)
+    stream.read(cbTextProps)
+
+def consume_GuidAndFont(stream):
+    # GuidAndFont: [MS-OFORMS] 2.4.7
+    UUIDS = stream.unpacks('<LHH', 8) + stream.unpacks('>Q', 8)
+    if UUIDS == (199447043, 36753, 4558, 11376937813817407569L):
+        # UUID == {0BE35203-8F91-11CE-9DE300AA004BB851}
+        # StdFont: [MS-OFORMS] 2.4.12
+        stream.check_value('StdFont (version)', '<B', 1, 1)
+        # Skip sCharset, bFlags, sWeight, ulHeight
+        stream.read(9)
+        bFaceLen = stream.unpack('<B', 1)
+        stream.read(bFaceLen)
+    elif UUIDs == (2948729120, 55886, 4558, 13349514450607572916L):
+        # UUID == {AFC20920-DA4E-11CE-B94300AA006887B4}
+        consume_TextProps(stream)
+    else:
+        stream.raise_error('Invalid GuidAndFont (UUID)', 16)
+
+def consume_GuidAndPicture(stream):
+    # GuidAndPicture: [MS-OFORMS] 2.4.8
+    # UUID == {0BE35204-8F91-11CE-9DE3-00AA004BB851}
+    stream.check_values('GuidAndPicture (UUID part 1)', '<LHH', 8, (199447044, 36753, 4558))
+    stream.check_value('GuidAndPicture (UUID part 1)', '>Q', 8, 11376937813817407569L)
+    # StdPicture: [MS-OFORMS] 2.4.13
+    stream.check_value('StdPicture (Preamble)', '<L', 4, 0x0000746C)
+    size = stream.unpack('<L', 4)
+    stream.read(size)
+
+def consume_CountOfBytesWithCompressionFlag(stream):
+    # CountOfBytesWithCompressionFlag or CountOfCharsWithCompressionFlag: [MS-OFORMS] 2.4.14.2 or 2.4.14.3
+    count = stream.unpack('<L', 4)
+    if not count & 0x80000000 and count != 0:
+        stream.aise_error('Uncompress string length', 4)
+    return count & 0x7FFFFFFF
+
+def consume_SiteClassInfo(stream):
+   # SiteClassInfo: [MS-OFORMS] 2.2.10.10.1
+   stream.check_value('SiteClassInfo (version)', '<H', 2, 0)
+   cbClassTable = stream.unpack('<H', 2)
+   stream.read(cbClassTable)
+
+def consume_FormObjectDepthTypeCount(stream):
+   # FormObjectDepthTypeCount: [MS-OFORMS] 2.2.10.7
+   (depth, mixed) = stream.unpacks('<BB', 2)
+   if mixed & 0x80:
+       stream.check_value('FormObjectDepthTypeCount (SITE_TYPE)', '<B', 1, 1)
+       return mixed ^ 0x80
+   if mixed != 1:
+       stream.raise_error('Invalid FormObjectDepthTypeCount (SITE_TYPE): expected 1 got {0}'.format(str(mixed)))
+   return 1
+
+def consume_OleSiteConcreteControl(stream):
+    # OleSiteConcreteControl: [MS-OFORMS] 2.2.10.12.1
+    stream.check_value('OleSiteConcreteControl (version)', '<H', 2, 0)
+    cbSite = stream.unpack('<H', 2)
+    with stream.will_jump_to(cbSite):
+        propmask = SitePropMask(stream.unpack('<L', 4))
         # SiteDataBlock: [MS-OFORMS] 2.2.10.12.3
         name_len = tag_len = id = 0
         if propmask.fName:
-            name_len = self.consume_CountOfBytesWithCompressionFlag()
+            name_len = consume_CountOfBytesWithCompressionFlag(stream)
         if propmask.fTag:
-            tag_len = self.consume_CountOfBytesWithCompressionFlag()
+            tag_len = consume_CountOfBytesWithCompressionFlag(stream)
         if propmask.fID:
-            id = self.unpack('<L', 4)
+            id = stream.unpack('<L', 4)
         for prop in ['fHelpContextID', 'fBitFlags', 'fObjectStreamSize']:
             if propmask[prop]:
-                self.read(4)
+                stream.read(4)
         tabindex = ClsidCacheIndex = 0
-        self.freeze()
-        if propmask.fTabIndex:
-            tabindex = self.unpack('<H', 2)
-        if propmask.fClsidCacheIndex:
-            ClsidCacheIndex = self.unpack('<H', 2)
-        if propmask.fGroupID:
-            self.read(2)
-        self.unfreeze_pad()
+        with stream.will_pad():
+            if propmask.fTabIndex:
+                tabindex = stream.unpack('<H', 2)
+            if propmask.fClsidCacheIndex:
+                ClsidCacheIndex = stream.unpack('<H', 2)
+            if propmask.fGroupID:
+                stream.read(2)
         # For the next 4 entries, the documentation adds padding, but it should already be aligned??
         for prop in ['fControlTipText', 'fRuntimeLicKey', 'fControlSource', 'fRowSource']:
             if propmask[prop]:
-                self.read(4)
+                stream.read(4)
         # SiteExtraDataBlock: [MS-OFORMS] 2.2.10.12.4
-        name = self.read(name_len)
-        tag = self.read(tag_len)
-        self.variables.append({'name': name, 'tag': tag, 'id': id,
-                               'tabindex': tabindex,
-                               'ClsidCacheIndex': ClsidCacheIndex})
-        self.unfreeze(cbSite)
-
-    def consume_FormControl(self):
-        # FormControl: [MS-OFORMS] 2.2.10.1
-        self.check_values('FormControl (versions)', '<BB', 2, (0, 4))
-        cbform = self.unpack('<H', 2)
-        self.freeze()
-        propmask = FormPropMask(self.unpack('<L', 4))
+        name = stream.read(name_len)
+        tag = stream.read(tag_len)
+        return {'name': name, 'tag': tag, 'id': id, 'tabindex': tabindex,
+               'ClsidCacheIndex': ClsidCacheIndex}
+
+def consume_FormControl(stream):
+    # FormControl: [MS-OFORMS] 2.2.10.1
+    stream.check_values('FormControl (versions)', '<BB', 2, (0, 4))
+    cbform = stream.unpack('<H', 2)
+    with stream.will_jump_to(cbform):
+        propmask = FormPropMask(stream.unpack('<L', 4))
         # FormDataBlock: [MS-OFORMS] 2.2.10.3
         for prop in ['fBackColor', 'fForeColor', 'fNextAvailableID']:
             if propmask[prop]:
-                self.read(4)
+                stream.read(4)
         if propmask.fBooleanProperties:
-            BooleanProperties = self.unpack('<L', 4)
+            BooleanProperties = stream.unpack('<L', 4)
             FORM_FLAG_DONTSAVECLASSTABLE = (BooleanProperties & (1<<15)) >> 15
         else:
             FORM_FLAG_DONTSAVECLASSTABLE = 0
         # Skip the rest of DataBlock and ExtraDataBlock
-        self.unfreeze(cbform)
-        # FormStreamData: [MS-OFORMS] 2.2.10.5
-        if propmask.fMouseIcon:
-            self.consume_GuidAndPicture()
-        if propmask.fFont:
-            self.consume_GuidAndFont()
-        if propmask.fPicture:
-            self.consume_GuidAndPicture()
-        # FormSiteData: [MS-OFORMS] 2.2.10.6
-        if not FORM_FLAG_DONTSAVECLASSTABLE:
-            CountOfSiteClassInfo = self.unpack('<H', 2)
-            for i in range(CountOfSiteClassInfo):
-                self.consume_SiteClassInfo()
-        (CountOfSites, CountOfBytes) = self.unpacks('<LL', 8)
-        remaining_SiteDepthsAndTypes = CountOfSites
-        self.freeze()
+    # FormStreamData: [MS-OFORMS] 2.2.10.5
+    if propmask.fMouseIcon:
+        consume_GuidAndPicture(stream)
+    if propmask.fFont:
+        consume_GuidAndFont(stream)
+    if propmask.fPicture:
+        consume_GuidAndPicture(stream)
+    # FormSiteData: [MS-OFORMS] 2.2.10.6
+    if not FORM_FLAG_DONTSAVECLASSTABLE:
+        CountOfSiteClassInfo = stream.unpack('<H', 2)
+        for i in range(CountOfSiteClassInfo):
+            consume_SiteClassInfo(stream)
+    (CountOfSites, CountOfBytes) = stream.unpacks('<LL', 8)
+    remaining_SiteDepthsAndTypes = CountOfSites
+    with stream.will_pad():
         while remaining_SiteDepthsAndTypes > 0:
-            remaining_SiteDepthsAndTypes -= self.consume_FormObjectDepthTypeCount()
-        self.unfreeze_pad()
-        for i in range(CountOfSites):
-            self.consume_OleSiteConcreteControl()
-
-    def consume_MorphDataControl(self):
-        # MorphDataControl: [MS-OFORMS] 2.2.5.1
-        self.check_values('MorphDataControl (versions)', '<BB', 2, (0, 2))
-        cbMorphData = self.unpack('<H', 2)
-        self.freeze()
-        propmask = MorphDataPropMask(self.unpack('<Q', 8))
+            remaining_SiteDepthsAndTypes -= consume_FormObjectDepthTypeCount(stream)
+    for i in range(CountOfSites):
+        yield consume_OleSiteConcreteControl(stream)
+
+def consume_MorphDataControl(stream):
+    # MorphDataControl: [MS-OFORMS] 2.2.5.1
+    stream.check_values('MorphDataControl (versions)', '<BB', 2, (0, 2))
+    cbMorphData = stream.unpack('<H', 2)
+    with stream.will_jump_to(cbMorphData):
+        propmask = MorphDataPropMask(stream.unpack('<Q', 8))
         # MorphDataDataBlock: [MS-OFORMS] 2.2.5.3
         for prop in ['fVariousPropertyBits', 'fBackColor', 'fForeColor', 'fMaxLength']:
             if propmask[prop]:
-                self.read(4)
-        self.freeze()
-        for prop in ['fBorderStyle', 'fScrollBars', 'fDisplayStyle', 'fMousePointer']:
-            if propmask[prop]:
-                self.read(1)
-        self.unfreeze_pad()
+                stream.read(4)
+        with stream.will_pad():
+            for prop in ['fBorderStyle', 'fScrollBars', 'fDisplayStyle', 'fMousePointer']:
+                if propmask[prop]:
+                    stream.read(1)
         # PasswordChar, BoundColumn, TextColumn, ColumnCount, and ListRows are 2B + pad = 4B
         # ListWidth is 4B + pad = 4B
         for prop in ['fPasswordChar', 'fListWidth', 'fBoundColumn', 'fTextColumn', 'fColumnCount',
                      'fListRows']:
             if propmask[prop]:
-                self.read(4)
-        self.freeze()
-	if propmask.fcColumnInfo:
-            self.read(2)
-        for prop in ['fMatchEntry', 'fListStyle', 'fShowDropButtonWhen', 'fDropButtonStyle',
-                     'fMultiSelect']:
-            if propmask[prop]:
-                self.read(1)
-        self.unfreeze_pad()
+                stream.read(4)
+        with stream.will_pad():
+            if propmask.fcColumnInfo:
+                stream.read(2)
+            for prop in ['fMatchEntry', 'fListStyle', 'fShowDropButtonWhen', 'fDropButtonStyle',
+                         'fMultiSelect']:
+                if propmask[prop]:
+                    stream.read(1)
         if propmask.fValue:
-            value_size = self.consume_CountOfBytesWithCompressionFlag()
+            value_size = consume_CountOfBytesWithCompressionFlag(stream)
         else:
             value_size = 0
         # Caption, PicturePosition, BorderColor, SpecialEffect, GroupName  are 4B + pad = 4B
@@ -261,27 +268,24 @@ class OleUserFormParser(object):
         for prop in ['fCaption', 'fPicturePosition', 'fBorderColor', 'fSpecialEffect',
                      'fMouseIcon', 'fPicture', 'fAccelerator', 'fGroupName']:
             if propmask[prop]:
-                self.read(4)
+                stream.read(4)
         # MorphDataExtraDataBlock: [MS-OFORMS] 2.2.5.4
-        self.read(8)
-        value = self.read(value_size)
-        self.unfreeze(cbMorphData)
-        # MorphDataStreamData: [MS-OFORMS] 2.2.5.5
-        if propmask.fMouseIcon:
-            self.consume_GuidAndPicture()
-        if propmask.fPicture:
-            self.consume_GuidAndPicture()
-        self.consume_TextProps()
-        return value
-
-    def consume_stored_data(self):
-        for var in self.variables:
-            if var['ClsidCacheIndex'] != 23:
-                raise OleFormParsingError('Unsupported stored type: {0}'.format(str(var['ClsidCacheIndex'])))
-            var['value'] = self.consume_MorphDataControl()
-
-def OleFormVariables(ole_file, stream_dir):
-    control_stream = ole_file.openstream('/'.join(stream_dir + ['f']))
-    data_stream = ole_file.openstream('/'.join(stream_dir + ['o']))
-    form = OleUserFormParser(control_stream, data_stream)
-    return form.variables
+        stream.read(8)
+        value = stream.read(value_size)
+    # MorphDataStreamData: [MS-OFORMS] 2.2.5.5
+    if propmask.fMouseIcon:
+        consume_GuidAndPicture(stream)
+    if propmask.fPicture:
+        consume_GuidAndPicture(stream)
+    consume_TextProps(stream)
+    return value
+
+def extract_OleFormVariables(ole_file, stream_dir):
+    control = ExtendedStream.open(ole_file, '/'.join(stream_dir + ['f']))
+    variables = list(consume_FormControl(control))
+    data = ExtendedStream.open(ole_file, '/'.join(stream_dir + ['o']))
+    for var in variables:
+        if var['ClsidCacheIndex'] != 23:
+            raise OleFormParsingError('Unsupported stored type: {0}'.format(str(var['ClsidCacheIndex'])))
+        var['value'] = consume_MorphDataControl(data)
+    return variables
diff --git a/oletools/olevba.py b/oletools/olevba.py
index 9395e9f..ae0032b 100644
--- a/oletools/olevba.py
+++ b/oletools/olevba.py
@@ -254,7 +254,7 @@ except ImportError:
                                + "see http://codespeak.net/lxml " \
                                + "or http://effbot.org/zone/element-index.htm")
 
-from oleform import OleFormVariables
+from oleform import extract_OleFormVariables
 
 import thirdparty.olefile as olefile
 from thirdparty.prettytable import prettytable
@@ -2924,7 +2924,7 @@ class VBA_Parser(object):
             self.find_vba_forms()
             ole = self.ole_file
             for form_storage in self.vba_forms:
-                for variable in OleFormVariables(ole, form_storage):
+                for variable in extract_OleFormVariables(ole, form_storage):
                     yield (self.filename, '/'.join(form_storage), variable)
 
     def close(self):
--
libgit2 0.21.4