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()