| #!/usr/bin/env python3 |
| # Copyright 2020 The Pigweed Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| # use this file except in compliance with the License. You may obtain a copy of |
| # the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| # License for the specific language governing permissions and limitations under |
| # the License. |
| """Tests for detokenize.""" |
| |
| import base64 |
| import concurrent |
| import datetime as dt |
| import functools |
| import io |
| import os |
| from pathlib import Path |
| import struct |
| import tempfile |
| from typing import Any, Callable, NamedTuple |
| import unittest |
| from unittest import mock |
| |
| from pw_tokenizer import database |
| from pw_tokenizer import detokenize |
| from pw_tokenizer import elf_reader |
| from pw_tokenizer import tokens |
| |
| |
| # This function is not part of this test. It was used to generate the binary |
| # strings for EMPTY_ELF and ELF_WITH_TOKENIZER_SECTIONS. It takes a path and |
| # returns a Python byte string suitable for copying into Python source code. |
| def path_to_byte_string(path): |
| with open(path, 'rb') as fd: |
| data = fd.read() |
| |
| output = [] |
| indices = iter(range(len(data))) |
| |
| while True: |
| line = '' |
| |
| while len(line) < 70: |
| try: |
| i = next(indices) |
| except StopIteration: |
| break |
| |
| line += repr(data[i : i + 1])[2:-1].replace("'", r'\'') |
| |
| if not line: |
| return ''.join(output) |
| |
| output.append(" b'{}'\n".format(''.join(line))) |
| |
| |
| # This is an empty ELF file. It was created from the ELF file for |
| # tokenize_test.cc with the command: |
| # |
| # arm-none-eabi-objcopy -S --only-section NO_SECTIONS_PLEASE <ELF> <OUTPUT> |
| # |
| # The resulting ELF was converted to a Python binary string using |
| # path_to_byte_string function above. |
| EMPTY_ELF = ( |
| b'\x7fELF\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00(\x00\x01' |
| b'\x00\x00\x00\xd1\x83\x00\x084\x00\x00\x00\xe0\x00\x00\x00\x00\x04\x00\x05' |
| b'4\x00 \x00\x05\x00(\x00\x02\x00\x01\x00\x01\x00\x00\x00\xd4\x00\x00\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x07\x00' |
| b'\x00\x00\x00\x00\x01\x00\x01\x00\x00\x00\xd4\x00\x00\x00\x00\x00\x00\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x06\x00\x00\x00\x00\x00' |
| b'\x01\x00\x01\x00\x00\x00\xd4\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x06\x00\x00\x00\x00\x00\x01\x00\x01\x00' |
| b'\x00\x00\xd4\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' |
| b'\x00\x00\x00\x00\x06\x00\x00\x00\x00\x00\x01\x00\x01\x00\x00\x00\xd4\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' |
| b'\x06\x00\x00\x00\x00\x00\x01\x00\x00.shstrtab\x00\x00\x00\x00\x00\x00\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' |
| b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01' |
| b'\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xd4\x00\x00' |
| b'\x00\x0b\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00' |
| b'\x00\x00\x00' |
| ) |
| |
| # This is an ELF file with only the pw_tokenizer sections. It was created |
| # from a tokenize_test binary built for the STM32F429i Discovery board. The |
| # pw_tokenizer sections were extracted with this command: |
| # |
| # arm-none-eabi-objcopy -S --only-section ".pw_tokenizer*" <ELF> <OUTPUT> |
| # |
| ELF_WITH_TOKENIZER_SECTIONS_PATH = Path(__file__).parent.joinpath( |
| 'example_binary_with_tokenized_strings.elf' |
| ) |
| ELF_WITH_TOKENIZER_SECTIONS = ELF_WITH_TOKENIZER_SECTIONS_PATH.read_bytes() |
| |
| TOKENS_IN_ELF = 22 |
| TOKENS_IN_ELF_WITH_TOKENIZER_SECTIONS = 26 |
| |
| # 0x2e668cd6 is 'Jello, world!' (which is also used in database_test.py). |
| JELLO_WORLD_TOKEN = b'\xd6\x8c\x66\x2e' |
| |
| |
| class DetokenizeTest(unittest.TestCase): |
| """Tests the detokenize.Detokenizer.""" |
| |
| def test_simple(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| 0xCDAB, '%02d %s %c%%', date_removed=dt.datetime.now() |
| ) |
| ] |
| ) |
| ) |
| self.assertEqual( |
| str(detok.detokenize(b'\xab\xcd\0\0\x02\x03Two\x66')), '01 Two 3%' |
| ) |
| |
| def test_detokenize_extra_data_is_unsuccessful(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| 1, 'no args', date_removed=dt.datetime(1, 1, 1) |
| ) |
| ] |
| ) |
| ) |
| |
| result = detok.detokenize(b'\x01\0\0\0\x04args') |
| self.assertEqual(len(result.failures), 1) |
| string, args, remaining = result.failures[0] |
| self.assertEqual('no args', string) |
| self.assertFalse(args) |
| self.assertEqual(b'\x04args', remaining) |
| self.assertEqual('no args', string) |
| self.assertEqual('no args', str(result)) |
| |
| def test_detokenize_zero_extend_short_token_with_no_args(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [tokens.TokenizedStringEntry(0xCDAB, 'This token is 16 bits')] |
| ) |
| ) |
| self.assertEqual( |
| str(detok.detokenize(b'\xab\xcd')), 'This token is 16 bits' |
| ) |
| |
| def test_detokenize_missing_data_is_unsuccessful(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| 2, '%s', date_removed=dt.datetime(1, 1, 1) |
| ) |
| ] |
| ) |
| ) |
| |
| result = detok.detokenize(b'\x02\0\0\0') |
| string, args, remaining = result.failures[0] |
| self.assertEqual('%s', string) |
| self.assertEqual(len(args), 1) |
| self.assertEqual(b'', remaining) |
| self.assertEqual(len(result.failures), 1) |
| self.assertEqual('%s', str(result)) |
| |
| def test_detokenize_missing_data_with_errors_is_unsuccessful(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| 2, '%s', date_removed=dt.datetime(1, 1, 1) |
| ) |
| ] |
| ), |
| show_errors=True, |
| ) |
| |
| result = detok.detokenize(b'\x02\0\0\0') |
| string, args, remaining = result.failures[0] |
| self.assertIn('%s MISSING', string) |
| self.assertEqual(len(args), 1) |
| self.assertEqual(b'', remaining) |
| self.assertEqual(len(result.failures), 1) |
| self.assertIn('%s MISSING', str(result)) |
| |
| def test_unparsed_data(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| 1, 'no args', date_removed=dt.datetime(100, 1, 1) |
| ), |
| ] |
| ) |
| ) |
| result = detok.detokenize(b'\x01\0\0\0o_o') |
| self.assertFalse(result.ok()) |
| self.assertEqual('no args', str(result)) |
| self.assertIn('o_o', repr(result)) |
| self.assertIn('decoding failed', result.error_message()) |
| |
| def test_empty_db(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF)) |
| self.assertFalse(detok.detokenize(b'\x12\x34\0\0').ok()) |
| self.assertIn( |
| 'unknown token', detok.detokenize(b'1234').error_message() |
| ) |
| self.assertIn('unknown token', repr(detok.detokenize(b'1234'))) |
| |
| self.assertEqual( |
| '$' + base64.b64encode(b'1234').decode(), |
| str(detok.detokenize(b'1234')), |
| ) |
| |
| self.assertIsNone(detok.detokenize(b'').token) |
| |
| def test_empty_db_show_errors(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF), show_errors=True) |
| self.assertFalse(detok.detokenize(b'\x12\x34\0\0').ok()) |
| self.assertIn( |
| 'unknown token', detok.detokenize(b'1234').error_message() |
| ) |
| self.assertIn('unknown token', repr(detok.detokenize(b'1234'))) |
| self.assertIn('unknown token', str(detok.detokenize(b'1234'))) |
| |
| self.assertIsNone(detok.detokenize(b'').token) |
| |
| def test_missing_token_show_errors(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF), show_errors=True) |
| self.assertIn('missing token', detok.detokenize(b'').error_message()) |
| self.assertIn('missing token', str(detok.detokenize(b''))) |
| |
| def test_missing_token(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF)) |
| self.assertIn('missing token', detok.detokenize(b'').error_message()) |
| self.assertEqual('$', str(detok.detokenize(b''))) |
| |
| def test_unknown_shorter_token_show_error(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF), show_errors=True) |
| |
| self.assertIn('unknown token', detok.detokenize(b'1').error_message()) |
| self.assertIn('unknown token', str(detok.detokenize(b'1'))) |
| self.assertIn('unknown token', repr(detok.detokenize(b'1'))) |
| |
| self.assertIn('unknown token', detok.detokenize(b'123').error_message()) |
| self.assertIn('unknown token', str(detok.detokenize(b'123'))) |
| self.assertIn('unknown token', repr(detok.detokenize(b'123'))) |
| |
| def test_unknown_shorter_token(self): |
| detok = detokenize.Detokenizer(io.BytesIO(EMPTY_ELF)) |
| |
| self.assertEqual( |
| 'unknown token 00000001', detok.detokenize(b'\1').error_message() |
| ) |
| self.assertEqual( |
| '$' + base64.b64encode(b'\1\0\0\0').decode(), |
| str(detok.detokenize(b'\1')), |
| ) |
| self.assertIn('unknown token 00000001', repr(detok.detokenize(b'\1'))) |
| |
| self.assertEqual( |
| 'unknown token 00030201', |
| detok.detokenize(b'\1\2\3').error_message(), |
| ) |
| self.assertEqual( |
| '$' + base64.b64encode(b'\1\2\3\0').decode(), |
| str(detok.detokenize(b'\1\2\3')), |
| ) |
| self.assertIn( |
| 'unknown token 00030201', repr(detok.detokenize(b'\1\2\3')) |
| ) |
| |
| def test_decode_from_elf_data(self): |
| detok = detokenize.Detokenizer(io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS)) |
| |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| self.assertEqual( |
| str(detok.detokenize(JELLO_WORLD_TOKEN)), 'Jello, world!' |
| ) |
| |
| undecoded_args = detok.detokenize(JELLO_WORLD_TOKEN + b'some junk') |
| self.assertFalse(undecoded_args.ok()) |
| self.assertEqual(str(undecoded_args), 'Jello, world!') |
| |
| self.assertTrue(detok.detokenize(b'\0\0\0\0').ok()) |
| self.assertEqual(str(detok.detokenize(b'\0\0\0\0')), '') |
| |
| def test_decode_from_elf_file(self): |
| """Test decoding from an elf file.""" |
| detok = detokenize.Detokenizer(io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS)) |
| expected_tokens = frozenset(detok.database.token_to_entries.keys()) |
| |
| with tempfile.NamedTemporaryFile('wb', delete=False) as elf: |
| try: |
| elf.write(ELF_WITH_TOKENIZER_SECTIONS) |
| elf.close() |
| |
| # Open ELF by file object |
| with open(elf.name, 'rb') as fd: |
| detok = detokenize.Detokenizer(fd) |
| |
| self.assertEqual( |
| expected_tokens, |
| frozenset(detok.database.token_to_entries.keys()), |
| ) |
| |
| # Open ELF by path |
| detok = detokenize.Detokenizer(elf.name) |
| self.assertEqual( |
| expected_tokens, |
| frozenset(detok.database.token_to_entries.keys()), |
| ) |
| |
| # Open ELF by elf_reader.Elf |
| with open(elf.name, 'rb') as fd: |
| detok = detokenize.Detokenizer(elf_reader.Elf(fd)) |
| |
| self.assertEqual( |
| expected_tokens, |
| frozenset(detok.database.token_to_entries.keys()), |
| ) |
| finally: |
| os.unlink(elf.name) |
| |
| def test_decode_from_csv_file(self): |
| detok = detokenize.Detokenizer(io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS)) |
| expected_tokens = frozenset(detok.database.token_to_entries.keys()) |
| |
| csv_database = str(detok.database) |
| self.assertEqual(len(csv_database.splitlines()), TOKENS_IN_ELF) |
| |
| with tempfile.NamedTemporaryFile('w', delete=False) as csv_file: |
| try: |
| csv_file.write(csv_database) |
| csv_file.close() |
| |
| # Open CSV by path |
| detok = detokenize.Detokenizer(csv_file.name) |
| self.assertEqual( |
| expected_tokens, |
| frozenset(detok.database.token_to_entries.keys()), |
| ) |
| |
| # Open CSV by file object |
| with open(csv_file.name) as fd: |
| detok = detokenize.Detokenizer(fd) |
| |
| self.assertEqual( |
| expected_tokens, |
| frozenset(detok.database.token_to_entries.keys()), |
| ) |
| finally: |
| os.unlink(csv_file.name) |
| |
| def test_create_detokenizer_with_token_database(self): |
| detok = detokenize.Detokenizer(io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS)) |
| expected_tokens = frozenset(detok.database.token_to_entries.keys()) |
| |
| detok = detokenize.Detokenizer(detok.database) |
| self.assertEqual( |
| expected_tokens, frozenset(detok.database.token_to_entries.keys()) |
| ) |
| |
| |
| class DetokenizeWithCollisions(unittest.TestCase): |
| """Tests collision resolution.""" |
| |
| def setUp(self): |
| super().setUp() |
| token = 0xBAAD |
| |
| # Database with several conflicting tokens. |
| self.detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry( |
| token, 'REMOVED', date_removed=dt.datetime(9, 1, 1) |
| ), |
| tokens.TokenizedStringEntry(token, 'newer'), |
| tokens.TokenizedStringEntry( |
| token, 'A: %d', date_removed=dt.datetime(30, 5, 9) |
| ), |
| tokens.TokenizedStringEntry( |
| token, 'B: %c', date_removed=dt.datetime(30, 5, 10) |
| ), |
| tokens.TokenizedStringEntry(token, 'C: %s'), |
| tokens.TokenizedStringEntry(token, '%d%u'), |
| tokens.TokenizedStringEntry(token, '%s%u %d'), |
| tokens.TokenizedStringEntry(1, '%s'), |
| tokens.TokenizedStringEntry(1, '%d'), |
| tokens.TokenizedStringEntry(2, 'Three %s %s %s'), |
| tokens.TokenizedStringEntry(2, 'Five %d %d %d %d %s'), |
| ] |
| ) |
| ) |
| |
| def test_collision_no_args_favors_most_recently_present(self): |
| no_args = self.detok.detokenize(b'\xad\xba\0\0') |
| self.assertFalse(no_args.ok()) |
| self.assertEqual(len(no_args.successes), 2) |
| self.assertEqual(len(no_args.failures), 5) |
| self.assertEqual(len(no_args.matches()), 7) |
| self.assertEqual(str(no_args), 'newer') |
| self.assertEqual(len(no_args.best_result()[1]), 0) |
| self.assertEqual(no_args.best_result()[0], 'newer') |
| |
| def test_collision_one_integer_arg_favors_most_recently_present(self): |
| multiple_correct = self.detok.detokenize(b'\xad\xba\0\0\x7a') |
| self.assertFalse(multiple_correct.ok()) |
| self.assertIn('ERROR', repr(multiple_correct)) |
| self.assertEqual(len(multiple_correct.successes), 2) |
| self.assertEqual(len(multiple_correct.failures), 5) |
| self.assertEqual(len(multiple_correct.matches()), 7) |
| self.assertEqual(str(multiple_correct), 'B: =') |
| |
| def test_collision_one_integer_arg_favor_successful_decode(self): |
| # One string decodes successfully, since the arg is out of range for %c. |
| int_arg = self.detok.detokenize(b'\xad\xba\0\0\xfe\xff\xff\xff\x0f') |
| self.assertTrue(int_arg.ok()) |
| self.assertEqual(str(int_arg), 'A: 2147483647') |
| |
| def test_collision_one_string_arg_favors_successful_decode(self): |
| # One string decodes successfully, since decoding the argument as an |
| # integer does not decode all the data. |
| string_arg = self.detok.detokenize(b'\xad\xba\0\0\x02Hi') |
| self.assertTrue(string_arg.ok()) |
| self.assertEqual(str(string_arg), 'C: Hi') |
| |
| def test_collision_one_string_arg_favors_decoding_all_data(self): |
| result = self.detok.detokenize(b'\1\0\0\0\x83hi') |
| self.assertEqual(len(result.failures), 2) |
| # Should resolve to the string since %d would leave one byte behind. |
| self.assertEqual(str(result), '%s') |
| |
| def test_collision_multiple_args_favors_decoding_more_arguments(self): |
| result = self.detok.detokenize(b'\2\0\0\0\1\2\1\4\5') |
| self.assertEqual(len(result.matches()), 2) |
| self.assertEqual(result.matches()[0][0], 'Five -1 1 -1 2 %s') |
| self.assertEqual(result.matches()[1][0], 'Three \2 \4 %s') |
| |
| def test_collision_multiple_args_favors_decoding_all_arguments(self): |
| unambiguous = self.detok.detokenize(b'\xad\xba\0\0\x01#\x00\x01') |
| self.assertTrue(unambiguous.ok()) |
| self.assertEqual(len(unambiguous.matches()), 7) |
| self.assertEqual('#0 -1', str(unambiguous)) |
| self.assertIn('#0 -1', repr(unambiguous)) |
| |
| |
| class ManualPoolExecutor(concurrent.futures.Executor): |
| """A stubbed pool executor that captures the most recent work request |
| and holds it until the public process method is manually called.""" |
| |
| def __init__(self): |
| super().__init__() |
| self._func = None |
| |
| # pylint: disable=arguments-differ |
| def submit(self, func, *args, **kwargs): |
| """Submits work to the pool, stashing the partial for later use.""" |
| self._func = functools.partial(func, *args, **kwargs) |
| |
| def process(self): |
| """Processes the latest func submitted to the pool.""" |
| if self._func is not None: |
| self._func() |
| self._func = None |
| |
| |
| class InlinePoolExecutor(concurrent.futures.Executor): |
| """A stubbed pool executor that runs work immediately, inline.""" |
| |
| # pylint: disable=arguments-differ |
| def submit(self, func, *args, **kwargs): |
| """Submits work to the pool, stashing the partial for later use.""" |
| func(*args, **kwargs) |
| |
| |
| @mock.patch('os.path.getmtime') |
| class AutoUpdatingDetokenizerTest(unittest.TestCase): |
| """Tests the AutoUpdatingDetokenizer class.""" |
| |
| def test_update(self, mock_getmtime): |
| """Tests the update command.""" |
| |
| db = database.load_token_database( |
| io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS) |
| ) |
| self.assertEqual(len(db), TOKENS_IN_ELF) |
| |
| the_time = [100] |
| |
| def move_back_time_if_file_exists(path): |
| if os.path.exists(path): |
| the_time[0] -= 1 |
| return the_time[0] |
| |
| raise FileNotFoundError |
| |
| mock_getmtime.side_effect = move_back_time_if_file_exists |
| |
| with tempfile.NamedTemporaryFile('wb', delete=False) as file: |
| try: |
| file.close() |
| |
| pool = ManualPoolExecutor() |
| detok = detokenize.AutoUpdatingDetokenizer( |
| file.name, min_poll_period_s=0, pool=pool |
| ) |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| with open(file.name, 'wb') as fd: |
| tokens.write_binary(db, fd) |
| |
| # After the change but before the pool runs in another thread, |
| # the token should not exist. |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| # After the pool is allowed to process, it should. |
| pool.process() |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| finally: |
| os.unlink(file.name) |
| |
| def test_update_with_directory(self, mock_getmtime): |
| """Tests the update command with a directory format database.""" |
| db = database.load_token_database( |
| io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS) |
| ) |
| self.assertEqual(len(db), TOKENS_IN_ELF) |
| |
| the_time = [100] |
| |
| def move_back_time_if_file_exists(path): |
| if os.path.exists(path): |
| the_time[0] -= 1 |
| return the_time[0] |
| |
| raise FileNotFoundError |
| |
| mock_getmtime.side_effect = move_back_time_if_file_exists |
| |
| with tempfile.TemporaryDirectory() as dbdir: |
| with tempfile.NamedTemporaryFile( |
| 'wb', delete=False, suffix='.pw_tokenizer.csv', dir=dbdir |
| ) as matching_suffix_file, tempfile.NamedTemporaryFile( |
| 'wb', delete=False, suffix='.not.right', dir=dbdir |
| ) as mismatched_suffix_file: |
| try: |
| matching_suffix_file.close() |
| mismatched_suffix_file.close() |
| |
| pool = ManualPoolExecutor() |
| detok = detokenize.AutoUpdatingDetokenizer( |
| dbdir, min_poll_period_s=0, pool=pool |
| ) |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| with open(mismatched_suffix_file.name, 'wb') as fd: |
| tokens.write_csv(db, fd) |
| pool.process() |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| with open(matching_suffix_file.name, 'wb') as fd: |
| tokens.write_csv(db, fd) |
| |
| # After the change but before the pool runs in another |
| # thread, the token should not exist. |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| pool.process() |
| |
| # After the pool is allowed to process, it should. |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| finally: |
| os.unlink(mismatched_suffix_file.name) |
| os.unlink(matching_suffix_file.name) |
| os.rmdir(dbdir) |
| |
| # The database stays around if the file is deleted. |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| def test_no_update_if_time_is_same(self, mock_getmtime): |
| mock_getmtime.return_value = 100 |
| |
| with tempfile.NamedTemporaryFile('wb', delete=False) as file: |
| try: |
| tokens.write_csv( |
| database.load_token_database( |
| io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS) |
| ), |
| file, |
| ) |
| file.close() |
| |
| detok = detokenize.AutoUpdatingDetokenizer( |
| file.name, min_poll_period_s=0, pool=InlinePoolExecutor() |
| ) |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| # Empty the database, but keep the mock modified time the same. |
| with open(file.name, 'wb'): |
| pass |
| |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| self.assertTrue(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| |
| # Move back time so the now-empty file is reloaded. |
| mock_getmtime.return_value = 50 |
| self.assertFalse(detok.detokenize(JELLO_WORLD_TOKEN).ok()) |
| finally: |
| os.unlink(file.name) |
| |
| def test_token_domain_in_str(self, _) -> None: |
| """Tests a str containing a domain""" |
| detok = detokenize.AutoUpdatingDetokenizer( |
| f'{ELF_WITH_TOKENIZER_SECTIONS_PATH}#.*', |
| min_poll_period_s=0, |
| pool=InlinePoolExecutor(), |
| ) |
| self.assertEqual( |
| len(detok.database), TOKENS_IN_ELF_WITH_TOKENIZER_SECTIONS |
| ) |
| |
| def test_token_domain_in_path(self, _) -> None: |
| """Tests a Path() containing a domain""" |
| detok = detokenize.AutoUpdatingDetokenizer( |
| Path(f'{ELF_WITH_TOKENIZER_SECTIONS_PATH}#.*'), |
| min_poll_period_s=0, |
| pool=InlinePoolExecutor(), |
| ) |
| self.assertEqual( |
| len(detok.database), TOKENS_IN_ELF_WITH_TOKENIZER_SECTIONS |
| ) |
| |
| def test_token_no_domain_in_str(self, _) -> None: |
| """Tests a str without a domain""" |
| detok = detokenize.AutoUpdatingDetokenizer( |
| str(ELF_WITH_TOKENIZER_SECTIONS_PATH), |
| min_poll_period_s=0, |
| pool=InlinePoolExecutor(), |
| ) |
| self.assertEqual(len(detok.database), TOKENS_IN_ELF) |
| |
| def test_token_no_domain_in_path(self, _) -> None: |
| """Tests a Path() without a domain""" |
| detok = detokenize.AutoUpdatingDetokenizer( |
| ELF_WITH_TOKENIZER_SECTIONS_PATH, |
| min_poll_period_s=0, |
| pool=InlinePoolExecutor(), |
| ) |
| self.assertEqual(len(detok.database), TOKENS_IN_ELF) |
| |
| |
| def _next_char(message: bytes) -> bytes: |
| return bytes(b + 1 for b in message) |
| |
| |
| class NestedMessageParserTest(unittest.TestCase): |
| """Tests parsing prefixed messages.""" |
| |
| class _Case(NamedTuple): |
| data: bytes |
| expected: bytes |
| title: str |
| transform: Callable[[bytes], bytes] = _next_char |
| |
| TRANSFORM_TEST_CASES = ( |
| _Case(b'$abcd', b'%bcde', 'single message'), |
| _Case( |
| b'$$WHAT?$abc$WHY? is this $ok $', |
| b'%%WHAT?%bcd%WHY? is this %ok %', |
| 'message and non-message', |
| ), |
| _Case(b'$1$', b'%1%', 'empty message'), |
| _Case(b'$abc$defgh', b'%bcd%efghh', 'sequential message'), |
| _Case( |
| b'w$abcx$defygh$$abz', |
| b'w$ABCx$DEFygh$$ABz', |
| 'interspersed start/end non-message', |
| bytes.upper, |
| ), |
| _Case( |
| b'$abcx$defygh$$ab', |
| b'$ABCx$DEFygh$$AB', |
| 'interspersed start/end message ', |
| bytes.upper, |
| ), |
| ) |
| |
| def setUp(self) -> None: |
| self.decoder = detokenize.NestedMessageParser('$', 'abcdefg') |
| |
| def test_transform_io(self) -> None: |
| for data, expected, title, transform in self.TRANSFORM_TEST_CASES: |
| self.assertEqual( |
| expected, |
| b''.join( |
| self.decoder.transform_io(io.BytesIO(data), transform) |
| ), |
| f'{title}: {data!r}', |
| ) |
| |
| def test_transform_bytes_with_flush(self) -> None: |
| for data, expected, title, transform in self.TRANSFORM_TEST_CASES: |
| self.assertEqual( |
| expected, |
| self.decoder.transform(data, transform, flush=True), |
| f'{title}: {data!r}', |
| ) |
| |
| def test_transform_bytes_sequential(self) -> None: |
| transform = lambda message: message.upper().replace(b'$', b'*') |
| |
| self.assertEqual(self.decoder.transform(b'abc$abcd', transform), b'abc') |
| self.assertEqual(self.decoder.transform(b'$', transform), b'*ABCD') |
| self.assertEqual(self.decoder.transform(b'$b', transform), b'*') |
| self.assertEqual(self.decoder.transform(b'', transform), b'') |
| self.assertEqual(self.decoder.transform(b' ', transform), b'*B ') |
| self.assertEqual(self.decoder.transform(b'hello', transform), b'hello') |
| self.assertEqual(self.decoder.transform(b'?? $ab', transform), b'?? ') |
| self.assertEqual( |
| self.decoder.transform(b'123$ab4$56$a', transform), b'*AB123*AB4*56' |
| ) |
| self.assertEqual( |
| self.decoder.transform(b'bc', transform, flush=True), b'*ABC' |
| ) |
| |
| MESSAGES_TEST: Any = ( |
| (b'123$abc456$a', (False, b'123'), (True, b'$abc'), (False, b'456')), |
| (b'7$abcd', (True, b'$a'), (False, b'7')), |
| (b'e',), |
| (b'',), |
| (b'$', (True, b'$abcde')), |
| (b'$', (True, b'$')), |
| (b'$a$b$c', (True, b'$'), (True, b'$a'), (True, b'$b')), |
| (b'1', (True, b'$c'), (False, b'1')), |
| (b'',), |
| (b'?', (False, b'?')), |
| (b'!@', (False, b'!@')), |
| (b'%^&', (False, b'%^&')), |
| ) |
| |
| def test_read_messages(self) -> None: |
| for step in self.MESSAGES_TEST: |
| data: bytes = step[0] |
| pieces: tuple[tuple[bool, bytes], ...] = step[1:] |
| self.assertEqual(tuple(self.decoder.read_messages(data)), pieces) |
| |
| def test_read_messages_flush(self) -> None: |
| self.assertEqual( |
| list(self.decoder.read_messages(b'123$a')), [(False, b'123')] |
| ) |
| self.assertEqual(list(self.decoder.read_messages(b'b')), []) |
| self.assertEqual( |
| list(self.decoder.read_messages(b'', flush=True)), [(True, b'$ab')] |
| ) |
| |
| def test_read_messages_io(self) -> None: |
| # Rework the read_messages test data for stream input. |
| data = io.BytesIO(b''.join(step[0] for step in self.MESSAGES_TEST)) |
| expected_pieces = sum((step[1:] for step in self.MESSAGES_TEST), ()) |
| |
| result = self.decoder.read_messages_io(data) |
| for expected_is_message, expected_data in expected_pieces: |
| if expected_is_message: |
| is_message, piece = next(result) |
| self.assertTrue(is_message) |
| self.assertEqual(expected_data, piece) |
| else: # the IO version yields non-messages byte by byte |
| for byte in expected_data: |
| is_message, piece = next(result) |
| self.assertFalse(is_message) |
| self.assertEqual(bytes([byte]), piece) |
| |
| |
| class DetokenizeNested(unittest.TestCase): |
| """Tests detokenizing nested tokens""" |
| |
| def test_nested_hashed_arg(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry(0xA, 'tokenized argument'), |
| tokens.TokenizedStringEntry( |
| 2, |
| 'This is a ' + '$#%08x', |
| ), |
| ] |
| ) |
| ) |
| self.assertEqual( |
| str(detok.detokenize(b'\x02\0\0\0\x14')), |
| 'This is a tokenized argument', |
| ) |
| |
| def test_nested_base64_arg(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry(1, 'base64 argument'), |
| tokens.TokenizedStringEntry(2, 'This is a %s'), |
| ] |
| ) |
| ) |
| self.assertEqual( |
| str(detok.detokenize(b'\x02\0\0\0\x09$AQAAAA==')), # token for 1 |
| 'This is a base64 argument', |
| ) |
| |
| def test_deeply_nested_arg(self): |
| detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry(1, '$10#0000000005'), |
| tokens.TokenizedStringEntry(2, 'This is a $#%08x'), |
| tokens.TokenizedStringEntry(3, 'deeply nested argument'), |
| tokens.TokenizedStringEntry(4, '$AQAAAA=='), |
| tokens.TokenizedStringEntry(5, '$AwAAAA=='), |
| ] |
| ) |
| ) |
| self.assertEqual( |
| str(detok.detokenize(b'\x02\0\0\0\x08')), # token for 4 |
| 'This is a deeply nested argument', |
| ) |
| |
| |
| class DetokenizeBase64(unittest.TestCase): |
| """Tests detokenizing Base64 messages.""" |
| |
| JELLO = b'$' + base64.b64encode(JELLO_WORLD_TOKEN) |
| |
| RECURSION_STRING = f'The secret message is "{JELLO.decode()}"' |
| RECURSION = b'$' + base64.b64encode( |
| struct.pack('I', tokens.c_hash(RECURSION_STRING)) |
| ) |
| |
| RECURSION_STRING_2 = f"'{RECURSION.decode()}', said the spy." |
| RECURSION_2 = b'$' + base64.b64encode( |
| struct.pack('I', tokens.c_hash(RECURSION_STRING_2)) |
| ) |
| |
| TEST_CASES = ( |
| (b'', b''), |
| (b'nothing here', b'nothing here'), |
| (JELLO, b'Jello, world!'), |
| (JELLO + b'a', b'Jello, world!a'), |
| (JELLO + b'abc', b'Jello, world!abc'), |
| (JELLO + b'abc=', b'Jello, world!abc='), |
| (b'$a' + JELLO + b'a', b'$aJello, world!a'), |
| (b'Hello ' + JELLO + b'?', b'Hello Jello, world!?'), |
| (b'$' + JELLO, b'$Jello, world!'), |
| (JELLO + JELLO, b'Jello, world!Jello, world!'), |
| (JELLO + b'$' + JELLO, b'Jello, world!$Jello, world!'), |
| (JELLO + b'$a' + JELLO + b'bcd', b'Jello, world!$aJello, world!bcd'), |
| (b'$3141', b'$3141'), |
| (JELLO + b'$3141', b'Jello, world!$3141'), |
| ( |
| JELLO + b'$a' + JELLO + b'b' + JELLO + b'c', |
| b'Jello, world!$aJello, world!bJello, world!c', |
| ), |
| (RECURSION, b'The secret message is "Jello, world!"'), |
| ( |
| RECURSION_2, |
| b'\'The secret message is "Jello, world!"\', said the spy.', |
| ), |
| ) |
| |
| def setUp(self): |
| super().setUp() |
| db = database.load_token_database( |
| io.BytesIO(ELF_WITH_TOKENIZER_SECTIONS) |
| ) |
| db.add( |
| tokens.TokenizedStringEntry(tokens.c_hash(s), s) |
| for s in [self.RECURSION_STRING, self.RECURSION_STRING_2] |
| ) |
| self.detok = detokenize.Detokenizer(db) |
| |
| def test_detokenize_base64_live(self): |
| for data, expected in self.TEST_CASES: |
| output = io.BytesIO() |
| self.detok.detokenize_base64_live(io.BytesIO(data), output, '$') |
| |
| self.assertEqual(expected, output.getvalue(), f'Input: {data!r}') |
| |
| def test_detokenize_base64_to_file(self): |
| for data, expected in self.TEST_CASES: |
| output = io.BytesIO() |
| self.detok.detokenize_base64_to_file(data, output, '$') |
| |
| self.assertEqual(expected, output.getvalue()) |
| |
| def test_detokenize_base64(self): |
| for data, expected in self.TEST_CASES: |
| self.assertEqual(expected, self.detok.detokenize_base64(data, b'$')) |
| |
| def test_detokenize_base64_str(self): |
| for data, expected in self.TEST_CASES: |
| self.assertEqual( |
| expected.decode(), self.detok.detokenize_base64(data.decode()) |
| ) |
| |
| |
| class DetokenizeInfiniteRecursion(unittest.TestCase): |
| """Tests that infinite Base64 token recursion resolves.""" |
| |
| def setUp(self): |
| super().setUp() |
| self.detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry(0, '$AAAAAA=='), # token for 0 |
| tokens.TokenizedStringEntry(1, '$AgAAAA=='), # token for 2 |
| tokens.TokenizedStringEntry(2, '$#00000003'), # token for 3 |
| tokens.TokenizedStringEntry(3, '$AgAAAA=='), # token for 2 |
| ] |
| ) |
| ) |
| |
| def test_detokenize_self_recursion(self): |
| for depth in range(5): |
| self.assertEqual( |
| self.detok.detokenize_text( |
| b'This one is deep: $AAAAAA==', recursion=depth |
| ), |
| b'This one is deep: $AAAAAA==', |
| ) |
| |
| def test_detokenize_self_recursion_default(self): |
| self.assertEqual( |
| self.detok.detokenize_text( |
| b'This one is deep: $AAAAAA==', |
| ), |
| b'This one is deep: $AAAAAA==', |
| ) |
| |
| def test_detokenize_cyclic_recursion_even(self): |
| self.assertEqual( |
| self.detok.detokenize_text(b'I said "$AQAAAA=="', recursion=6), |
| b'I said "$AgAAAA=="', |
| ) |
| |
| def test_detokenize_cyclic_recursion_odd(self): |
| self.assertEqual( |
| self.detok.detokenize_text(b'I said "$AQAAAA=="', recursion=7), |
| b'I said "$#00000003"', |
| ) |
| |
| |
| class DetokenizeBase64InfiniteRecursion(unittest.TestCase): |
| """Tests that infinite Bas64 token recursion resolves.""" |
| |
| def setUp(self): |
| super().setUp() |
| self.detok = detokenize.Detokenizer( |
| tokens.Database( |
| [ |
| tokens.TokenizedStringEntry(0, '$AAAAAA=='), # token for 0 |
| tokens.TokenizedStringEntry(1, '$AgAAAA=='), # token for 2 |
| tokens.TokenizedStringEntry(2, '$AwAAAA=='), # token for 3 |
| tokens.TokenizedStringEntry(3, '$AgAAAA=='), # token for 2 |
| ] |
| ) |
| ) |
| |
| def test_detokenize_self_recursion(self): |
| for depth in range(5): |
| self.assertEqual( |
| self.detok.detokenize_base64( |
| b'This one is deep: $AAAAAA==', recursion=depth |
| ), |
| b'This one is deep: $AAAAAA==', |
| ) |
| |
| def test_detokenize_self_recursion_default(self): |
| self.assertEqual( |
| self.detok.detokenize_base64(b'This one is deep: $64#AAAAAA=='), |
| b'This one is deep: $AAAAAA==', |
| ) |
| |
| def test_detokenize_cyclic_recursion_even(self): |
| self.assertEqual( |
| self.detok.detokenize_base64(b'I said "$AQAAAA=="', recursion=2), |
| b'I said "$AgAAAA=="', |
| ) |
| |
| def test_detokenize_cyclic_recursion_odd(self): |
| self.assertEqual( |
| self.detok.detokenize_base64(b'I said "$AQAAAA=="', recursion=3), |
| b'I said "$AwAAAA=="', |
| ) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
