75 lines
3.3 KiB
Python
75 lines
3.3 KiB
Python
|
import base64
|
||
|
|
import hashlib
|
||
|
|
import io
|
||
|
|
from pathlib import Path
|
||
|
|
import struct
|
||
|
|
import unittest
|
||
|
|
import zipfile
|
||
|
|
from engine import scan_download
|
||
|
|
from office_macros import bounded_decompress, macro_payloads, MAX_CODE
|
||
|
|
|
||
|
|
|
||
|
|
def fixture():
|
||
|
|
path = Path(__file__).resolve().parents[2] / 'tests/fixtures/office/benign-vba-project.base64.txt'
|
||
|
|
data = base64.b64decode(path.read_text(), validate=False)
|
||
|
|
assert hashlib.sha256(data).hexdigest() == '0ced1464b3677e98f5e3a8c5d80135e18dc98dca39299f1a8cfd2a00999fbf9f'
|
||
|
|
return data
|
||
|
|
|
||
|
|
|
||
|
|
def office(project=None, extra=None):
|
||
|
|
output = io.BytesIO()
|
||
|
|
with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as archive:
|
||
|
|
archive.writestr('xl/workbook.xml', '<workbook/>')
|
||
|
|
if project is not None: archive.writestr('xl/vbaProject.bin', project)
|
||
|
|
if extra: archive.writestr(*extra)
|
||
|
|
return output.getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
def policy(source='rule clean { condition: false }', enabled=True):
|
||
|
|
return {'security': {'yara': {'enabled': True, 'scanOfficeMacros': enabled,
|
||
|
|
'rules': [{'id': 'macro-marker', 'source': source}]}}}
|
||
|
|
|
||
|
|
|
||
|
|
class OfficeMacroTests(unittest.TestCase):
|
||
|
|
def test_real_compressed_source_is_scanned(self):
|
||
|
|
data = office(fixture())
|
||
|
|
self.assertNotIn(b'Hello from Python!', data)
|
||
|
|
result = scan_download(data, policy('rule source_marker { strings: $m = "Hello from Python!" condition: $m }'))
|
||
|
|
self.assertEqual(result['decision'], 'block')
|
||
|
|
self.assertEqual(result['ruleId'], 'macro-marker')
|
||
|
|
self.assertEqual(result['officeMacros']['modules'], 5)
|
||
|
|
self.assertTrue(result['officeMacros']['present'])
|
||
|
|
self.assertNotIn('Hello from Python!', str(result))
|
||
|
|
|
||
|
|
def test_benign_macros_are_not_blanket_blocked(self):
|
||
|
|
result = scan_download(office(fixture()), policy())
|
||
|
|
self.assertEqual(result['decision'], 'allow')
|
||
|
|
self.assertTrue(result['officeMacros']['complete'])
|
||
|
|
|
||
|
|
def test_corrupt_project_fails_closed_and_can_be_disabled(self):
|
||
|
|
data = office(b'bad project')
|
||
|
|
self.assertEqual(scan_download(data, policy())['ruleId'], 'office-macro-inspection-failed')
|
||
|
|
self.assertEqual(scan_download(data, policy(enabled=False))['decision'], 'allow')
|
||
|
|
|
||
|
|
def test_macro_free_office(self):
|
||
|
|
summary, payloads = macro_payloads(office())
|
||
|
|
self.assertEqual(summary, {'office': True, 'present': False, 'modules': 0, 'complete': True})
|
||
|
|
self.assertEqual(payloads, [])
|
||
|
|
|
||
|
|
def test_legacy_project(self):
|
||
|
|
self.assertEqual(macro_payloads(fixture())[0]['modules'], 5)
|
||
|
|
|
||
|
|
def test_unsupported_embedded_or_xlm(self):
|
||
|
|
for name in ('xl/macrosheets/sheet1.xml', 'word/embeddings/object.bin'):
|
||
|
|
with self.subTest(name=name):
|
||
|
|
self.assertEqual(scan_download(office(extra=(name, 'x')), policy())['decision'], 'block')
|
||
|
|
|
||
|
|
def test_strict_compression_and_size_limits(self):
|
||
|
|
self.assertEqual(bounded_decompress(b'\x01\x03\xb0\x00abc'), b'abc')
|
||
|
|
for data in (b'', b'\x01\x03\xb0\x00a', b'\x01\x03\xb0\x01\xff\xff',
|
||
|
|
b'\x01' + (struct.pack('<H', 0x3fff) + b'x' * 4096) * (MAX_CODE // 4096 + 1)):
|
||
|
|
with self.subTest(length=len(data)):
|
||
|
|
with self.assertRaises(ValueError): bounded_decompress(data)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__': unittest.main()
|