Nymiro
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
import struct
|
|
from typing import List, BinaryIO
|
|
from io import BytesIO
|
|
|
|
from base import Box # Import the base Box class
|
|
from heic_types import ItemLocationBox, PrimaryItemBox, ItemInfoBox
|
|
|
|
# The factory map remains here
|
|
BOX_TYPE_MAP = {
|
|
'iloc': ItemLocationBox,
|
|
'pitm': PrimaryItemBox,
|
|
'iinf': ItemInfoBox,
|
|
}
|
|
|
|
CONTAINER_BOXES = {'meta', 'moov', 'trak', 'iprp', 'ipco', 'dinf', 'fiinf', 'ipro'}
|
|
FULL_BOXES = {'meta', 'hdlr', 'pitm', 'iinf', 'iloc'}
|
|
|
|
def parse_boxes(stream: BinaryIO, max_size: int) -> List[Box]:
|
|
"""
|
|
Parses boxes from a file stream up to a maximum size.
|
|
Now also handles recursive parsing.
|
|
"""
|
|
boxes = []
|
|
start_pos_in_stream = stream.tell()
|
|
|
|
while stream.tell() - start_pos_in_stream < max_size:
|
|
current_offset_in_stream = stream.tell()
|
|
|
|
header = stream.read(8)
|
|
if len(header) < 8: break
|
|
|
|
size = struct.unpack('>I', header[:4])[0]
|
|
box_type = header[4:].decode('ascii', errors='ignore')
|
|
|
|
header_size = 8
|
|
if size == 1:
|
|
largesize_header = stream.read(8)
|
|
if len(largesize_header) < 8: break
|
|
size = struct.unpack('>Q', largesize_header)[0]
|
|
header_size = 16
|
|
elif size == 0:
|
|
size = max_size - (current_offset_in_stream - start_pos_in_stream)
|
|
|
|
if size < header_size: break
|
|
|
|
content_size = size - header_size
|
|
|
|
stream.seek(current_offset_in_stream + header_size)
|
|
raw_data = stream.read(content_size)
|
|
if len(raw_data) < content_size: break
|
|
|
|
# --- Factory Pattern ---
|
|
box_class = BOX_TYPE_MAP.get(box_type, Box)
|
|
box = box_class(size, box_type, current_offset_in_stream, raw_data)
|
|
|
|
# --- NEW: Recursive parsing logic moved here ---
|
|
if box.type in CONTAINER_BOXES:
|
|
child_stream = BytesIO(box.raw_data)
|
|
parse_size = len(box.raw_data)
|
|
if box.type in FULL_BOXES:
|
|
child_stream.read(4) # Skip version/flags
|
|
parse_size -= 4
|
|
box.children = parse_boxes(child_stream, parse_size)
|
|
# --- End of new logic ---
|
|
|
|
boxes.append(box)
|
|
stream.seek(current_offset_in_stream + size)
|
|
|
|
return boxes |