"""Bounded in-memory Office VBA extraction. Never executes macros or writes files.

Source matching supplements original/compiled-stream YARA and enterprise ICAP;
it is not VBA execution, deobfuscation, p-code equivalence or an AV verdict.
"""
import io
import logging
import struct
import zipfile

OLE_MAGIC = b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1'
MAX_EXPANDED = 20 * 1024 * 1024
MAX_CODE = 4 * 1024 * 1024
MAX_MODULES = 128


def bounded_decompress(data):
    """Strict MS-OVBA 2.4.1 decompression with per-chunk and total limits."""
    if not data or data[0] != 1:
        raise ValueError('macro-compression-invalid')
    output = bytearray(); cursor = 1
    while cursor < len(data):
        if cursor + 2 > len(data): raise ValueError('macro-compression-invalid')
        header = struct.unpack_from('<H', data, cursor)[0]
        size = (header & 4095) + 3; end = cursor + size
        if ((header >> 12) & 7) != 3 or end > len(data):
            raise ValueError('macro-compression-invalid')
        cursor += 2; start = len(output)
        if not header & 0x8000:
            if size != 4098: raise ValueError('macro-compression-invalid')
            output.extend(data[cursor:end]); cursor = end
        else:
            while cursor < end:
                flags = data[cursor]; cursor += 1
                for bit in range(8):
                    if cursor >= end: break
                    if flags & (1 << bit):
                        if cursor + 2 > end: raise ValueError('macro-compression-invalid')
                        token = struct.unpack_from('<H', data, cursor)[0]; cursor += 2
                        count = len(output) - start
                        bits = max(4, (count - 1).bit_length())
                        mask = 0xffff >> bits
                        length = (token & mask) + 3
                        offset = (token >> (16 - bits)) + 1
                        if offset > count or count + length > 4096:
                            raise ValueError('macro-compression-invalid')
                        for _ in range(length): output.append(output[-offset])
                    else:
                        if len(output) - start >= 4096: raise ValueError('macro-compression-invalid')
                        output.append(data[cursor]); cursor += 1
        if len(output) > MAX_CODE: raise ValueError('macro-code-limit')
    return bytes(output)


def macro_payloads(data):
    """Return (metadata, YARA payloads); reject ambiguous or unsupported Office."""
    summary = {'office': False, 'present': False, 'modules': 0, 'complete': True}
    projects = []; payloads = []; ooxml = False; payload_bytes = 0
    def append_payload(payload):
        nonlocal payload_bytes
        payload_bytes += len(payload)
        if len(payloads) >= 512 or payload_bytes > 2 * MAX_EXPANDED:
            raise ValueError('macro-payload-limit')
        payloads.append(payload)
    if data.startswith(OLE_MAGIC):
        summary['office'] = True; projects.append(data)
    elif data.startswith(b'PK'):
        with zipfile.ZipFile(io.BytesIO(data)) as archive:
            entries = archive.infolist()
            if len(entries) > 2048 or sum(x.file_size for x in entries) > MAX_EXPANDED:
                raise ValueError('macro-archive-limit')
            names = [x.filename.lower() for x in entries]
            if len(set(names)) != len(names): raise ValueError('macro-archive-ambiguous')
            summary['office'] = any(n in names for n in ('word/document.xml', 'xl/workbook.xml', 'xl/workbook.bin', 'ppt/presentation.xml'))
            if not summary['office']: return summary, []
            ooxml = True
            if 'xl/workbook.bin' in names: raise ValueError('macro-xlsb-unsupported')
            for entry, name in zip(entries, names):
                if entry.flag_bits & 1: raise ValueError('macro-encrypted')
                if '/embeddings/' in name or '/macrosheets/' in name or name.endswith('vbadata.xml'):
                    raise ValueError('macro-format-unsupported')
                if name.endswith('vbaproject.bin'):
                    if len(projects) >= 16 or entry.file_size > MAX_CODE: raise ValueError('macro-project-limit')
                    project = archive.read(entry)
                    if not project.startswith(OLE_MAGIC): raise ValueError('macro-project-invalid')
                    projects.append(project); append_payload(project)
    else:
        return summary, []
    if not projects: return summary, payloads
    import olefile
    from oletools import olevba
    # Silent third-party parser: module/path/code fragments must never reach logs.
    for name in ('olevba', 'olefile', 'crypto', 'ppt', 'ppt_parser', 'ftguess', 'oleform', 'rtfobj'):
        logging.getLogger(name).disabled = True
    previous = olevba.decompress_stream
    olevba.decompress_stream = bounded_decompress
    total = 0
    try:
        for project in projects:
            with olefile.OleFileIO(io.BytesIO(project), raise_defects=olefile.DEFECT_INCORRECT) as ole:
                streams = ole.listdir(streams=True, storages=False)
                if len(streams) > 2048 or sum(ole.get_size(s) for s in streams) > MAX_EXPANDED:
                    raise ValueError('macro-ole-limit')
                roots = {s[0].lower() for s in streams}
                if 'encryptedpackage' in roots or 'encryptioninfo' in roots:
                    raise ValueError('macro-encrypted')
                # Unsupported legacy Excel XLM / embedded Office containers are not safe negatives.
                if any(s[-1].lower() in ('workbook', 'book') for s in streams):
                    raise ValueError('macro-xlm-unsupported')
                has_vba = any('vba' in [p.lower() for p in s] for s in streams)
                if ooxml and not has_vba: raise ValueError('macro-project-invalid')
                # Contiguous streams expose compiled/p-code bytes to existing rules.
                # They do not establish source/p-code equivalence (VBA stomping).
                for stream in streams:
                    append_payload(ole.openstream(stream).read())
            summary['present'] = summary['present'] or has_vba
            parser = olevba.VBA_Parser('memory-office', data=project, relaxed=False)
            try:
                before = summary['modules']
                for _, _, _, code in parser.extract_macros():
                    if not isinstance(code, str): raise ValueError('macro-source-invalid')
                    encoded = code.encode('utf-8'); total += len(encoded)
                    summary['modules'] += 1
                    if summary['modules'] > MAX_MODULES or total > MAX_CODE:
                        raise ValueError('macro-code-limit')
                    append_payload(encoded)
                if has_vba and summary['modules'] == before:
                    raise ValueError('macro-source-unavailable')
            finally:
                parser.close()
    finally:
        olevba.decompress_stream = previous
    return summary, payloads