refactor: office: make nicer, add comments

plugins/analyzer_office.py

@@ -1,19 +1,19 @@
 
-import copy
 import logging
-from re import I
 import olefile
 from typing import List
 from reducer import Reducer
 from utils import *
-from model.model import Match, FileInfo
+from model.model import Match, FileInfo, Scanner
 import pcodedmp.pcodedmp as pcodedmp
-from plugins.file_office import FileOffice, VbaAddressConverter, AddressConverter
+from plugins.file_office import FileOffice, VbaAddressConverter, OleStructurizer
 from pcodedmp.disasm import DisasmEntry
 from intervaltree import Interval, IntervalTree
 
 
-def analyzeFileWord(fileOffice: FileOffice, scanner, analyzerOptions={}):
+def analyzeFileWord(fileOffice: FileOffice, scanner: Scanner, analyzerOptions={}) -> IntervalTree:
+    # Scans a office file given with fileOffice with Scanner scanner. 
+    # Returns all matches.
     makroData = fileOffice.data
 
     reducer = Reducer(fileOffice, scanner)
@@ -21,7 +21,46 @@ def analyzeFileWord(fileOffice: FileOffice, scanner, analyzerOptions={}):
     return matchesIntervalTree
 
 
-def convertResults(ole, results) -> IntervalTree:
+def augmentFileWord(fileOffice: FileOffice, matches: List[Match]) -> FileInfo:
+    # Augment the matches with VBA decompilation and section information.
+    # Returns a FileInfo object with detailed file information too.
+
+    # dump all makros as disassembled code
+    fd = open('/dev/null', 'w')
+    disasmList = pcodedmp.processFile(fileOffice.filepath, output_file=fd)
+    fd.close()
+
+    oleFile = olefile.OleFileIO(fileOffice.data)
+    disasmList = convertDisasmAddr(oleFile, disasmList)
+    ac = OleStructurizer(oleFile)
+
+    # correlate the matches with the dumped code
+    m: Match
+    for m in matches:
+        data = fileOffice.data[m.start():m.end()]
+        dataHexdump = hexdump.hexdump(data, result='return')
+        sectionName = ac.getSectionsForAddr(m.start(), m.size)
+
+        disasmMatches = disasmList.overlap(m.fileOffset, m.fileOffset+m.size)
+        details = []
+        for item in iter(disasmMatches):
+            detail = {}
+            detail['part'] = True
+            detail['textHtml'] = "{}-{} (line #{}): ".format(item.data.begin, item.data.end, item.data.lineNr) + "\n" + item.data.text
+            details.append(detail)
+
+        m.setData(data)
+        m.setDataHexdump(dataHexdump)
+        m.setInfo(sectionName)
+        m.setDetail(details)
+
+    fileInfo = FileInfo(fileOffice.filename, 0, ac.getStructure())
+    return fileInfo
+
+
+def convertDisasmAddr(ole: olefile.olefile.OleFileIO, results: List[DisasmEntry]) -> IntervalTree:
+    # the output of pcodedmp is wrong. Convert results to real physical addresses.
+    # use the extracted vbaProject.bin from fileOffice.data
     ac = VbaAddressConverter(ole)
     it = IntervalTree()
 
@@ -38,40 +77,4 @@ def convertResults(ole, results) -> IntervalTree:
             else:
                 it.add(Interval(physBegin, physEnd, ite.data))
       
-    return it
-
-
-def augmentFileWord(fileOffice: FileOffice, matches: List[Match]):
-    # dump makros as disassembled code
-    fd = open('/dev/null', 'w')
-    results = pcodedmp.processFile(fileOffice.filepath, output_file=fd)
-    fd.close()
-
-    # the output of pcodedmp is wrong. Convert results to real physical addresses.
-    # use the extracted vbaProject.bin from fileOffice.data
-    oleFile = olefile.OleFileIO(fileOffice.data)
-    results = convertResults(oleFile, results)
-    ac = AddressConverter(oleFile)
-
-    # correlate the matches with the dumped code
-    for m in matches:
-        data = fileOffice.data[m.start():m.end()]
-        dataHexdump = hexdump.hexdump(data, result='return')
-        sectionName = ac.getSectionsForAddr(m.start(), m.size)
-        detail = ''
-
-        itemSet = results.overlap(m.fileOffset, m.fileOffset+m.size)
-        details = []
-        for item in iter(itemSet):
-            detail = {}
-            detail['part'] = True
-            detail['textHtml'] = "{} {} {}: ".format(item.data.lineNr, item.data.begin, item.data.end) + "\n" + item.data.text
-            details.append(detail)
-
-        m.setData(data)
-        m.setDataHexdump(dataHexdump)
-        m.setInfo(sectionName)
-        m.setDetail(details)
-
-    fileInfo = FileInfo(fileOffice.filename, 0, ac.getStructure())
-    return fileInfo
+    return it

plugins/file_office.py

@@ -1,14 +1,18 @@
-import os
 import zipfile
 import io
 from model.model import PluginFileFormat
 import olefile
 from math import floor
+from typing import List
 
 MAKRO_PATH = 'word/vbaProject.bin'
 
 
 class FileOffice(PluginFileFormat):
+    # Represents an office file.
+    #   - fileData: the complete file content
+    #   - data: vbaProject.bin file content 
+
     def __init__(self):
         super().__init__()
 
@@ -31,10 +35,17 @@ class FileOffice(PluginFileFormat):
         return False
 
 
-    def getFileWithNewData(self, data):
+    def getFileWithNewData(self, data: bytes) -> bytes:
+        # get office file with vbaProject replaced
         return self.getPatchedByReplacement(data)
 
 
+    def getPatchedByOffset(self, offset: int, patch: bytes) -> bytes:
+        # get office file with parts of vbaProject replaced
+        goat = self.data[:offset] + patch + self.data[offset+len(patch):]
+        return self.getPatchedByReplacement(goat)
+
+
     def getPatchedByReplacement(self, data: bytes) -> bytes:
         outData = io.BytesIO()
 
@@ -56,12 +67,10 @@ class FileOffice(PluginFileFormat):
         return outData.getvalue()
 
 
-    def getPatchedByOffset(self, offset: int, patch: bytes) -> bytes:
-        goat = self.data[:offset] + patch + self.data[offset+len(patch):]
-        return self.getPatchedByReplacement(goat)
-
-
 class VbaAddressConverter():
+    # Given a OLE file, it can convert section-relative addresses 
+    # (e.g. those from pcodedmp output) into file-based offset addresses
+
     def __init__(self, ole: olefile.olefile.OleFileIO):
         self.ole = ole
         self.ministream = None
@@ -88,22 +97,15 @@ class VbaAddressConverter():
         self.ministream = arr
 
 
-    def _getDirForName(self, name:str) -> olefile.olefile.OleDirectoryEntry:
-        for id in range(len(self.ole.direntries)):
-            d: olefile.olefile.OleDirectoryEntry = self.ole.direntries[id]
-            if d is None:
-                continue
-            if d.name == name:
-                return d
-
-
     def physicalAddressFor(self, modulepath: str, offset: int) -> int:
+        # return the physical address of offset of modulepath (e.g. "VBA/ThisDocument")
+
         # sanity checks
         mp = modulepath.split('/')
         if len (mp) != 2:
-            return 0
+            return -1
         if mp[0] != 'VBA':
-            return 0
+            return -1
         moduleName = mp[1]
 
         # If the stream is >4096: use normal sectors
@@ -146,7 +148,19 @@ class VbaAddressConverter():
         return result
 
 
-class AddressConverter():
+    def _getDirForName(self, name:str) -> olefile.olefile.OleDirectoryEntry:
+        for id in range(len(self.ole.direntries)):
+            d: olefile.olefile.OleDirectoryEntry = self.ole.direntries[id]
+            if d is None:
+                continue
+            if d.name == name:
+                return d
+
+
+class OleStructurizer():
+    # Parses an OLE file, so it is possible to find the section
+    # of an address/offset into the file
+
     def __init__(self, ole: olefile.olefile.OleFileIO):
         self.ole = ole
         self.sector = None
@@ -189,7 +203,9 @@ class AddressConverter():
                 self._paintMinistreamSectorChain(ministreamSect, d.name, d.isectStart, d.size)
 
 
-    def getSectionForAddr(self, addr):
+    def getSectionForAddr(self, addr: int) -> str:
+        # given an offset into the OLE file, find the section containing it
+
         # find sector
         sector = roundTo(addr, self.ole.sectorsize) // self.ole.sector_size
         if sector == 0:
@@ -208,7 +224,8 @@ class AddressConverter():
             return self.sector[sector]
 
 
-    def getSectionsForAddr(self, addr, size):
+    def getSectionsForAddr(self, addr: int, size: int) -> List[str]:
+        # given a offset and its size into a file, find all sections covered by it
         res = {}
 
         # just brute force it...

tests/test_makro_disasm.py

@@ -3,7 +3,7 @@
 import unittest
 import pcodedmp.pcodedmp as pcodedmp
 from plugins.analyzer_office import augmentFileWord
-from plugins.file_office import FileOffice, AddressConverter, VbaAddressConverter
+from plugins.file_office import FileOffice, OleStructurizer, VbaAddressConverter
 import olefile
 from model.model import Match
 
@@ -39,7 +39,7 @@ class DisasmMakroTest(unittest.TestCase):
     def test_AddressConverterGetSection(self):
         file = 'tests/data/test.docm.vbaProject.bin'
         ole = olefile.OleFileIO(file)
-        ac = AddressConverter(ole)
+        ac = OleStructurizer(ole)
         self.assertEqual(ac.getSectionForAddr(0), "Header")
         self.assertEqual(ac.getSectionForAddr(1), "Header")
         self.assertEqual(ac.getSectionForAddr(512), "FAT Sector")
@@ -51,7 +51,7 @@ class DisasmMakroTest(unittest.TestCase):
     def test_AddressConverterGetSections(self):
         file = 'tests/data/test.docm.vbaProject.bin'
         ole = olefile.OleFileIO(file)
-        ac = AddressConverter(ole)
+        ac = OleStructurizer(ole)
         
         sections = ac.getSectionsForAddr(3584, 1024)
         print(str(sections))