test_ese.py

#!/usr/bin/env python
# Impacket - Collection of Python classes for working with network protocols.
#
# Copyright Fortra, LLC and its affiliated companies
#
# All rights reserved.
#
# This software is provided under a slightly modified version
# of the Apache Software License. See the accompanying LICENSE file
# for more information.
#
import struct
import unittest

from impacket.ese import (
    ESENT_PAGE,
    ESENT_PAGE_HEADER,
    FIRST_AVAILABLE_PAGE_TAG_MASK,
    FIRST_AVAILABLE_PAGE_TAG_RESERVED_SHIFT,
)


class TestESENTLargePageTags(unittest.TestCase):
    PAGE_SIZE = 32768
    VERSION = 0x620
    REVISION = 0x122

    def _build_page(self, raw_tag_state, payloads=None):
        tag_count = raw_tag_state & FIRST_AVAILABLE_PAGE_TAG_MASK
        header = ESENT_PAGE_HEADER(self.VERSION, self.REVISION, self.PAGE_SIZE)
        header['FirstAvailableDataOffset'] = tag_count * 4
        header['FirstAvailablePageTag'] = raw_tag_state
        header['PageFlags'] = 0x2001

        header_bytes = header.getData()
        page = bytearray(self.PAGE_SIZE)
        page[:len(header_bytes)] = header_bytes

        base_offset = len(header_bytes)
        if payloads is None:
            payloads = [bytes([i, i, i, i]) for i in range(tag_count)]
        self.assertEqual(len(payloads), tag_count)

        tags = []
        offset = 0
        for payload in payloads:
            page[base_offset + offset:base_offset + offset + len(payload)] = payload
            tags.append(struct.pack('<HH', len(payload), offset))
            offset += len(payload)

        page[-4 * tag_count:] = b''.join(reversed(tags))

        return ESENT_PAGE({
            'Version': self.VERSION,
            'FileFormatRevision': self.REVISION,
            'PageSize': self.PAGE_SIZE,
        }, bytes(page)), payloads

    def test_splits_large_page_tag_state(self):
        page, _ = self._build_page(0x100c)

        self.assertEqual(page.record['FirstAvailablePageTag'], 0x100c)
        self.assertEqual(page.tagReserved, (0x100c >> FIRST_AVAILABLE_PAGE_TAG_RESERVED_SHIFT) or 1)
        self.assertEqual(page.tagCount, 0x100c & FIRST_AVAILABLE_PAGE_TAG_MASK)

    def test_iterates_large_page_tags_with_high_bits_set(self):
        page, payloads = self._build_page(0x100c)

        for tag_num in range(1, page.tagCount):
            self.assertEqual(page.getTag(tag_num), (0, payloads[tag_num]))

        with self.assertRaisesRegex(Exception, r'unknown tag'):
            page.getTag(page.tagCount)

    def test_iter_data_tags_skips_all_reserved_large_page_tags(self):
        page, payloads = self._build_page(0x200c)

        self.assertEqual(page.tagReserved, 2)
        self.assertEqual(page.firstDataTag, 2)
        self.assertEqual(list(page.iterDataTagNums()), list(range(2, page.tagCount)))

        iterated_payloads = [page.getTag(tag_num)[1] for tag_num in page.iterDataTagNums()]
        self.assertEqual(iterated_payloads, payloads[page.firstDataTag:])

    def test_large_page_tag_handles_empty_and_single_byte_values(self):
        page, _ = self._build_page(0x1003, [
            b'\x11\x22',
            b'',
            b'\xab',
        ])

        self.assertEqual(page.getTag(1), (0, b''))
        self.assertEqual(page.getTag(2), (0, b'\xab'))


if __name__ == '__main__':
    unittest.main(verbosity=1)