From 833615b7ce3af086e5448653fc25e9fabae3137c Mon Sep 17 00:00:00 2001 From: Ketan Patel Date: Thu, 14 Dec 2023 13:39:44 -0500 Subject: [PATCH] Add BoxUtil.search --- src/pymp4/cli.py | 15 ++----- src/pymp4/util.py | 20 +++++++++ tests/test_box.py | 10 ++++- tests/test_cli.py | 35 +++++++-------- tests/test_util.py | 108 +++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 156 insertions(+), 32 deletions(-) diff --git a/src/pymp4/cli.py b/src/pymp4/cli.py index 81a3422..1629a6d 100644 --- a/src/pymp4/cli.py +++ b/src/pymp4/cli.py @@ -3,10 +3,9 @@ import io import logging import argparse -import codecs from pymp4.parser import Box -from pymp4.util import BoxUtil +from pymp4.util import BoxUtil, escape_decode from construct import setGlobalPrintFullStrings log = logging.getLogger(__name__) @@ -23,20 +22,14 @@ def dump(args=None): setGlobalPrintFullStrings(args.full) - boxes_to_dump = {codecs.escape_decode(b)[0] for b in args.box or [] if b} + boxes_to_dump = {escape_decode(b) for b in args.box or [] if b} fd = args.input_file fd.seek(0, io.SEEK_END) eof = fd.tell() fd.seek(0) - def dump_box(box): - for t in boxes_to_dump: - for b in BoxUtil.find(box, t): - print(b) - - print_fn = dump_box if boxes_to_dump else print - while fd.tell() < eof: box = Box.parse_stream(fd) - print_fn(box) + for b in BoxUtil.search(box, *boxes_to_dump, decode=False): + print(b) diff --git a/src/pymp4/util.py b/src/pymp4/util.py index b72b4b4..394c800 100644 --- a/src/pymp4/util.py +++ b/src/pymp4/util.py @@ -15,6 +15,7 @@ See the License for the specific language governing permissions and limitations under the License. """ +import codecs import logging from pymp4.exceptions import BoxNotFound @@ -22,6 +23,10 @@ log = logging.getLogger(__name__) +def escape_decode(s): + return codecs.escape_decode(s)[0] + + class BoxUtil(object): @classmethod def child(cls, box, type_): @@ -59,6 +64,21 @@ def find(cls, box, type_): for fbox in cls.find(sbox, type_): yield fbox + @classmethod + def search(cls, box, *types, decode=True, key='type', children='children'): + if decode: + types = {escape_decode(t) for t in types if t} + if not types: + yield box + else: + if getattr(box, key, None) in types: + yield box + for sbox in (getattr(box, children, None) or []): + for fbox in cls.search(sbox, *types, decode=False, + key=key, children=children): + yield fbox + + @classmethod def find_extended(cls, box, extended_type_): if hasattr(box, "extended_type"): diff --git a/tests/test_box.py b/tests/test_box.py index 6759846..e8f29b6 100644 --- a/tests/test_box.py +++ b/tests/test_box.py @@ -15,15 +15,23 @@ limitations under the License. """ import logging +import pytest import unittest -from construct import Container, ListContainer +from construct import Container, ListContainer, StreamError from pymp4.parser import Box log = logging.getLogger(__name__) class BoxTests(unittest.TestCase): + def test_build_non_fourcc_type(self): + with pytest.raises(StreamError) as err_info: + Box.build(Container( + type=b"MP4", + )) + assert 'bytes object of wrong length, expected 4, found 3' in err_info.value.args[0] + def test_ftyp_parse(self): self.assertEqual( Box.parse(b'\x00\x00\x00\x18ftypiso5\x00\x00\x00\x01iso5avc1'), diff --git a/tests/test_cli.py b/tests/test_cli.py index 183183b..0f8dae0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -14,8 +14,7 @@ def test_missing_arg(capsys): captured = capsys.readouterr() assert "the following arguments are required: FILE" in captured.err -def test_dump(capsys): - cli.dump(args=[TEST_FILE]) +def __assert_full_dump(capsys): captured = capsys.readouterr() assert "" == captured.err all_types = re.findall(r'type = [^\s]+', captured.out) @@ -24,15 +23,13 @@ def test_dump(capsys): assert captured.out.rstrip().endswith('end = 60') assert 'truncated,' not in captured.out +def test_dump(capsys): + cli.dump(args=[TEST_FILE]) + __assert_full_dump(capsys) + def test_full_dump(capsys): cli.dump(args=['--full', TEST_FILE]) - captured = capsys.readouterr() - assert "" == captured.err - all_types = re.findall(r'type = [^\s]+', captured.out) - assert all_types == ["type = b'ftyp'", "type = b'free'"] - assert "major_brand = b'M4A ' " in captured.out - assert captured.out.rstrip().endswith('end = 60') - assert 'truncated,' not in captured.out + __assert_full_dump(capsys) def test_truncated_dump(capsys): cli.dump(args=['--truncated', TEST_FILE]) @@ -42,8 +39,7 @@ def test_truncated_dump(capsys): assert captured.out.rstrip().endswith('end = 60') assert '\\x16\'... (truncated, total 24)' in captured.out -def test_dumping_specific_box(capsys): - cli.dump(args=['-b', 'free', TEST_FILE]) +def __assert_free_box_only(capsys): captured = capsys.readouterr() assert "" == captured.err @@ -55,15 +51,14 @@ def test_dumping_specific_box(capsys): assert captured.out.rstrip().endswith('end = 60') assert "!\"#$' (total 24)" in captured.out +def test_dumping_specific_box(capsys): + cli.dump(args=['-b', 'free', TEST_FILE]) + __assert_free_box_only(capsys) + def test_dumping_specific_box_with_escaped_input(capsys): cli.dump(args=['-b', '\\x66ree', TEST_FILE]) - captured = capsys.readouterr() - assert "" == captured.err + __assert_free_box_only(capsys) - all_types = re.findall(r'type = [^\s]+', captured.out) - assert all_types == ["type = b'free'"] - - assert "major_brand = b'M4A ' " not in captured.out - assert "type = b'free'" in captured.out - assert captured.out.rstrip().endswith('end = 60') - assert "!\"#$' (total 24)" in captured.out +def test_dumping_specific_box_with_empty_string(capsys): + cli.dump(args=['-b', '', TEST_FILE]) + __assert_full_dump(capsys) diff --git a/tests/test_util.py b/tests/test_util.py index 615bc8b..25f507c 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -76,6 +76,114 @@ class BoxTests(unittest.TestCase): ]) ) + box_search_data = Container( + type=b"demo", + children=ListContainer([ + Container(type=b"a ", id=1), + Container(type=b"b ", id=2), + Container( + type=b"c ", + children=ListContainer([ + Container(type="a ", id=3), + Container(type=b"b ", id=4), + Container(type=b"\xa9too", id=99), + Container(type=b"c ", id=98), + ]) + ), + Container(type=b"d ", id=5), + Container(type="zzzz", id=100, + entries=ListContainer([ + Container(type=b"zzzz", id=101), + Container(type="xyzz", id=102) + ]) + ), + ]) + ) + + def test_search_convert_to_bytes_false(self): + actual = list(BoxUtil.search(self.box_data, "a ", decode=False)) + expect = [Container(type="a ", id=1), Container(type="a ", id=3)] + assert actual == expect + + def test_search_not_found(self): + # convert_to_bytes true so "a " converts to b"a " + actual = list(BoxUtil.search(self.box_data, "a ")) + expect = [] + assert actual == expect + + def test_search(self): + # convert_to_bytes true so "a " converts to b"a " + actual = list(BoxUtil.search(self.box_search_data, "a ")) + expect = [Container(type=b"a ", id=1)] + assert actual == expect + + def test_search_converts_bytes2bytes(self): + # convert_to_bytes true so b"\xa9too" converts to b"\xa9to" + actual = list(BoxUtil.search(self.box_search_data, b"\xa9too")) + expect = [Container(type=b"\xa9too", id=99)] + assert actual == expect + + def test_search_converts_str2bytes(self): + # convert_to_bytes true so "\\xa9too" converts to b"\xa9too" + actual = list(BoxUtil.search(self.box_search_data, "\\xa9too")) + expect = [Container(type=b"\xa9too", id=99)] + assert actual == expect + + def test_search_converts_utf8_str2bytes(self): + # convert_to_bytes true so "\xa9too" converts to b"\xc2\xa9too" + actual = list(BoxUtil.search(self.box_search_data, "\xa9too")) + expect = [] + assert actual == expect + + def test_search_includes_childrens(self): + # convert_to_bytes true so "b " converts to b"b " + actual = list(BoxUtil.search(self.box_search_data, "b ", decode=True)) + expect = [Container(type=b"b ", id=2), Container(type=b"b ", id=4)] + assert actual == expect + + def test_search_includes_childrens_of_matched(self): + # convert_to_bytes true so "c " converts to b"c " + actual = list(BoxUtil.search(self.box_search_data, "c ", decode=True)) + expect = [self.box_search_data.children[2], Container(type=b"c ", id=98)] + assert actual == expect + + def test_search_key_is_id(self): + actual = list(BoxUtil.search(self.box_search_data, 2, 3, key='id', decode=False)) + expect = [Container(type=b"b ", id=2), Container(type="a ", id=3)] + assert actual == expect + + def test_search_children_key_is_entries(self): + actual = list(BoxUtil.search(self.box_search_data.children[-1], 'zzzz', children='entries')) + expect = [Container(type=b"zzzz", id=101)] + assert actual == expect + + def test_search_key_does_not_exist(self): + # convert_to_bytes true so "a " converts to b"a " + actual = list(BoxUtil.search(self.box_search_data, "a ", key='NOTE')) + expect = [] + assert actual == expect + + def test_search_childen_key_does_not_exist_in_root_box(self): + # convert_to_bytes true so "a " converts to b"a " + actual = list(BoxUtil.search(self.box_search_data, "a ", children='entries')) + expect = [] + assert actual == expect + + def test_search_empty_key_returns_self(self): + actual = list(BoxUtil.search(self.box_search_data.children[0])) + expect = [self.box_search_data.children[0]] + assert actual == expect + + def test_search_none_key_returns_self(self): + actual = list(BoxUtil.search(self.box_search_data.children[0], None)) + expect = [self.box_search_data.children[0]] + assert actual == expect + + def test_search_empty_str_key_returns_self(self): + actual = list(BoxUtil.search(self.box_search_data.children[0], "")) + expect = [self.box_search_data.children[0]] + assert actual == expect + def test_child(self): self.assertListEqual( list(BoxUtil.child(self.box_data, "a ")),